API REFERENCE/omnibias-torch

omnibias.torch.implicit

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.implicit

Module contract & mathematical conventions

Implicit / DEQ Newton (theory 08-08), PyTorch 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 solver loop is a Python while (do not wrap in torch.compile). JAX uses lax.while_loop; see the twin docstring. IFT is the chain rule at a fixed point, not an absence of the chain rule. Not a global min

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.torch.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	TensorRequired	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: Tensor

Declared Tensor 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: Tensor
    residual: float
    n_iter: int
    spectral_radius_bound: float

Other public import paths

• omnibias.torch.DEQResult
• omnibias.torch.implicit.DEQResult

FUNCTION

deq_du_dW

omnibias.torch.deq_du_dWSource ↗

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

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

Indexing is [u_component, W_row, W_col] after any batch axis.

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
xpositional-or-keyword	TensorRequired	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	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · g2 implicit residual skill

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_g2_implicit_residual_skill() -> None:
    torch.set_default_dtype(torch.float64)
    from omnibias.torch.implicit import deq_du_dW
    from omnibias.torch.optim import gauss_newton_direction

    xs = torch.linspace(0.0, 1.0, 21)
    target = torch.sin(math.pi * xs)
    x_inj = torch.stack([xs, torch.ones_like(xs)], dim=-1)
    cfg = DEQConfig(tol=1e-12, max_iter=30)
    zero_mse = float((target**2).mean())
    mses: list[float] = []
    for seed in range(5):
        gen = torch.Generator().manual_seed(seed)
        w = 0.15 * torch.randn(2, 2, generator=gen)
        sol = deq_solve(w, x_inj, "tanh", config=cfg)
        best = float(((sol.u[:, 0] - target) ** 2).mean())
        for _ in range(15):
            res = sol.u[:, 0] - target
            jac = deq_du_dW(w, x_inj, "tanh", config=cfg, u=sol.u)
            j = jac[:, 0, :, :].reshape(xs.numel(), 4)
            trial = w + gauss_newton_direction(j, res, 1e-3).reshape(2, 2)
            trial_sol = deq_solve(trial, x_inj, "tanh", config=cfg)
            trial_mse = float(((trial_sol.u[:, 0] - target) ** 2).mean())
            if trial_mse <= best:
                w, sol, best = trial, trial_sol, trial_mse
        mses.append(best)
    assert all(m < zero_mse for m in mses)
    assert statistics.mean(mses) <= 0.20
    assert max(mses) <= 0.25

python -m pytest "packages/omnibias-torch/tests/test_implicit.py::test_g2_implicit_residual_skill" -q

Read the implementation

def deq_du_dW(
    W: Tensor,
    x: Tensor,
    spec: str = "tanh",
    *,
    config: DEQConfig | None = None,
    u: Tensor | None = None,
) -> Tensor:
    """IFT Jacobian ``du*/dW`` with shape ``(..., d, d, d)``.

    Indexing is ``[u_component, W_row, W_col]`` after any batch axis.
    """
    cfg = config if config is not None else DEQConfig()
    w = _as_W(W)
    x_t = _as_x(x, int(w.shape[0])).to(dtype=w.dtype, device=w.device)
    if u is None:
        u_star = deq_solve(w, x_t, spec, config=cfg).u
    else:
        u_star = torch.as_tensor(u, dtype=w.dtype, device=w.device)
    sigp = _linear_map(u_star, w, spec, x_t)
    d = int(w.shape[0])
    eye = torch.eye(d, dtype=w.dtype, device=w.device)
    if u_star.ndim == 1:
        a = eye - sigp.reshape(d, 1) * w
        ainv = torch.linalg.inv(a)
        return cast(
            Tensor,
            ainv.unsqueeze(-1) * sigp.reshape(1, d, 1) * u_star.reshape(1, 1, d),
        )
    n = int(u_star.shape[0])
    a = eye.unsqueeze(0) - sigp.reshape(n, d, 1) * w.unsqueeze(0)
    ainv = torch.linalg.inv(a)
    return cast(
        Tensor,
        ainv.unsqueeze(-1) * sigp.reshape(n, 1, d, 1) * u_star.reshape(n, 1, 1, d),
    )

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

Other public import paths

• omnibias.torch.deq_du_dW
• omnibias.torch.implicit.deq_du_dW

FUNCTION

deq_solve

omnibias.torch.deq_solveSource ↗

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

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

torch deq_solve uses a Python while loop; do not wrap it in torch.compile. The IFT VJP is one linear solve, not unrolled BPTT.

Parameters

Name	Type & default	Contract
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
xpositional-or-keyword	TensorRequired	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:
    torch.set_default_dtype(torch.float64)
    cfg = DEQConfig(require_contraction=True, max_iter=2)
    with pytest.raises(DEQNotContractive, match="unroll"):
        deq_solve(torch.tensor([[3.0]]), torch.tensor([0.1]), "tanh", config=cfg)

python -m pytest "packages/omnibias-torch/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:
    torch.set_default_dtype(torch.float64)
    with pytest.raises(DEQSolverUnknown, match="newton"):
        deq_solve(
            torch.tensor([[0.2]]),
            torch.tensor([0.5]),
            "tanh",
            config=DEQConfig(solver="anderson"),
        )

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

Read the implementation

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

    torch deq_solve uses a Python while loop; do not wrap it in
    torch.compile. 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])).to(dtype=w.dtype, device=w.device)
    u = torch.zeros_like(x_t)
    bound = _maybe_contract(u, w, x_t, spec, cfg)
    res = _residual_norm(_residual(u, w, x_t, spec))
    n_iter = 0
    while res > cfg.tol and n_iter < cfg.max_iter:
        if cfg.solver == "newton":
            u = _newton_step(u, w, x_t, spec)
        else:
            u = _iterate_step(u, w, x_t, spec)
        bound = _maybe_contract(u, w, x_t, spec, cfg)
        res = _residual_norm(_residual(u, w, x_t, spec))
        n_iter += 1
    return DEQResult(
        u=u,
        residual=res,
        n_iter=n_iter,
        spectral_radius_bound=bound,
    )

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

Other public import paths

• omnibias.torch.deq_solve
• omnibias.torch.implicit.deq_solve

FUNCTION

deq_vjp

omnibias.torch.deq_vjpSource ↗

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

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	TensorRequired	See the operation contract and implementation below.
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	strRequired	See the operation contract and implementation below.
gpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
configkeyword-only	DEQConfig | NoneDefault: None	See the operation contract and implementation below.
ukeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

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:
    torch.set_default_dtype(torch.float64)
    cfg = DEQConfig(solver="newton", tol=1e-14, max_iter=30)
    w = torch.tensor([[0.2]])
    x = torch.tensor([0.5])
    g = torch.tensor([1.0])
    vjp = float(deq_vjp(w, x, "tanh", g, 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-torch/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: Tensor,
    x: Tensor,
    spec: str,
    g: Tensor,
    *,
    config: DEQConfig | None = None,
    u: Tensor | None = None,
) -> Tensor:
    """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])).to(dtype=w.dtype, device=w.device)
    if u is None:
        u_star = deq_solve(w, x_t, spec, config=cfg).u
    else:
        u_star = torch.as_tensor(u, dtype=w.dtype, device=w.device)
        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 = torch.as_tensor(g, dtype=w.dtype, device=w.device)
    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)
    # A^T v = g  with A = I - D W, A^T = I - W^T D
    if u_star.ndim == 1:
        d = int(w.shape[0])
        eye = torch.eye(d, dtype=w.dtype, device=w.device)
        at = eye - w.T * sigp.reshape(1, d)
        v = torch.linalg.solve(at, g_t)
        return torch.outer(v * sigp, u_star)
    n = int(u_star.shape[0])
    d = int(w.shape[0])
    eye = torch.eye(d, dtype=w.dtype, device=w.device)
    at = eye.unsqueeze(0) - w.T.unsqueeze(0) * sigp.reshape(n, 1, d)
    v = torch.linalg.solve(at, g_t)
    return torch.einsum("bi,bj->ij", v * sigp, u_star)

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

Other public import paths

• omnibias.torch.deq_vjp
• omnibias.torch.implicit.deq_vjp

← Package overviewRunnable learning examples →