API REFERENCE/omnibias-convex

omnibias.convex.jax.layer

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.convex.jax.layer

Module contract & mathematical conventions

Differentiable argmin for LP/QP -- the NN-native convex layer (JAX).

qp_layer(Q, c, A, b) returns the QP optimiser x*(Q, c, A, b) and is differentiable in all of Q, c, A, b via the **KKT implicit-function theorem** (OptNet / cvxpylayers style): the forward solve is the interior-point method from omnibias.convex.jax.solver, and the backward pass solves the linearised KKT system once for the adjoint -- no unrolling through the solver iterations.

At the optimum the KKT residual

r(x, \lambda) = \begin{bmatrix} Q x + c + A^\top \lambda \\
                 \operatorname{diag}(\lambda)(A x - b) \end{bmatrix} = 0

has Jacobian (w.r.t. (x, lambda))

K = \begin{bmatrix} Q & A^\top \\
    \operatorname{diag}(\lambda) A & \operatorname{diag}(A x - b) \end{bmatrix},

and for an upstream cotangent g = dL/dx we solve K^T [y_x; y_l] = -[g; 0] and read off the parameter gradients (Amos & Kolter, 2017).

The forward solve runs eagerly through jax.pure_callback, so qp_layer is usable under jax.grad / jax.jacobian without making the (data-dependent, Python-control-flow) interior-point loop itself trace-able.

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 →

FUNCTION

lp_layer

omnibias.convex.jax.lp_layerSource ↗

lp_layer(c: Array, A: Array, b: Array, options: BarrierOptions | None=None) -> Array

Differentiable LP optimiser x* of min c^T x s.t. A x <= b (Q = 0).

Parameters

Name	Type & default	Contract
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
optionspositional-or-keyword	BarrierOptions | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · lp layer grad b vs finite diff and c is flat

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_lp_layer_grad_b_vs_finite_diff_and_c_is_flat() -> None:
    # LP optimum is a vertex: it moves with b but is locally constant in c.
    rng = np.random.default_rng(5)
    n = 2
    c = rng.standard_normal(n)
    A = np.vstack([rng.standard_normal((2, n)), np.eye(n), -np.eye(n)])
    b = np.concatenate([rng.uniform(0.5, 1.5, size=2), 3.0 * np.ones(2 * n)])
    w = rng.standard_normal(n)

    def loss_b(bv: jnp.ndarray) -> jnp.ndarray:
        return jnp.dot(_arr(w), lp_layer(_arr(c), _arr(A), bv))

    def eager(bb: np.ndarray) -> float:
        return float(np.dot(w, np.asarray(solve_lp(_arr(c), _arr(A), _arr(bb)).x)))

    grad_b = np.asarray(jax.grad(loss_b)(_arr(b)))
    for _ in range(3):
        d = rng.standard_normal(b.shape[0])
        np.testing.assert_allclose(
            float(np.dot(grad_b, d)), _dir_fd(b, d, eager), atol=_ATOL, rtol=_RTOL
        )

    def loss_c(cv: jnp.ndarray) -> jnp.ndarray:
        return jnp.dot(_arr(w), lp_layer(cv, _arr(A), _arr(b)))

    grad_c = np.asarray(jax.grad(loss_c)(_arr(c)))
    assert np.max(np.abs(grad_c)) < 1e-3  # vertex is flat in c

python -m pytest "packages/omnibias-convex/tests/test_layer_jax.py::test_lp_layer_grad_b_vs_finite_diff_and_c_is_flat" -q

Read the implementation

def lp_layer(
    c: Array, A: Array, b: Array, options: BarrierOptions | None = None
) -> Array:
    """Differentiable LP optimiser ``x*`` of ``min c^T x s.t. A x <= b`` (``Q = 0``)."""
    c = jnp.asarray(c)
    n = c.shape[0]
    Q = jnp.zeros((n, n), dtype=c.dtype)
    return qp_layer(Q, c, A, b, options)

Related APIsomnibias.convex.jax.qp_layer↗

Other public import paths

• omnibias.convex.jax.lp_layer
• omnibias.convex.jax.layer.lp_layer

FUNCTION

qp_layer

omnibias.convex.jax.qp_layerSource ↗

qp_layer(Q: Array, c: Array, A: Array, b: Array, options: BarrierOptions | None=None) -> Array

Differentiable QP optimiser x* of min 1/2 x^T Q x + c^T x s.t. A x <= b.

Parameters

Name	Type & default	Contract
Qpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
cpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
Apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
optionspositional-or-keyword	BarrierOptions | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · qp layer matches eager solver

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_qp_layer_matches_eager_solver() -> None:
    Q, c, A, b = _random_qp(0)
    x_layer = np.asarray(qp_layer(_arr(Q), _arr(c), _arr(A), _arr(b)))
    x_eager = np.asarray(solve_qp(_arr(Q), _arr(c), _arr(A), _arr(b)).x)
    np.testing.assert_allclose(x_layer, x_eager, atol=1e-9)

python -m pytest "packages/omnibias-convex/tests/test_layer_jax.py::test_qp_layer_matches_eager_solver" -q

Regression example · qp layer grad c vs finite diff

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_qp_layer_grad_c_vs_finite_diff() -> None:
    Q, c, A, b = _random_qp(1)
    rng = np.random.default_rng(10)
    w = rng.standard_normal(c.shape[0])

    def loss_c(cv: jnp.ndarray) -> jnp.ndarray:
        return jnp.dot(_arr(w), qp_layer(_arr(Q), cv, _arr(A), _arr(b)))

    grad = np.asarray(jax.grad(loss_c)(_arr(c)))
    for _ in range(3):
        d = rng.standard_normal(c.shape[0])
        ad = float(np.dot(grad, d))
        fd = _dir_fd(c, d, lambda cc: _loss_eager(Q, cc, A, b, w))
        np.testing.assert_allclose(ad, fd, atol=_ATOL, rtol=_RTOL)

python -m pytest packages/omnibias-convex/tests/test_layer_jax.py -k test_qp_layer_grad_c_vs_finite_diff -q

Read the implementation

def qp_layer(
    Q: Array, c: Array, A: Array, b: Array, options: BarrierOptions | None = None
) -> Array:
    """Differentiable QP optimiser ``x*`` of ``min 1/2 x^T Q x + c^T x s.t. A x <= b``."""
    x, _, _ = _solve_callback(Q, c, A, b, options)
    return x

Other public import paths

• omnibias.convex.jax.qp_layer
• omnibias.convex.jax.layer.qp_layer

← Package overviewRunnable learning examples →