API REFERENCE/omnibias-convex

omnibias.convex.torch.layer

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.convex.torch.layer

Module contract & mathematical conventions

Differentiable argmin for LP/QP (torch) -- the KKT implicit-function layer.

Torch twin of omnibias.convex.jax.layer: a torch.autograd.Function whose forward is the interior-point solve and whose backward solves the linearised KKT system K^T [y_x; y_l] = -[g; 0] once for the adjoint (OptNet style). Runs in float64.

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.torch.lp_layerSource ↗

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

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

Parameters

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

ReturnsTensor

Executed regression · lp layer runs and is flat in c

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_runs_and_is_flat_in_c() -> None:
    rng = np.random.default_rng(4)
    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)])
    ct = _t(c, grad=True)
    lp_layer(ct, _t(A), _t(b)).sum().backward()
    assert np.max(np.abs(ct.grad.numpy())) < 1e-3  # LP optimum is a vertex, flat in c

python -m pytest "packages/omnibias-convex/tests/test_layer_torch.py::test_lp_layer_runs_and_is_flat_in_c" -q

Read the implementation

def lp_layer(
    c: Tensor, A: Tensor, b: Tensor, options: BarrierOptions | None = None
) -> Tensor:
    """Differentiable LP optimiser ``x*`` of ``min c^T x s.t. A x <= b`` (``Q = 0``)."""
    n = c.shape[0]
    Q = torch.zeros((n, n), dtype=torch.float64)
    x: Tensor = _QPLayerFn.apply(Q, c, A, b, options)  # type: ignore[no-untyped-call]
    return x

Other public import paths

• omnibias.convex.torch.lp_layer
• omnibias.convex.torch.layer.lp_layer

FUNCTION

qp_layer

omnibias.convex.torch.qp_layerSource ↗

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

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	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Apositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
optionspositional-or-keyword	BarrierOptions | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · qp layer grad c finite diff

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_grad_c_finite_diff() -> None:
    Q, c, A, b = _random_qp(1)
    rng = np.random.default_rng(21)
    w = rng.standard_normal(c.shape[0])
    wt = torch.as_tensor(w, dtype=torch.float64)

    ct = _t(c, grad=True)
    torch.dot(wt, qp_layer(_t(Q), ct, _t(A), _t(b))).backward()
    grad = ct.grad.numpy()

    def eager(cc: np.ndarray) -> float:
        return float(np.dot(w, solve_qp(_t(Q), _t(cc), _t(A), _t(b)).x.numpy()))

    eps = 1e-3
    for _ in range(3):
        d = rng.standard_normal(c.shape[0])
        ad = float(np.dot(grad, d))
        fd = (eager(c + eps * d) - eager(c - eps * d)) / (2.0 * eps)
        np.testing.assert_allclose(ad, fd, atol=3e-4, rtol=3e-3)

python -m pytest "packages/omnibias-convex/tests/test_layer_torch.py::test_qp_layer_grad_c_finite_diff" -q

Regression example · qp layer grad matches jax

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_matches_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.convex.jax import qp_layer as qp_layer_jax

    Q, c, A, b = _random_qp(2)
    rng = np.random.default_rng(20)
    w = rng.standard_normal(c.shape[0])
    wt = torch.as_tensor(w, dtype=torch.float64)

    # ---- torch grads (one backward per parameter) ----
    def torch_grad(which: str) -> np.ndarray:
        Qt, ct, At, bt = (_t(Q, which == "Q"), _t(c, which == "c"),
                          _t(A, which == "A"), _t(b, which == "b"))
        leaf = {"Q": Qt, "c": ct, "A": At, "b": bt}[which]
        torch.dot(wt, qp_layer(Qt, ct, At, bt)).backward()
        return leaf.grad.numpy()

    # ---- jax grads ----
    def jax_grad(argnum: int) -> np.ndarray:
        f = lambda Qv, cv, Av, bv: jnp.dot(  # noqa: E731
            jnp.asarray(w), qp_layer_jax(Qv, cv, Av, bv)
        )
        args = (jnp.asarray(Q), jnp.asarray(c), jnp.asarray(A), jnp.asarray(b))
        return np.asarray(jax.grad(f, argnums=argnum)(*args))

    for which, argnum in [("c", 1), ("b", 3), ("A", 2), ("Q", 0)]:
        np.testing.assert_allclose(torch_grad(which), jax_grad(argnum), rtol=1e-7, atol=1e-8)

python -m pytest packages/omnibias-convex/tests/test_layer_torch.py -k test_qp_layer_grad_matches_jax -q

Read the implementation

def qp_layer(
    Q: Tensor, c: Tensor, A: Tensor, b: Tensor, options: BarrierOptions | None = None
) -> Tensor:
    """Differentiable QP optimiser ``x*`` of ``min 1/2 x^T Q x + c^T x s.t. A x <= b``."""
    x: Tensor = _QPLayerFn.apply(Q, c, A, b, options)  # type: ignore[no-untyped-call]
    return x

Other public import paths

• omnibias.convex.torch.qp_layer
• omnibias.convex.torch.layer.qp_layer

← Package overviewRunnable learning examples →