API REFERENCE/omnibias-convex

omnibias.convex.arrangement.jax

0.1.0a2AlphaPython >=3.10AGPL / commercial

Differentiable LP and QP solver paths, barrier methods, KKT sensitivities, and scoped optimality evidence.

import omnibias.convex.arrangement.jax

Module contract & mathematical conventions

JAX twin of arrangement-LP soft membership (theory 03-02).

soft_membership matches the numpy / torch twins at float64 (jax_enable_x64). Soft mode is temperature collapse (beta -> inf, feasibility), not the founding bias collapse (delta -> 0). Do not conflate the two.

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
DiffMode	omnibias.convex.arrangement.DiffMode
LPOutput	omnibias.convex.arrangement.LPOutput
LearnedPolytope	omnibias.convex.arrangement.LearnedPolytope
named_pentagon	omnibias.convex.arrangement.named_pentagon
soft_cell_gap_bound	omnibias.convex.arrangement.soft_cell_gap_bound
solve_arrangement_lp	omnibias.convex.arrangement.solve_arrangement_lp
sound_lower_bound	omnibias.convex.arrangement.sound_lower_bound

CLASS

LPLayer

omnibias.convex.arrangement.jax.LPLayerSource ↗

class LPLayer

Learned-polytope LP. mode is returned so the gradient is interpretable.

Parameters

Name	Type & default	Contract
modepositional-or-keyword	DiffModeRequired	See the operation contract and implementation below.
betakeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Fields & methods

__init__method

__init__(self, mode: DiffMode, *, beta: float | None=None) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
modepositional-or-keyword	DiffModeRequired	See the operation contract and implementation below.
betakeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When self.mode is DiffMode.SOFT and self.beta <= 0.0, raises ValueError('SOFT mode needs beta > 0').

Method implementation

    def __init__(self, mode: DiffMode, *, beta: float | None = None) -> None:
        self.mode = DiffMode(mode)
        self.beta = 8.0 if beta is None else float(beta)
        if self.mode is DiffMode.SOFT and self.beta <= 0.0:
            raise ValueError("SOFT mode needs beta > 0")

View source ↗

forwardmethod

forward(self, polytope: LearnedPolytope, c: Array) -> LPOutput

See the class contract above and the exact method implementation below.

Parameters

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

Returns: LPOutput

Method implementation

    def forward(self, polytope: LearnedPolytope, c: Array) -> LPOutput:
        cv = jnp.asarray(c, dtype=jnp.float64)
        return solve_arrangement_lp(polytope, cv, mode=self.mode, beta=self.beta)

View source ↗

modeattribute

mode = DiffMode(mode)

Public instance attribute assigned by __init__. The source assignment is mode = DiffMode(mode); it describes instance state, not a literal constructor default.

View source ↗

betaattribute

beta = 8.0 if beta is None else float(beta)

Public instance attribute assigned by __init__. The source assignment is beta = 8.0 if beta is None else float(beta); it describes instance state, not a literal constructor default.

View source ↗

Read the implementation

class LPLayer:
    """Learned-polytope LP. ``mode`` is returned so the gradient is interpretable."""

    def __init__(self, mode: DiffMode, *, beta: float | None = None) -> None:
        self.mode = DiffMode(mode)
        self.beta = 8.0 if beta is None else float(beta)
        if self.mode is DiffMode.SOFT and self.beta <= 0.0:
            raise ValueError("SOFT mode needs beta > 0")

    def forward(self, polytope: LearnedPolytope, c: Array) -> LPOutput:
        cv = jnp.asarray(c, dtype=jnp.float64)
        return solve_arrangement_lp(polytope, cv, mode=self.mode, beta=self.beta)

Related APIsomnibias.convex.arrangement.DiffMode↗omnibias.convex.arrangement.solve_arrangement_lp↗

FUNCTION

soft_membership

omnibias.convex.arrangement.jax.soft_membershipSource ↗

soft_membership(normals: Array, offsets: Array, x: Array, *, beta: float) -> Array

prod_i sigma(beta (b_i - a_i · x)).

Parameters

Name	Type & default	Contract
normalspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
offsetspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betakeyword-only	floatRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · g6 soft membership parity

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_g6_soft_membership_parity() -> None:
    poly, _c, _opt = named_pentagon()
    x = np.array([1.5, 1.5], dtype=np.float64)
    ref = float(soft_membership(poly, x, beta=5.0))
    torch = pytest.importorskip("torch")
    tw = __import__("omnibias.convex.arrangement.torch", fromlist=["soft_membership"]).soft_membership
    got_t = float(tw(torch.as_tensor(poly.normals), torch.as_tensor(poly.offsets), torch.as_tensor(x), beta=5.0))
    assert math.isclose(got_t, ref, rel_tol=0.0, abs_tol=1e-15)
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    jw = __import__("omnibias.convex.arrangement.jax", fromlist=["soft_membership"]).soft_membership
    got_j = float(np.asarray(jw(poly.normals, poly.offsets, x, beta=5.0)))
    assert math.isclose(got_j, ref, rel_tol=0.0, abs_tol=2e-15)

python -m pytest "packages/omnibias-convex/tests/test_arrangement_lp.py::test_g6_soft_membership_parity" -q

Read the implementation

def soft_membership(normals: Array, offsets: Array, x: Array, *, beta: float) -> Array:
    """``prod_i sigma(beta (b_i - a_i · x))``."""
    a = jnp.asarray(normals, dtype=jnp.float64)
    b = jnp.asarray(offsets, dtype=jnp.float64).reshape(-1)
    xv = jnp.asarray(x, dtype=jnp.float64)
    if xv.ndim == 1:
        slack = b - a @ xv
        return jnp.prod(jax_sigmoid(float(beta) * slack))
    slack = b[None, :] - xv @ a.T
    return jnp.prod(jax_sigmoid(float(beta) * slack), axis=-1)

← Package overviewRunnable learning examples →