API REFERENCE/omnibias-boolean

omnibias.boolean.jax.ops.solver

0.1.0a2AlphaPython >=3.10Apache-2.0

Boolean truth tables, multilinear extensions, algebraic spectra, and finite-temperature gate systems.

import omnibias.boolean.jax.ops.solver

Module contract & mathematical conventions

Differentiable Boolean equation/system solver (jax): propose-and-verify.

Twin of omnibias.boolean.torch.ops.solver. The exact GF(2) fast-path and the propose-and-verify contract are shared (pure-Python omnibias.boolean._core.systems); the soft search anneals x = sigmoid(beta * theta) with a hand-rolled Adam over jax.grad of the multilinear-extension residual.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
BetaAnnealScheduler	omnibias.binary.BetaAnnealScheduler

CLASS

BooleanSystem

omnibias.boolean.jax.ops.BooleanSystemSource ↗

class BooleanSystem

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

A system of Boolean equations phi_k(x) = 0 (each phi_k a truth table).

Parameters

Name	Type & default	Contract
ndataclass field	intRequired	See the operation contract and implementation below.
constraintsdataclass field	tuple[TruthTable, ...]Required	See the operation contract and implementation below.

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Fields & methods

nfield

n: int

Declared int field on BooleanSystem. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

constraintsfield

constraints: tuple[TruthTable, ...]

Declared tuple[TruthTable, ...] field on BooleanSystem. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

from_constraintsmethod

from_constraints(cls, constraints: Sequence[TruthTable]) -> BooleanSystem

Build from constraint truth tables (1 = violated).

Parameters

Name	Type & default	Contract
constraintspositional-or-keyword	Sequence[TruthTable]Required	See the operation contract and implementation below.

Returns: BooleanSystem

Validation & failure cases

• When not constraints, raises ValueError('need at least one constraint').
• When num_vars(c) != n, raises ValueError('all constraints must share the same arity').

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Method implementation

    def from_constraints(cls, constraints: Sequence[TruthTable]) -> BooleanSystem:
        """Build from constraint truth tables (``1`` = violated)."""
        if not constraints:
            raise ValueError("need at least one constraint")
        n = num_vars(constraints[0])
        for c in constraints:
            if num_vars(c) != n:
                raise ValueError("all constraints must share the same arity")
        return cls(n=n, constraints=tuple(tuple(c) for c in constraints))

View source ↗

from_predicatesmethod

from_predicates(cls, predicates: Sequence[Callable[..., bool]], n: int) -> BooleanSystem

Build from predicates; phi_k is 0 exactly where pred_k holds.

Parameters

Name	Type & default	Contract
predicatespositional-or-keyword	Sequence[Callable[..., bool]]Required	See the operation contract and implementation below.
npositional-or-keyword	intRequired	See the operation contract and implementation below.

Returns: BooleanSystem

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Method implementation

    def from_predicates(
        cls, predicates: Sequence[Callable[..., bool]], n: int
    ) -> BooleanSystem:
        """Build from predicates; ``phi_k`` is ``0`` exactly where ``pred_k`` holds."""
        constraints = [
            tuple(0 if pred(*bits) else 1 for bits in all_assignments(n))
            for pred in predicates
        ]
        return cls.from_constraints(constraints)

View source ↗

residual_softmethod

residual_soft(self, x: Array) -> Array

Sum of squared multilinear-extension violations at a soft point x.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

Returns: Array

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Method implementation

    def residual_soft(self, x: Array) -> Array:
        """Sum of squared multilinear-extension violations at a soft point ``x``."""
        total = jnp.zeros((), dtype=x.dtype)
        for c in self.constraints:
            phi = _multilinear_eval_point(jnp.asarray(c, dtype=x.dtype), x)
            total = total + phi * phi
        return total

View source ↗

verifymethod

verify(self, bits: Sequence[int]) -> bool

Exact check: every equation holds at the (hard) assignment bits.

Parameters

Name	Type & default	Contract
bitspositional-or-keyword	Sequence[int]Required	See the operation contract and implementation below.

Returns: bool

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Method implementation

    def verify(self, bits: Sequence[int]) -> bool:
        """Exact check: every equation holds at the (hard) assignment ``bits``."""
        return verify_assignment(self.constraints, bits)

View source ↗

is_linearmethod

is_linear(self) -> bool

True iff every equation is XOR-linear (GF(2) fast-path applies).

Returns: bool

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Method implementation

    def is_linear(self) -> bool:
        """``True`` iff every equation is XOR-linear (GF(2) fast-path applies)."""
        return constraints_are_linear(self.constraints)

View source ↗

combinedmethod

combined(self) -> TruthTable

The single OR-combined constraint (1 where any equation is violated).

Returns: TruthTable

Method implementation

    def combined(self) -> TruthTable:
        """The single OR-combined constraint (``1`` where any equation is violated)."""
        return system_constraint(self.constraints)

View source ↗

Read the implementation

class BooleanSystem:
    """A system of Boolean equations ``phi_k(x) = 0`` (each ``phi_k`` a truth table)."""

    n: int
    constraints: tuple[TruthTable, ...]

    @classmethod
    def from_constraints(cls, constraints: Sequence[TruthTable]) -> BooleanSystem:
        """Build from constraint truth tables (``1`` = violated)."""
        if not constraints:
            raise ValueError("need at least one constraint")
        n = num_vars(constraints[0])
        for c in constraints:
            if num_vars(c) != n:
                raise ValueError("all constraints must share the same arity")
        return cls(n=n, constraints=tuple(tuple(c) for c in constraints))

    @classmethod
    def from_predicates(
        cls, predicates: Sequence[Callable[..., bool]], n: int
    ) -> BooleanSystem:
        """Build from predicates; ``phi_k`` is ``0`` exactly where ``pred_k`` holds."""
        constraints = [
            tuple(0 if pred(*bits) else 1 for bits in all_assignments(n))
            for pred in predicates
        ]
        return cls.from_constraints(constraints)

    def residual_soft(self, x: Array) -> Array:
        """Sum of squared multilinear-extension violations at a soft point ``x``."""
        total = jnp.zeros((), dtype=x.dtype)
        for c in self.constraints:
            phi = _multilinear_eval_point(jnp.asarray(c, dtype=x.dtype), x)
            total = total + phi * phi
        return total

    def verify(self, bits: Sequence[int]) -> bool:
        """Exact check: every equation holds at the (hard) assignment ``bits``."""
        return verify_assignment(self.constraints, bits)

    def is_linear(self) -> bool:
        """``True`` iff every equation is XOR-linear (GF(2) fast-path applies)."""
        return constraints_are_linear(self.constraints)

    def combined(self) -> TruthTable:
        """The single OR-combined constraint (``1`` where any equation is violated)."""
        return system_constraint(self.constraints)

Related APIsomnibias.boolean.system_constraint↗omnibias.boolean.constraints_are_linear↗omnibias.boolean.verify_assignment↗omnibias.boolean.all_assignments↗omnibias.boolean.num_vars↗

Other public import paths

• omnibias.boolean.jax.ops.BooleanSystem
• omnibias.boolean.jax.ops.solver.BooleanSystem

CLASS

SolveResult

omnibias.boolean.jax.ops.SolveResultSource ↗

class SolveResult

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Outcome of solve: a hardened assignment plus its exact verification.

Parameters

Name	Type & default	Contract
assignmentdataclass field	tuple[int, ...] | NoneRequired	See the operation contract and implementation below.
verifieddataclass field	boolRequired	See the operation contract and implementation below.
residualdataclass field	floatRequired	See the operation contract and implementation below.
methoddataclass field	strRequired	See the operation contract and implementation below.
relaxeddataclass field	tuple[float, ...] | NoneDefault: None	See the operation contract and implementation below.

Fields & methods

assignmentfield

assignment: tuple[int, ...] | None

Declared tuple[int, ...] | None field on SolveResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

verifiedfield

verified: bool

Declared bool field on SolveResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

residualfield

residual: float

Declared float field on SolveResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

methodfield

method: str

Declared str field on SolveResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

relaxedfield

relaxed: tuple[float, ...] | None = None

Declared tuple[float, ...] | None field on SolveResult. Default expression: None.

View source ↗

Read the implementation

class SolveResult:
    """Outcome of :func:`solve`: a hardened assignment plus its exact verification."""

    assignment: tuple[int, ...] | None
    verified: bool
    residual: float
    method: str
    relaxed: tuple[float, ...] | None = None

Other public import paths

• omnibias.boolean.jax.ops.SolveResult
• omnibias.boolean.jax.ops.solver.SolveResult

FUNCTION

brute_force_solutions

omnibias.boolean.jax.ops.brute_force_solutionsSource ↗

brute_force_solutions(system: BooleanSystem) -> list[tuple[int, ...]]

All exact solutions (for tests / verification on small systems).

Parameters

Name	Type & default	Contract
systempositional-or-keyword	BooleanSystemRequired	See the operation contract and implementation below.

Returnslist[tuple[int, ...]]

Read the implementation

def brute_force_solutions(system: BooleanSystem) -> list[tuple[int, ...]]:
    """All exact solutions (for tests / verification on small systems)."""
    out = []
    for idx in range(1 << system.n):
        bits = assignment(idx, system.n)
        if system.verify(bits):
            out.append(bits)
    return out

Related APIsomnibias.boolean.assignment↗

Other public import paths

• omnibias.boolean.jax.ops.brute_force_solutions
• omnibias.boolean.jax.ops.solver.brute_force_solutions

FUNCTION

solve

omnibias.boolean.jax.ops.solveSource ↗

solve(system: BooleanSystem, *, schedule: BetaAnnealScheduler | None=None, steps: int=300, lr: float=0.2, restarts: int=16, seed: int=0, use_gf2: bool=True) -> SolveResult

Solve system (GF(2) fast-path when linear, else annealed soft search).

Parameters

Name	Type & default	Contract
systempositional-or-keyword	BooleanSystemRequired	See the operation contract and implementation below.
schedulekeyword-only	BetaAnnealScheduler | NoneDefault: None	See the operation contract and implementation below.
stepskeyword-only	intDefault: 300	See the operation contract and implementation below.
lrkeyword-only	floatDefault: 0.2	See the operation contract and implementation below.
restartskeyword-only	intDefault: 16	See the operation contract and implementation below.
seedkeyword-only	intDefault: 0	See the operation contract and implementation below.
use_gf2keyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsSolveResult

Executed regression · jax solver gf2 and anneal

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_jax_solver_gf2_and_anneal() -> None:
    lin = jb.BooleanSystem.from_predicates(
        [lambda a, b: (a ^ b) == 1, lambda a, b: b == 1], n=2
    )
    res = jb.solve(lin)
    assert res.method == "gf2" and res.verified and res.assignment == (0, 1)

    nonlin = jb.BooleanSystem.from_predicates(
        [lambda a, b, c: (a & b) == 1, lambda a, b, c: (b | c) == 1], n=3
    )
    res = jb.solve(nonlin, steps=200, restarts=12, seed=0)
    assert res.verified and res.assignment is not None
    assert nonlin.verify(res.assignment)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_jax_solver_gf2_and_anneal" -q

Read the implementation

def solve(
    system: BooleanSystem,
    *,
    schedule: BetaAnnealScheduler | None = None,
    steps: int = 300,
    lr: float = 0.2,
    restarts: int = 16,
    seed: int = 0,
    use_gf2: bool = True,
) -> SolveResult:
    """Solve ``system`` (GF(2) fast-path when linear, else annealed soft search)."""
    if use_gf2 and system.is_linear():
        gsol = gf2_solve(system.constraints)
        if gsol is not None:
            if not gsol.consistent:
                return SolveResult(None, False, float("inf"), "gf2")
            bits = gsol.particular
            return SolveResult(bits, system.verify(bits), 0.0, "gf2")

    if schedule is None:
        schedule = BetaAnnealScheduler(1.0, 20.0, steps, "exp")
    rng = np.random.default_rng(seed)
    best: SolveResult | None = None
    for _ in range(restarts):
        theta0 = jnp.asarray(rng.standard_normal(system.n), dtype=jnp.float64)
        theta = _adam_solve(system, theta0, schedule, steps, lr)
        x = jax.nn.sigmoid(schedule.value(steps) * theta)
        bits = tuple(int(v >= 0.5) for v in np.asarray(x))
        res = float(system.residual_soft(x))
        cand = SolveResult(
            bits, system.verify(bits), res, "anneal", tuple(float(v) for v in np.asarray(x))
        )
        if best is None or (cand.verified and not best.verified) or (
            cand.verified == best.verified and cand.residual < best.residual
        ):
            best = cand
        if best.verified:
            break
    assert best is not None
    return best

Related APIsomnibias.binary.BetaAnnealScheduler↗omnibias.boolean.gf2_solve↗omnibias.boolean.jax.ops.SolveResult↗

Other public import paths

• omnibias.boolean.jax.ops.solve
• omnibias.boolean.jax.ops.solver.solve

← Package overviewRunnable learning examples →