API REFERENCE/omnibias-jax

omnibias.jax.implicit

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.implicit

Module contract & mathematical conventions

Implicit / DEQ Newton (theory 08-08), JAX twin.

Solve u = sigma(W u + x) to a named residual, then differentiate the root by the implicit-function theorem with exact sigma' (bias collapse, delta -> 0). One linear solve replaces unrolled backprop through a fixed-point iteration.

The default solver loop is lax.while_loop (G4). A Python while is used only when require_contraction=True so a bound >= 1 can raise instead of silently unrolling. Do not wrap a data-dependent Python break in jax.jit. IFT is the chain rule at a fixed point, not an absence of the chain rule. The fixed point is local. Enable 64-bit JAX before the first array for torch parity.

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 →

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
DEQConfig	omnibias.core.DEQConfig
DEQNotContractive	omnibias.core.DEQNotContractive
DEQSolverUnknown	omnibias.core.DEQSolverUnknown
honesty_payload	omnibias.core.honesty_payload

CLASS

DEQResult

omnibias.jax.DEQResultSource ↗

class DEQResult

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Fixed point, residual, iteration count, and contraction bound.

Parameters

Name	Type & default	Contract
udataclass field	ArrayRequired	See the operation contract and implementation below.
residualdataclass field	floatRequired	See the operation contract and implementation below.
n_iterdataclass field	intRequired	See the operation contract and implementation below.
spectral_radius_bounddataclass field	floatRequired	See the operation contract and implementation below.

Fields & methods

ufield

u: Array

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

View source ↗

residualfield

residual: float

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

View source ↗

n_iterfield

n_iter: int

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

View source ↗

spectral_radius_boundfield

spectral_radius_bound: float

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

View source ↗

Read the implementation

class DEQResult:
    """Fixed point, residual, iteration count, and contraction bound."""

    u: Array
    residual: float
    n_iter: int
    spectral_radius_bound: float

Other public import paths

• omnibias.jax.DEQResult
• omnibias.jax.implicit.DEQResult

FUNCTION

deq_du_dW

omnibias.jax.deq_du_dWSource ↗

deq_du_dW(W: Array, x: Array, spec: str='tanh', *, config: DEQConfig | None=None, u: Array | None=None) -> Array

IFT Jacobian du*/dW with shape (..., d, d, d).

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
specpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
configkeyword-only	DEQConfig | NoneDefault: None	See the operation contract and implementation below.
ukeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def deq_du_dW(
    W: Array,
    x: Array,
    spec: str = "tanh",
    *,
    config: DEQConfig | None = None,
    u: Array | None = None,
) -> Array:
    """IFT Jacobian ``du*/dW`` with shape ``(..., d, d, d)``."""
    cfg = config if config is not None else DEQConfig()
    w = _as_W(W)
    x_t = _as_x(x, int(w.shape[0])).astype(w.dtype)
    if u is None:
        u_star = deq_solve(w, x_t, spec, config=cfg).u
    else:
        u_star = jnp.asarray(u, dtype=w.dtype)
    sigp = _linear_map(u_star, w, spec, x_t)
    d = int(w.shape[0])
    eye = jnp.eye(d, dtype=w.dtype)
    if u_star.ndim == 1:
        a = eye - sigp.reshape((d, 1)) * w
        ainv = jnp.linalg.inv(a)
        return cast(
            Array,
            ainv[..., None] * sigp.reshape((1, d, 1)) * u_star.reshape((1, 1, d)),
        )
    n = int(u_star.shape[0])
    a = eye[None, :, :] - sigp.reshape((n, d, 1)) * w[None, :, :]
    ainv = jnp.linalg.inv(a)
    return cast(
        Array,
        ainv[..., None]
        * sigp.reshape((n, 1, d, 1))
        * u_star.reshape((n, 1, 1, d)),
    )

Related APIsomnibias.core.DEQConfig↗omnibias.jax.deq_solve↗

Other public import paths

• omnibias.jax.deq_du_dW
• omnibias.jax.implicit.deq_du_dW

FUNCTION

deq_solve

omnibias.jax.deq_solveSource ↗

deq_solve(W: Array, x: Array, spec: str='tanh', *, config: DEQConfig | None=None) -> DEQResult

Solve u = sigma(W u + x) by Newton or Banach iteration.

jax deq_solve uses lax.while_loop for the Newton / Banach loop (G4). require_contraction=True falls back to a Python while so a bound >= 1 can raise. The IFT VJP is one linear solve, not unrolled BPTT.

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
specpositional-or-keyword	strDefault: 'tanh'	See the operation contract and implementation below.
configkeyword-only	DEQConfig | NoneDefault: None	See the operation contract and implementation below.

ReturnsDEQResult

Executed regression · contraction 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.

Full file, imports & fixtures ↗

def test_contraction_raises() -> None:
    cfg = DEQConfig(require_contraction=True, max_iter=2)
    with pytest.raises(DEQNotContractive, match="unroll"):
        deq_solve(jnp.asarray([[3.0]]), jnp.asarray([0.1]), "tanh", config=cfg)

python -m pytest "packages/omnibias-jax/tests/test_implicit.py::test_contraction_raises" -q

Regression example · anderson 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.

Full file, imports & fixtures ↗

def test_anderson_raises() -> None:
    with pytest.raises(DEQSolverUnknown, match="newton"):
        deq_solve(
            jnp.asarray([[0.2]]),
            jnp.asarray([0.5]),
            "tanh",
            config=DEQConfig(solver="anderson"),
        )

python -m pytest packages/omnibias-jax/tests/test_implicit.py -k test_anderson_raises -q

Read the implementation

def deq_solve(
    W: Array,
    x: Array,
    spec: str = "tanh",
    *,
    config: DEQConfig | None = None,
) -> DEQResult:
    """Solve ``u = sigma(W u + x)`` by Newton or Banach iteration.

    jax deq_solve uses lax.while_loop for the Newton / Banach loop
    (G4). require_contraction=True falls back to a Python while so a
    bound >= 1 can raise. The IFT VJP is one linear solve, not
    unrolled BPTT.
    """
    cfg = config if config is not None else DEQConfig()
    reject_anderson(cfg.solver)
    w = _as_W(W)
    x_t = _as_x(x, int(w.shape[0])).astype(w.dtype)
    u = jnp.zeros_like(x_t)
    if cfg.require_contraction:
        bound = _maybe_contract(u, w, x_t, spec, cfg)
        res = float(jnp.asarray(_residual_norm(_residual(u, w, x_t, spec))).reshape(-1)[0])
        n_iter = 0
        while res > cfg.tol and n_iter < cfg.max_iter:
            u = _step(u, w, x_t, spec, cfg.solver)
            bound = _maybe_contract(u, w, x_t, spec, cfg)
            res = float(
                jnp.asarray(_residual_norm(_residual(u, w, x_t, spec))).reshape(-1)[0]
            )
            n_iter += 1
        return DEQResult(
            u=u, residual=res, n_iter=n_iter, spectral_radius_bound=bound
        )
    u, n_i, res_a = _while_loop_solve(
        w, x_t, spec, cfg.solver, cfg.max_iter, cfg.tol, u
    )
    bound = _maybe_contract(u, w, x_t, spec, cfg)
    return DEQResult(
        u=u,
        residual=float(jnp.asarray(res_a).reshape(-1)[0]),
        n_iter=int(jnp.asarray(n_i).reshape(-1)[0]),
        spectral_radius_bound=bound,
    )

Related APIsomnibias.core.DEQConfig↗omnibias.core.reject_anderson↗omnibias.jax.DEQResult↗

Other public import paths

• omnibias.jax.deq_solve
• omnibias.jax.implicit.deq_solve

FUNCTION

deq_vjp

omnibias.jax.deq_vjpSource ↗

deq_vjp(W: Array, x: Array, spec: str, g: Array, *, config: DEQConfig | None=None, u: Array | None=None) -> Array

VJP of u* with respect to W through the IFT.

g is the cotangent of u* (same shape as u). The backward linear solve uses exact sigma', not unrolled BPTT.

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
specpositional-or-keyword	strRequired	See the operation contract and implementation below.
gpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
configkeyword-only	DEQConfig | NoneDefault: None	See the operation contract and implementation below.
ukeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When u_star.shape != x_t.shape, raises ValueError(f'u shape {tuple(u_star.shape)} must match x shape {tuple(x_t.shape)}').
• When g_t.shape != u_star.shape, raises ValueError(f'g shape {tuple(g_t.shape)} must match u shape {tuple(u_star.shape)}').

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

Executed regression · g1 ift matches central fd

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_g1_ift_matches_central_fd() -> None:
    cfg = DEQConfig(solver="newton", tol=1e-14, max_iter=30)
    w = jnp.asarray([[0.2]])
    x = jnp.asarray([0.5])
    vjp = float(deq_vjp(w, x, "tanh", jnp.asarray([1.0]), config=cfg).reshape(-1)[0])
    eps = 1e-6
    up = float(deq_solve(w + eps, x, "tanh", config=cfg).u.reshape(-1)[0])
    um = float(deq_solve(w - eps, x, "tanh", config=cfg).u.reshape(-1)[0])
    fd = (up - um) / (2.0 * eps)
    rel = abs(vjp - fd) / max(abs(fd), 1e-12)
    assert rel <= 1e-8
    assert deq_solve(w, x, "tanh", config=cfg).residual <= 1e-10

python -m pytest "packages/omnibias-jax/tests/test_implicit.py::test_g1_ift_matches_central_fd" -q

Regression example · g3 width2 parity

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_g3_width2_parity() -> None:
    torch.set_default_dtype(torch.float64)
    cfg = DEQConfig(solver="newton", tol=1e-14, max_iter=40)
    w = [[0.15, -0.05], [0.08, 0.12]]
    x = [0.4, -0.2]
    g = [1.0, 0.0]
    t_res = torch_solve(torch.tensor(w), torch.tensor(x), "tanh", config=cfg)
    j_res = jax_solve(jnp.asarray(w), jnp.asarray(x), "tanh", config=cfg)
    t_g = torch_vjp(
        torch.tensor(w), torch.tensor(x), "tanh", torch.tensor(g), config=cfg
    )
    j_g = jax_vjp(jnp.asarray(w), jnp.asarray(x), "tanh", jnp.asarray(g), config=cfg)
    assert abs(float(t_res.u[0]) - float(j_res.u.reshape(-1)[0])) <= 1e-12
    assert abs(float(t_res.u[1]) - float(j_res.u.reshape(-1)[1])) <= 1e-12
    assert float((t_g - torch.tensor(j_g.tolist())).abs().max()) <= 1e-12

python -m pytest tests/test_implicit_parity.py -k test_g3_width2_parity -q

Read the implementation

def deq_vjp(
    W: Array,
    x: Array,
    spec: str,
    g: Array,
    *,
    config: DEQConfig | None = None,
    u: Array | None = None,
) -> Array:
    """VJP of ``u*`` with respect to ``W`` through the IFT.

    ``g`` is the cotangent of ``u*`` (same shape as ``u``). The
    backward linear solve uses exact ``sigma'``, not unrolled BPTT.
    """
    cfg = config if config is not None else DEQConfig()
    w = _as_W(W)
    x_t = _as_x(x, int(w.shape[0])).astype(w.dtype)
    if u is None:
        u_star = deq_solve(w, x_t, spec, config=cfg).u
    else:
        u_star = jnp.asarray(u, dtype=w.dtype)
        if u_star.shape != x_t.shape:
            raise ValueError(
                f"u shape {tuple(u_star.shape)} must match x shape {tuple(x_t.shape)}"
            )
    g_t = jnp.asarray(g, dtype=w.dtype)
    if g_t.shape != u_star.shape:
        raise ValueError(
            f"g shape {tuple(g_t.shape)} must match u shape {tuple(u_star.shape)}"
        )
    sigp = _linear_map(u_star, w, spec, x_t)
    if u_star.ndim == 1:
        d = int(w.shape[0])
        eye = jnp.eye(d, dtype=w.dtype)
        at = eye - w.T * sigp.reshape((1, d))
        v = jnp.linalg.solve(at, g_t)
        return jnp.outer(v * sigp, u_star)
    n = int(u_star.shape[0])
    d = int(w.shape[0])
    eye = jnp.eye(d, dtype=w.dtype)
    at = eye[None, :, :] - w.T[None, :, :] * sigp.reshape((n, 1, d))
    v = jnp.linalg.solve(at, g_t)
    return jnp.einsum("bi,bj->ij", v * sigp, u_star)

Related APIsomnibias.core.DEQConfig↗omnibias.jax.deq_solve↗

Other public import paths

• omnibias.jax.deq_vjp
• omnibias.jax.implicit.deq_vjp

← Package overviewRunnable learning examples →