omnibias.convex.jax.solver
Differentiable LP and QP solver paths, barrier methods, KKT sensitivities, and scoped optimality evidence.
import omnibias.convex.jax.solverModule contract & mathematical conventions
Closed-form-Hessian log-barrier interior-point LP/QP solver (JAX).
Solves
\min_x \tfrac12 x^\top Q x + c^\top x \quad\text{s.t.}\quad A x \le bwith Q positive semidefinite (Q = 0 => LP). The log-barrier subproblem
\varphi_t(x) = t\,\big(\tfrac12 x^\top Q x + c^\top x\big) - \sum_i \log s_i,
\qquad s = b - A x > 0is minimised by a damped Newton method whose gradient and **closed-form Hessian**
\nabla\varphi_t = t (Q x + c) + A^\top (1/s), \qquad
\nabla^2\varphi_t = t\,Q + A^\top \operatorname{diag}(1/s^2) Aare assembled directly and solved with jax.numpy.linalg.solve -- the same
closed-form-Hessian Newton pattern as omnibias.curvature.mse_newton_step. The
barrier weight t is increased along a short central path; the surrogate
duality gap is m / t and the dual multipliers are lambda = 1/(t s).
A strictly feasible start is found from x = 0 when b > 0; otherwise a
phase-1 min tau s.t. A x - tau <= b LP (the same solver) supplies one.
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.
InfeasibleProblemError
omnibias.convex.jax.InfeasibleProblemErrorSource ↗class InfeasibleProblemError(ValueError)Raised when no strictly feasible point A x < b can be found.
Regression example · infeasible raises
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.
def test_infeasible_raises() -> None:
# x <= -1 and x >= 1 simultaneously: empty feasible set.
c = _arr([1.0])
A = _arr([[1.0], [-1.0]])
b = _arr([-1.0, -1.0])
with pytest.raises(InfeasibleProblemError):
solve_lp(c, A, b)python -m pytest packages/omnibias-convex/tests/test_solver_jax.py -k test_infeasible_raises -qRegression example · warm start strict feasibility check
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.
def test_warm_start_strict_feasibility_check() -> None:
c = _arr([1.0, 1.0])
A = _arr([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]])
b = _arr([5.0, 5.0, 0.0, 0.0])
with pytest.raises(InfeasibleProblemError):
solve_lp(c, A, b, x0=_arr([10.0, 10.0])) # outside the boxpython -m pytest packages/omnibias-convex/tests/test_solver_jax.py -k test_warm_start_strict_feasibility_check -qRead the implementation
class InfeasibleProblemError(ValueError):
"""Raised when no strictly feasible point ``A x < b`` can be found."""Other public import paths
omnibias.convex.jax.InfeasibleProblemErroromnibias.convex.jax.solver.InfeasibleProblemError
solve_lp
omnibias.convex.jax.solve_lpSource ↗solve_lp(c: Array, A: Array, b: Array, *, x0: Array | None=None, options: BarrierOptions | None=None) -> ConvexSolution[Array]Solve the LP min c^T x s.t. A x <= b (Q = 0 special case).
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. |
x0keyword-only | Array | NoneDefault: None | See the operation contract and implementation below. |
optionskeyword-only | BarrierOptions | NoneDefault: None | See the operation contract and implementation below. |
ConvexSolution[Array]Executed regression · infeasible raises
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.
def test_infeasible_raises() -> None:
# x <= -1 and x >= 1 simultaneously: empty feasible set.
c = _arr([1.0])
A = _arr([[1.0], [-1.0]])
b = _arr([-1.0, -1.0])
with pytest.raises(InfeasibleProblemError):
solve_lp(c, A, b)python -m pytest "packages/omnibias-convex/tests/test_solver_jax.py::test_infeasible_raises" -qRegression example · lp matches newton interior point
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.
def test_lp_matches_newton_interior_point() -> None:
c, A, b = _lp_2d()
sol = solve_lp_penalty(c, A, b)
ref = solve_lp(c, A, b)
np.testing.assert_allclose(np.asarray(sol.x), np.asarray(ref.x), atol=5e-3)
np.testing.assert_allclose(float(sol.obj), float(ref.obj), atol=5e-3)python -m pytest packages/omnibias-convex/tests/test_penalty_jax.py -k test_lp_matches_newton_interior_point -qRead the implementation
def solve_lp(
c: Array,
A: Array,
b: Array,
*,
x0: Array | None = None,
options: BarrierOptions | None = None,
) -> ConvexSolution[Array]:
r"""Solve the LP ``min c^T x s.t. A x <= b`` (``Q = 0`` special case)."""
c = jnp.asarray(c)
n = c.shape[0]
Q = jnp.zeros((n, n), dtype=c.dtype)
return solve_qp(Q, c, A, b, x0=x0, options=options)Other public import paths
omnibias.convex.jax.solve_lpomnibias.convex.jax.solver.solve_lp
solve_qp
omnibias.convex.jax.solve_qpSource ↗solve_qp(Q: Array, c: Array, A: Array, b: Array, *, x0: Array | None=None, options: BarrierOptions | None=None) -> ConvexSolution[Array]Solve min 1/2 x^T Q x + c^T x s.t. A x <= b (Q PSD).
Parameters
Q, c, A, b: Problem data with shapes (n, n), (n,), (m, n), (m,).
x0: Optional strictly feasible warm start (A x0 < b). If omitted, a start is found from x = 0 (when b > 0) or a phase-1 LP.
options: ~omnibias.convex.problem.BarrierOptions tuning.
Returns
~omnibias.convex.problem.ConvexSolution with primal x, dual
lambda = 1/(t s) >= 0, slacks s, objective, gap m/t and flags.
Parameters
| Name | Type & default | Contract |
|---|---|---|
Qpositional-or-keyword | ArrayRequired | Problem data with shapes (n, n), (n,), (m, n), (m,). |
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. |
x0keyword-only | Array | NoneDefault: None | Optional strictly feasible warm start (A x0 < b). If omitted, a start is found from x = 0 (when b > 0) or a phase-1 LP. |
optionskeyword-only | BarrierOptions | NoneDefault: None | ~omnibias.convex.problem.BarrierOptions tuning. |
ConvexSolution[Array]Explicit validation & failure cases
- When
bool(jnp.any(b - A @ start <= 0.0)), raisesInfeasibleProblemError('provided x0 is not strictly feasible (A x0 < b violated)').
These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.
Executed regression · infeasible raises
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.
def test_infeasible_raises() -> None:
# x <= -1 and x >= 1 simultaneously: empty feasible set.
c = _arr([1.0])
A = _arr([[1.0], [-1.0]])
b = _arr([-1.0, -1.0])
with pytest.raises(InfeasibleProblemError):
solve_lp(c, A, b)python -m pytest "packages/omnibias-convex/tests/test_solver_jax.py::test_infeasible_raises" -qRegression example · qp layer matches eager solver
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.
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 -k test_qp_layer_matches_eager_solver -qRead the implementation
def solve_qp(
Q: Array,
c: Array,
A: Array,
b: Array,
*,
x0: Array | None = None,
options: BarrierOptions | None = None,
) -> ConvexSolution[Array]:
r"""Solve ``min 1/2 x^T Q x + c^T x s.t. A x <= b`` (``Q`` PSD).
Parameters
----------
Q, c, A, b:
Problem data with shapes ``(n, n)``, ``(n,)``, ``(m, n)``, ``(m,)``.
x0:
Optional strictly feasible warm start (``A x0 < b``). If omitted, a start
is found from ``x = 0`` (when ``b > 0``) or a phase-1 LP.
options:
:class:`~omnibias.convex.problem.BarrierOptions` tuning.
Returns
-------
:class:`~omnibias.convex.problem.ConvexSolution` with primal ``x``, dual
``lambda = 1/(t s) >= 0``, slacks ``s``, objective, gap ``m/t`` and flags.
"""
opts = options or BarrierOptions()
Q = jnp.asarray(Q)
c = jnp.asarray(c)
A = jnp.asarray(A)
b = jnp.asarray(b)
n = c.shape[0]
m = b.shape[0]
validate_shapes(n, m, A_shape=A.shape, b_shape=b.shape, c_shape=c.shape, Q_shape=Q.shape)
if x0 is not None:
start = jnp.asarray(x0, dtype=b.dtype)
phase1_newton = 0
else:
start, phase1_newton = _strictly_feasible_start(A, b, opts)
if bool(jnp.any(b - A @ start <= 0.0)):
raise InfeasibleProblemError("provided x0 is not strictly feasible (A x0 < b violated)")
x, slack, t, iterations, gap, converged, newton_total = _central_path_solve(
Q, c, A, b, start, opts
)
dual = 1.0 / (t * slack)
return ConvexSolution(
x=x,
dual=dual,
slack=slack,
obj=_objective(Q, c, x),
gap=gap,
iterations=iterations,
converged=converged,
newton_iterations=newton_total + phase1_newton,
)Other public import paths
omnibias.convex.jax.solve_qpomnibias.convex.jax.solver.solve_qp