API REFERENCE/omnibias-convex

omnibias.convex.jax.penalty

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.convex.jax.penalty

Module contract & mathematical conventions

Tempered temperature-collapse penalty LP/QP solver by gradient descent (JAX).

Turns a linear (or convex-quadratic) program into a plain **gradient-descent** problem. Each constraint a_i^T x <= b_i is a hyperplane; the *temperature-collapse* unit sigma(beta (a_i^T x - b_i)) (the omnibias tempered heaviside) has that hyperplane as its decision boundary, and its integral is the smooth hinge penalty softplus(beta u) / beta -> max(u, 0) as beta -> inf. Summing the penalties gives the smooth *exterior* objective

F(x) = c^\top x + \tfrac12 x^\top Q x
       + \mu \sum_i \tfrac1\beta \operatorname{softplus}(\beta (a_i^\top x - b_i)),

whose gradient is **closed form** -- a sum of temperature-collapse sigmoids, no autodiff:

\nabla F = c + Q x + \mu A^\top \sigma(\beta (A x - b)).

The sigmoid is spelled jax.nn.sigmoid rather than routed through the omnibias fastpath: at order 1, softplus' = sigmoid *is* the tower's value, so the two are the same expression and the framework primitive is the cheaper spelling. Any higher-order term would have to go through omnibias.jax.

penalty_descent minimises F for a fixed (beta, mu) by accelerated (Nesterov) gradient descent with the closed-form Lipschitz step eta = step_safety / L, L = ||Q||_2 + (mu beta / 4) ||A||_2^2; it is pure (jit / grad friendly) -- the differentiable "LP as gradient descent" primitive. solve_qp_penalty / solve_lp_penalty run it along a beta / mu homotopy (each stage warm-starts the next) and return a ~omnibias.convex.problem.ConvexSolution whose dual estimate is the temperature-collapse activation itself, lambda = mu sigma(beta (A x - b)) >= 0.

Honest scope: a first-order, GPU/batch-friendly, differentiable homotopy solver that converges to a *tolerance* -- the complement of, not a replacement for, the Newton ~omnibias.convex.jax.solver.solve_qp (which reaches ~1e-12). The exterior penalty needs **no** strictly feasible start and no phase-1.

Terminology: the *temperature-collapse* unit here is the beta -> inf **feasibility** sense of "collapse" (the constraint sigmoid saturates to a 0/1 step); it is distinct from omnibias's **founding bias collapse** -- the multi-bias delta -> 0 limit sum_k s_k sigma(z + b_k) -> sigma^(K-1) that yields a smooth *derivative* (see docs/derivatives.md and omnibias.torch.unit).

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

penalty_descent

omnibias.convex.jax.penalty_descentSource ↗

penalty_descent(Q: Array, c: Array, A: Array, b: Array, x0: Array, *, beta: float, mu: float, steps: int, step_safety: float=0.9, prox: float=0.0, anchor: Array | None=None, A_eq: Array | None=None, b_eq: Array | None=None) -> Array

Minimise F (plus an optional proximal term) for fixed (beta, mu).

Accelerated (Nesterov) gradient descent with the closed-form Lipschitz step eta = step_safety / L. The optional proximal term (prox / 2) ||x - anchor||^2 (anchor defaults to x0) makes the subproblem coercive/strongly convex -- it is what stops the exterior penalty running off to infinity while mu is below the largest optimal multiplier, and it vanishes at a homotopy fixed point (x == anchor) so the recovered optimum is unbiased. Optional equality constraints A_eq x = b_eq are enforced by the quadratic exterior penalty (mu/2)||A_eq x - b_eq||^2 (closed-form gradient, Lipschitz bump).

Pure (no data-dependent control flow), so with prox = 0 on a bounded (strongly convex) problem it is jax.jit / jax.grad friendly -- the differentiable "LP as gradient descent" primitive. Returns the iterate x.

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.
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betakeyword-only	floatRequired	See the operation contract and implementation below.
mukeyword-only	floatRequired	See the operation contract and implementation below.
stepskeyword-only	intRequired	See the operation contract and implementation below.
step_safetykeyword-only	floatDefault: 0.9	See the operation contract and implementation below.
proxkeyword-only	floatDefault: 0.0	See the operation contract and implementation below.
anchorkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
A_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
b_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · lp matches newton interior point

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_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::test_lp_matches_newton_interior_point" -q

Regression example · penalty descent is differentiable

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_penalty_descent_is_differentiable() -> None:
    # jax.grad through a fixed-beta descent (strongly convex, prox=0) is finite/nonzero.
    _, A, b = _lp_2d()

    def solved_sum(cost: jnp.ndarray) -> jnp.ndarray:
        x = penalty_descent(
            jnp.eye(2), cost, A, b, jnp.zeros(2), beta=8.0, mu=6.0, steps=300, prox=0.0
        )
        return jnp.sum(x)

    grad = jax.grad(solved_sum)(jnp.array([-3.0, -2.0]))
    assert bool(jnp.all(jnp.isfinite(grad)))
    assert float(jnp.max(jnp.abs(grad))) > 1e-6

python -m pytest packages/omnibias-convex/tests/test_penalty_jax.py -k test_penalty_descent_is_differentiable -q

Read the implementation

def penalty_descent(
    Q: Array,
    c: Array,
    A: Array,
    b: Array,
    x0: Array,
    *,
    beta: float,
    mu: float,
    steps: int,
    step_safety: float = 0.9,
    prox: float = 0.0,
    anchor: Array | None = None,
    A_eq: Array | None = None,
    b_eq: Array | None = None,
) -> Array:
    r"""Minimise ``F`` (plus an optional proximal term) for fixed ``(beta, mu)``.

    Accelerated (Nesterov) gradient descent with the closed-form Lipschitz step
    ``eta = step_safety / L``. The optional proximal term ``(prox / 2) ||x - anchor||^2``
    (``anchor`` defaults to ``x0``) makes the subproblem coercive/strongly convex --
    it is what stops the exterior penalty running off to infinity while ``mu`` is
    below the largest optimal multiplier, and it vanishes at a homotopy fixed point
    (``x == anchor``) so the recovered optimum is unbiased. Optional equality
    constraints ``A_eq x = b_eq`` are enforced by the quadratic exterior penalty
    ``(mu/2)||A_eq x - b_eq||^2`` (closed-form gradient, Lipschitz bump).

    Pure (no data-dependent control flow), so with ``prox = 0`` on a bounded
    (strongly convex) problem it is ``jax.jit`` / ``jax.grad`` friendly -- the
    differentiable "LP as gradient descent" primitive. Returns the iterate ``x``.
    """
    anchor_pt = x0 if anchor is None else anchor
    eta = _lipschitz_step(Q, A, beta, mu, prox, step_safety, A_eq)

    def body(_: Array, carry: tuple[Array, Array, Array]) -> tuple[Array, Array, Array]:
        x, y, t_acc = carry
        grad = (
            penalty_gradient(Q, c, A, b, y, beta=beta, mu=mu, A_eq=A_eq, b_eq=b_eq)
            + prox * (y - anchor_pt)
        )
        x_next = y - eta * grad
        t_next = 0.5 * (1.0 + jnp.sqrt(1.0 + 4.0 * t_acc * t_acc))
        y_next = x_next + ((t_acc - 1.0) / t_next) * (x_next - x)
        return x_next, y_next, t_next

    one = jnp.asarray(1.0, dtype=x0.dtype)
    x_final, _, _ = jax.lax.fori_loop(0, steps, body, (x0, x0, one))
    return cast(Array, x_final)

Related APIsomnibias.convex.jax.penalty_gradient↗

Other public import paths

• omnibias.convex.jax.penalty_descent
• omnibias.convex.jax.penalty.penalty_descent

FUNCTION

penalty_gradient

omnibias.convex.jax.penalty_gradientSource ↗

penalty_gradient(Q: Array, c: Array, A: Array, b: Array, x: Array, *, beta: float, mu: float, A_eq: Array | None=None, b_eq: Array | None=None) -> Array

Closed-form gradient of F (inequality + optional equality penalty).

\nabla F = c + Q x + \mu A^\top \sigma(\beta (A x - b))
           + \mu A_{eq}^\top (A_{eq} x - b_{eq}).

Each term sigma(beta (a_i^T x - b_i)) is the temperature-collapse unit of inequality i; weighted by mu and summed through A^T it is the exterior-penalty force pushing x toward the feasible polytope. The optional equality block adds the **quadratic** exterior penalty (mu/2)||A_eq x - b_eq||^2 whose closed-form gradient is mu A_eq^T (A_eq x - b_eq).

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.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betakeyword-only	floatRequired	See the operation contract and implementation below.
mukeyword-only	floatRequired	See the operation contract and implementation below.
A_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
b_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · lp matches newton interior point

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_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::test_lp_matches_newton_interior_point" -q

Regression example · closed form gradient equals autograd

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_closed_form_gradient_equals_autograd() -> None:
    # penalty_gradient is exactly grad of F = c.x + 1/2 x'Qx + mu*sum softplus(beta u)/beta.
    _, A, b = _lp_2d()
    Q = jnp.eye(2)
    c = jnp.array([0.3, -1.2])
    x = jnp.array([0.7, 0.4])
    beta, mu = 5.3, 2.1

    def objective(z: jnp.ndarray) -> jnp.ndarray:
        u = A @ z - b
        return c @ z + 0.5 * z @ (Q @ z) + mu * jnp.sum(jax.nn.softplus(beta * u)) / beta

    g_auto = jax.grad(objective)(x)
    g_closed = penalty_gradient(Q, c, A, b, x, beta=beta, mu=mu)
    np.testing.assert_allclose(np.asarray(g_closed), np.asarray(g_auto), atol=1e-12)

python -m pytest packages/omnibias-convex/tests/test_penalty_jax.py -k test_closed_form_gradient_equals_autograd -q

Read the implementation

def penalty_gradient(
    Q: Array,
    c: Array,
    A: Array,
    b: Array,
    x: Array,
    *,
    beta: float,
    mu: float,
    A_eq: Array | None = None,
    b_eq: Array | None = None,
) -> Array:
    r"""Closed-form gradient of ``F`` (inequality + optional equality penalty).

    .. math::
        \nabla F = c + Q x + \mu A^\top \sigma(\beta (A x - b))
                   + \mu A_{eq}^\top (A_{eq} x - b_{eq}).

    Each term ``sigma(beta (a_i^T x - b_i))`` is the temperature-collapse unit of
    inequality ``i``; weighted by ``mu`` and summed through ``A^T`` it is the
    exterior-penalty force pushing ``x`` toward the feasible polytope. The optional
    equality block adds the **quadratic** exterior penalty ``(mu/2)||A_eq x - b_eq||^2``
    whose closed-form gradient is ``mu A_eq^T (A_eq x - b_eq)``.
    """
    u = A @ x - b
    activations = jax.nn.sigmoid(beta * u)
    grad = c + Q @ x + mu * (A.T @ activations)
    if A_eq is not None and b_eq is not None:
        grad = grad + mu * (A_eq.T @ (A_eq @ x - b_eq))
    return grad

Other public import paths

• omnibias.convex.jax.penalty_gradient
• omnibias.convex.jax.penalty.penalty_gradient

FUNCTION

solve_lp_penalty

omnibias.convex.jax.solve_lp_penaltySource ↗

solve_lp_penalty(c: Array, A: Array, b: Array, *, A_eq: Array | None=None, b_eq: Array | None=None, x0: Array | None=None, options: PenaltyOptions | None=None) -> ConvexSolution[Array]

Solve the LP min c^T x s.t. A x <= b [, A_eq x = b_eq] (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.
A_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
b_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
x0keyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
optionskeyword-only	PenaltyOptions | NoneDefault: None	See the operation contract and implementation below.

ReturnsConvexSolution[Array]

Executed regression · lp matches newton interior point

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_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::test_lp_matches_newton_interior_point" -q

Regression example · qp projection onto nonnegative orthant

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_projection_onto_nonnegative_orthant() -> None:
    # min 1/2||x - p||^2  s.t.  -x <= 0   =>   x* = relu(p).
    p = _arr([1.5, -2.0, 0.7, -0.1])
    sol = solve_qp_penalty(jnp.eye(4), -p, -jnp.eye(4), jnp.zeros(4))
    np.testing.assert_allclose(np.asarray(sol.x), np.maximum(np.asarray(p), 0.0), atol=5e-3)

python -m pytest packages/omnibias-convex/tests/test_penalty_jax.py -k test_qp_projection_onto_nonnegative_orthant -q

Read the implementation

def solve_lp_penalty(
    c: Array,
    A: Array,
    b: Array,
    *,
    A_eq: Array | None = None,
    b_eq: Array | None = None,
    x0: Array | None = None,
    options: PenaltyOptions | None = None,
) -> ConvexSolution[Array]:
    r"""Solve the LP ``min c^T x  s.t.  A x <= b [, A_eq x = b_eq]`` (``Q = 0``)."""
    ca = jnp.asarray(c)
    n = ca.shape[0]
    Q = jnp.zeros((n, n), dtype=ca.dtype)
    return solve_qp_penalty(Q, ca, A, b, A_eq=A_eq, b_eq=b_eq, x0=x0, options=options)

Related APIsomnibias.convex.jax.solve_qp_penalty↗

Other public import paths

• omnibias.convex.jax.solve_lp_penalty
• omnibias.convex.jax.penalty.solve_lp_penalty

FUNCTION

solve_qp_penalty

omnibias.convex.jax.solve_qp_penaltySource ↗

solve_qp_penalty(Q: Array, c: Array, A: Array, b: Array, *, A_eq: Array | None=None, b_eq: Array | None=None, x0: Array | None=None, options: PenaltyOptions | None=None) -> ConvexSolution[Array]

Solve min 1/2 x^T Q x + c^T x s.t. A x <= b [, A_eq x = b_eq].

First-order homotopy: accelerated gradient descent on the tempered penalty with beta / mu grown geometrically. x0 may be **any** point (the exterior penalty needs no feasible start); it defaults to the origin. Optional equality constraints A_eq x = b_eq are enforced by a quadratic exterior penalty. Returns a ~omnibias.convex.problem.ConvexSolution with the inequality dual lambda = mu sigma(beta (A x - b)) >= 0, the equality dual estimate nu = mu (A_eq x - b_eq) in eq_dual (None if no equalities), iterations = homotopy stages run and newton_iterations = total steps.

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.
A_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
b_eqkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
x0keyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
optionskeyword-only	PenaltyOptions | NoneDefault: None	See the operation contract and implementation below.

ReturnsConvexSolution[Array]

Explicit validation & failure cases

• When Aeq is None or beq is None, raises ValueError(f"equality constraints require both A_eq and b_eq; got A_eq={('set' if Aeq is not None else None)}, b_eq={('set' if beq is not None else None)}").
• When Aeq is None or beq is None, raises ValueError(f"equality constraints require both A_eq and b_eq; got A_eq={('set' if Aeq is not None else None)}, b_eq={('set' if beq is not None else None)}").

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · lp matches newton interior point

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_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::test_lp_matches_newton_interior_point" -q

Regression example · bad equality shape rejected

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_bad_equality_shape_rejected() -> None:
    Q, c, A, b, _, _ = _hyperplane_qp()
    with pytest.raises(ValueError):
        jpen.solve_qp_penalty(
            jnp.asarray(Q), jnp.asarray(c), jnp.asarray(A), jnp.asarray(b),
            A_eq=jnp.asarray(np.ones((1, 5))), b_eq=jnp.asarray([1.0]),
        )

python -m pytest packages/omnibias-convex/tests/test_penalty_equality.py -k test_bad_equality_shape_rejected -q

Read the implementation

def solve_qp_penalty(
    Q: Array,
    c: Array,
    A: Array,
    b: Array,
    *,
    A_eq: Array | None = None,
    b_eq: Array | None = None,
    x0: Array | None = None,
    options: PenaltyOptions | None = None,
) -> ConvexSolution[Array]:
    r"""Solve ``min 1/2 x^T Q x + c^T x  s.t.  A x <= b [, A_eq x = b_eq]``.

    First-order homotopy: accelerated gradient descent on the tempered penalty
    with ``beta`` / ``mu`` grown geometrically. ``x0`` may be **any** point (the
    exterior penalty needs no feasible start); it defaults to the origin. Optional
    equality constraints ``A_eq x = b_eq`` are enforced by a quadratic exterior
    penalty. Returns a :class:`~omnibias.convex.problem.ConvexSolution` with the
    inequality dual ``lambda = mu sigma(beta (A x - b)) >= 0``, the equality dual
    estimate ``nu = mu (A_eq x - b_eq)`` in ``eq_dual`` (``None`` if no equalities),
    ``iterations`` = homotopy stages run and ``newton_iterations`` = total steps.
    """
    opts = options or PenaltyOptions()
    Qa = jnp.asarray(Q)
    ca = jnp.asarray(c)
    Aa = jnp.asarray(A)
    ba = jnp.asarray(b)
    n = ca.shape[0]
    m = ba.shape[0]
    has_eq = A_eq is not None and b_eq is not None
    Aeq = jnp.asarray(A_eq, dtype=ba.dtype) if has_eq else None
    beq = jnp.asarray(b_eq, dtype=ba.dtype) if has_eq else None
    validate_shapes(
        n, m, A_shape=Aa.shape, b_shape=ba.shape, c_shape=ca.shape, Q_shape=Qa.shape,
        A_eq_shape=None if Aeq is None else Aeq.shape,
        b_eq_shape=None if beq is None else beq.shape,
    )

    x = jnp.zeros((n,), dtype=ba.dtype) if x0 is None else jnp.asarray(x0, dtype=ba.dtype)
    beta = opts.beta0
    mu = opts.penalty0
    beta_used = beta
    mu_used = mu
    total_steps = 0
    stages_run = 0
    dx = float("inf")
    viol = float("inf")
    for _ in range(opts.stages):
        beta_used = beta
        mu_used = mu
        x_prev = x
        x = penalty_descent(
            Qa, ca, Aa, ba, x_prev,
            beta=beta, mu=mu, steps=opts.gd_steps, step_safety=opts.step_safety,
            prox=opts.prox, anchor=x_prev, A_eq=Aeq, b_eq=beq,
        )
        total_steps += opts.gd_steps
        stages_run += 1
        dx = float(jnp.max(jnp.abs(x - x_prev)))
        viol = float(jnp.max(jnp.maximum(Aa @ x - ba, 0.0)))
        if has_eq:
            if Aeq is None or beq is None:
                raise ValueError(
                    "equality constraints require both A_eq and b_eq; "
                    f"got A_eq={'set' if Aeq is not None else None}, "
                    f"b_eq={'set' if beq is not None else None}"
                )
            viol = max(viol, float(jnp.max(jnp.abs(Aeq @ x - beq))))
        if dx <= opts.tol and viol <= opts.feas_tol:
            break
        beta *= opts.beta_growth
        mu *= opts.penalty_growth

    converged = viol <= opts.feas_tol
    u = Aa @ x - ba
    dual = mu_used * jax.nn.sigmoid(beta_used * u)
    eq_dual = None
    if has_eq:
        if Aeq is None or beq is None:
            raise ValueError(
                "equality constraints require both A_eq and b_eq; "
                f"got A_eq={'set' if Aeq is not None else None}, "
                f"b_eq={'set' if beq is not None else None}"
            )
        eq_dual = mu_used * (Aeq @ x - beq)
    slack = ba - Aa @ x
    obj = 0.5 * x @ (Qa @ x) + ca @ x
    return ConvexSolution(
        x=x,
        dual=dual,
        slack=slack,
        obj=obj,
        gap=max(viol, dx),
        iterations=stages_run,
        converged=converged,
        newton_iterations=total_steps,
        eq_dual=eq_dual,
    )

Related APIsomnibias.convex.jax.penalty_descent↗omnibias.convex.ConvexSolution↗omnibias.convex.problem.PenaltyOptions↗omnibias.convex.problem.validate_shapes↗

Other public import paths

• omnibias.convex.jax.solve_qp_penalty
• omnibias.convex.jax.penalty.solve_qp_penalty

← Package overviewRunnable learning examples →