Skip to main content

API REFERENCE/omnibias-curvature

omnibias.curvature.operators

0.1.0a2AlphaPython >=3.10AGPL / commercial

Scoped parameter-curvature formulas, Fisher and Gauss–Newton constructions, natural steps, and sharpness utilities.

import omnibias.curvature.operators

Module contract & mathematical conventions

Matrix-free symmetric solves with residual diagnostics and implicit gradients.

PCG is numerical. custom_linear_solve differentiates its defining linear system, so gradients are meaningful only when the reported solve converges.

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 →

CONSTANT

MatVec

omnibias.curvature.operators.MatVecSource ↗

MatVec = Callable[[Array], Array]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

MatVec = Callable[[Array], Array]

CLASS

PCGResult

omnibias.curvature.operators.PCGResultSource ↗

class PCGResult(NamedTuple)

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Fields & methods

solutionfield

solution: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

residual_normfield

residual_norm: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

relative_residualfield

relative_residual: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

iterationsfield

iterations: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

convergedfield

converged: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

breakdownfield

breakdown: Array

Declared Array field on PCGResult. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

Read the implementation

class PCGResult(NamedTuple):
    solution: Array
    residual_norm: Array
    relative_residual: Array
    iterations: Array
    converged: Array
    breakdown: Array

FUNCTION

pcg_solve

omnibias.curvature.operators.pcg_solveSource ↗

pcg_solve(matvec: MatVec, rhs: Array, *, preconditioner: MatVec | None=None, rtol: float=1e-08, atol: float=0.0, max_iterations: int=200) -> PCGResult

Solve A x=b for a user-supplied SPD operator without assembling A.

The preconditioner applies an SPD approximate inverse. Invalid curvature, nonfinite iterates and exhausted budgets are explicit nonconvergence. Options are static under JIT. Reverse derivatives use the symmetric implicit solve rather than differentiating an adaptive iteration count.

Parameters

Name	Type & default	Contract
`matvec`positional-or-keyword	`MatVec`Required	See the operation contract and implementation below.
`rhs`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`preconditioner`keyword-only	`MatVec \| None`Default: None	See the operation contract and implementation below.
`rtol`keyword-only	`float`Default: 1e-08	See the operation contract and implementation below.
`atol`keyword-only	`float`Default: 0.0	See the operation contract and implementation below.
`max_iterations`keyword-only	`int`Default: 200	See the operation contract and implementation below.

ReturnsPCGResult

Explicit validation & failure cases

When rhs.ndim != 1 or rhs.size == 0 or jnp.issubdtype(rhs.dtype, jnp.complexfloating), raises ValueError('rhs must be a nonempty real vector').
When rtol < 0 or atol < 0 or (not isfinite(rtol + atol)) or (max_iterations < 1), raises ValueError('finite nonnegative tolerances and positive iteration budget required').

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

Executed regression · pcg failure and zero rhs

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_pcg_failure_and_zero_rhs():
    bad = pcg_solve(lambda x: -x, jnp.ones(3))
    assert bad.breakdown and not bad.converged
    zero = pcg_solve(lambda x: x, jnp.zeros(3))
    assert zero.converged and zero.iterations == 0
    exhausted = pcg_solve(
        lambda x: jnp.arange(1.0, 11.0) * x, jnp.ones(10), max_iterations=1, rtol=1e-14
    )
    assert not exhausted.converged
    for invalid in (float("inf"), float("nan")):
        with pytest.raises(ValueError, match="finite"):
            pcg_solve(lambda x: x, jnp.ones(3), rtol=invalid)

python -m pytest "packages/omnibias-curvature/tests/test_matrixfree_operators.py::test_pcg_failure_and_zero_rhs" -q

Regression example · pcg solution jit and implicit gradient

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_pcg_solution_jit_and_implicit_gradient():
    a = jnp.array([[4.0, 1.0], [1.0, 3.0]])
    b = jnp.array([1.0, 2.0])

    def solve(b):
        return pcg_solve(lambda x: a @ x, b, rtol=1e-12)

    out = jax.jit(solve)(b)
    assert out.converged and not out.breakdown
    np.testing.assert_allclose(out.solution, np.linalg.solve(a, b), rtol=1e-12)
    jac = jax.jacrev(lambda b: solve(b).solution)(b)
    np.testing.assert_allclose(jac, np.linalg.inv(a), atol=1e-12)
    grad = jax.grad(lambda scale: jnp.sum(pcg_solve(lambda x: scale * (a @ x), b).solution))(2.0)
    assert np.allclose(grad, -np.sum(np.linalg.solve(a, b)) / 4)

python -m pytest packages/omnibias-curvature/tests/test_matrixfree_operators.py -k test_pcg_solution_jit_and_implicit_gradient -q

Read the implementation

def pcg_solve(
    matvec: MatVec,
    rhs: Array,
    *,
    preconditioner: MatVec | None = None,
    rtol: float = 1e-8,
    atol: float = 0.0,
    max_iterations: int = 200,
) -> PCGResult:
    """Solve A x=b for a user-supplied SPD operator without assembling A.

    The preconditioner applies an SPD approximate inverse. Invalid curvature,
    nonfinite iterates and exhausted budgets are explicit nonconvergence.
    Options are static under JIT. Reverse derivatives use the symmetric
    implicit solve rather than differentiating an adaptive iteration count.
    """
    if rhs.ndim != 1 or rhs.size == 0 or jnp.issubdtype(rhs.dtype, jnp.complexfloating):
        raise ValueError("rhs must be a nonempty real vector")
    if rtol < 0 or atol < 0 or not isfinite(rtol + atol) or max_iterations < 1:
        raise ValueError("finite nonnegative tolerances and positive iteration budget required")
    apply_m: MatVec = (lambda x: x) if preconditioner is None else preconditioner

    def solve(op: MatVec, b: Array) -> tuple[Array, tuple[Array, Array]]:
        x = jnp.zeros_like(b)
        r = b
        z = apply_m(r)
        p = z
        rz = jnp.vdot(r, z).real
        target = jnp.maximum(atol, rtol * jnp.linalg.norm(b))
        invalid = (~jnp.all(jnp.isfinite(b))) | (~jnp.all(jnp.isfinite(z)))
        initial: State = (jnp.asarray(0), x, r, p, rz, invalid)

        def cond(state: State) -> Array:
            k, _, res, _, _, bad = state
            return jnp.asarray((k < max_iterations) & (jnp.linalg.norm(res) > target) & (~bad))

        def body(state: State) -> State:
            k, x, r, p, rz, bad = state
            ap = op(p)
            denominator = jnp.vdot(p, ap).real
            valid = (denominator > 0) & (rz > 0) & jnp.isfinite(denominator) & jnp.isfinite(rz)
            alpha = jnp.where(valid, rz / jnp.where(valid, denominator, 1.0), 0.0)
            xn = x + alpha * p
            rn = r - alpha * ap
            zn = apply_m(rn)
            rzn = jnp.vdot(rn, zn).real
            beta = jnp.where(valid, rzn / jnp.where(valid, rz, 1.0), 0.0)
            pn = zn + beta * p
            bad = bad | (~valid) | (~jnp.all(jnp.isfinite(xn))) | (~jnp.all(jnp.isfinite(rn)))
            return k + 1, xn, rn, pn, rzn, bad

        final = jax.lax.while_loop(cond, body, initial)
        return final[1], (final[0], final[-1])

    x, (iterations, breakdown) = jax.lax.custom_linear_solve(
        matvec, rhs, solve=solve, symmetric=True, has_aux=True
    )
    residual = jnp.linalg.norm(matvec(x) - rhs)
    norm = jnp.linalg.norm(rhs)
    relative = residual / jnp.where(norm > 0, norm, 1.0)
    converged = (residual <= jnp.maximum(atol, rtol * norm)) & (~breakdown) & jnp.isfinite(residual)
    return PCGResult(x, residual, relative, iterations, converged, breakdown)

← Package overview Runnable learning examples →

MatVec
PCGResult
pcg_solve