API REFERENCE/omnibias-torch

omnibias.torch.jet

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.jet

Module contract & mathematical conventions

Exact multi-layer directional jets via Faà di Bruno (torch).

Bit-identical twin of omnibias.jax.jet; see that module for the full mathematical exposition. A *jet* stores Taylor coefficients a_k = f^(k)(0)/k! along the leading axis of a tensor of shape (N+1, ...). The activation composition uses the numerically-stable shifted-power identity

\sigma(u(t)) = \sum_{k=0}^{N} \frac{\sigma^{(k)}(u_0)}{k!}\,(u(t)-u_0)^k,

with series powers built by truncated convolution and exact sigma^(k) from the omnibias activation fast paths.

Truncation-order cost (no O(N^2) claim for the general kernel) is discussed in omnibias.jax.jet; the two modules run the same recurrences.

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

affine_jet

omnibias.torch.affine_jetSource ↗

affine_jet(z_jet: Tensor, W: Tensor, b: Tensor | None=None) -> Tensor

Push a jet through an affine map u = W z + b.

z_jet has shape (N+1, ..., D_in); W has shape (D_out, D_in); b (optional) has shape (D_out,). Linear in t, so it acts per order with the bias added to the zeroth coefficient only.

Parameters

Name	Type & default	Contract
z_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · hessian requires order two

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_hessian_requires_order_two() -> None:
    layers, x0 = _build_mlp(seed=6)
    jet = mlp_jet_mv(x0, layers, 1)
    with pytest.raises(ValueError, match="hessian needs order >= 2"):
        jet_hessian(jet, int(x0.shape[0]), 1)

python -m pytest "packages/omnibias-torch/tests/test_jet_mv.py::test_hessian_requires_order_two" -q

Regression example · affine jet is linear per order

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_affine_jet_is_linear_per_order() -> None:
    rng = np.random.default_rng(7)
    order = 4
    z_jet = torch.as_tensor(rng.normal(size=(order + 1, 5)))
    W = torch.as_tensor(rng.normal(size=(3, 5)))
    b = torch.as_tensor(rng.normal(size=(3,)))
    out = affine_jet(z_jet, W, b)
    assert torch.allclose(out[0], W @ z_jet[0] + b)
    for k in range(1, order + 1):
        assert torch.allclose(out[k], W @ z_jet[k])

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_affine_jet_is_linear_per_order -q

Read the implementation

def affine_jet(z_jet: Tensor, W: Tensor, b: Tensor | None = None) -> Tensor:
    """Push a jet through an affine map ``u = W z + b``.

    ``z_jet`` has shape ``(N+1, ..., D_in)``; ``W`` has shape ``(D_out, D_in)``;
    ``b`` (optional) has shape ``(D_out,)``. Linear in ``t``, so it acts per
    order with the bias added to the zeroth coefficient only.
    """
    z_jet = torch.as_tensor(z_jet)
    W = torch.as_tensor(W)
    out: Tensor = torch.tensordot(z_jet, W, dims=([-1], [-1]))  # (N+1, ..., D_out)
    if b is not None:
        b = torch.as_tensor(b)
        out0 = (out[0] + b).unsqueeze(0)
        out = torch.cat([out0, out[1:]], dim=0)
    return out

Other public import paths

• omnibias.torch.affine_jet
• omnibias.torch.jet.affine_jet

FUNCTION

affine_joint_jet

omnibias.torch.jet.affine_joint_jetSource ↗

affine_joint_jet(input_jet: Tensor, weight_jet: Tensor, bias_jet: Tensor | None=None) -> Tensor

Live directional product with jet-valued weights and biases.

Parameters

Name	Type & default	Contract
input_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
weight_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bias_jetpositional-or-keyword	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def affine_joint_jet(input_jet: Tensor, weight_jet: Tensor, bias_jet: Tensor | None = None) -> Tensor:
    """Live directional product with jet-valued weights and biases."""
    from omnibias.torch.realization import affine_joint_jet as impl

    return impl(input_jet, weight_jet, bias_jet)

FUNCTION

antiderivative_jet

omnibias.torch.antiderivative_jetSource ↗

antiderivative_jet(jet: Tensor, constant: float=0.0) -> Tensor

Taylor jet of F(t) = constant + \int_0^t f (term-by-term integration).

A_0 = constant; A_m = jet[m-1]/m for m >= 1, so the jet lengthens by one order. Works on scalar and vector jets; the two-sided partner of jet_to_tower and the exact float inverse of derivative_jet. The constant of integration is a free parameter (default 0).

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
constantpositional-or-keyword	floatDefault: 0.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · antiderivative jet ftc part1 roundtrip

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_antiderivative_jet_ftc_part1_roundtrip() -> None:
    """FTC part 1: d/dt of the antiderivative jet recovers the original vector jet."""
    layers, x0, v = _build_mlp(seed=4)
    order = 6
    jet = mlp_jet(x0, v, layers, order)  # (order+1, C)
    anti = antiderivative_jet(jet, constant=0.3)
    assert anti.shape[0] == jet.shape[0] + 1
    assert torch.allclose(anti[0], torch.full_like(jet[0], 0.3))
    back = derivative_jet(anti)
    assert back.shape == jet.shape
    assert torch.allclose(back, jet, rtol=1e-12, atol=1e-12)

python -m pytest "packages/omnibias-torch/tests/test_jet.py::test_antiderivative_jet_ftc_part1_roundtrip" -q

Regression example · antiderivative jet torch 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_antiderivative_jet_torch_parity() -> None:
    """Bit-identical antiderivative / derivative jets across the jax and torch twins."""
    torch = pytest.importorskip("torch")
    from omnibias.torch.jet import antiderivative_jet as t_anti
    from omnibias.torch.jet import derivative_jet as t_deriv

    rng = np.random.default_rng(13)
    arr = rng.normal(size=(6, 4))
    jx = jnp.asarray(arr)
    tx = torch.as_tensor(arr, dtype=torch.float64)

    aj = np.asarray(antiderivative_jet(jx, constant=0.7))
    at = t_anti(tx, constant=0.7).numpy()
    assert np.array_equal(aj, at)

    dj = np.asarray(derivative_jet(jx))
    dt = t_deriv(tx).numpy()
    assert np.array_equal(dj, dt)

python -m pytest packages/omnibias-jax/tests/test_jet.py -k test_antiderivative_jet_torch_parity -q

Read the implementation

def antiderivative_jet(jet: Tensor, constant: float = 0.0) -> Tensor:
    r"""Taylor jet of ``F(t) = constant + \int_0^t f`` (term-by-term integration).

    ``A_0 = constant``; ``A_m = jet[m-1]/m`` for ``m >= 1``, so the jet lengthens by
    one order. Works on scalar and vector jets; the two-sided partner of
    :func:`jet_to_tower` and the exact float inverse of :func:`derivative_jet`.
    The constant of integration is a free parameter (default 0).
    """
    jet = torch.as_tensor(jet)
    np1 = jet.shape[0]
    divisors = torch.arange(1, np1 + 1, dtype=jet.dtype, device=jet.device)
    scaled = _scale_axis0(jet, 1.0 / divisors)
    const_row = torch.full_like(jet[0], float(constant)).unsqueeze(0)
    return torch.cat([const_row, scaled], dim=0)

Other public import paths

• omnibias.torch.antiderivative_jet
• omnibias.torch.jet.antiderivative_jet

FUNCTION

compose_jet

omnibias.torch.compose_jetSource ↗

compose_jet(u_jet: Tensor, sigma_tower: Tensor) -> Tensor

Compose an elementwise activation onto a jet: b = sigma(u).

u_jet is the pre-activation jet of shape (N+1, ...); sigma_tower[k] = sigma^(k)(u_jet[0]) (same trailing shape). Returns the jet of sigma(u(t)).

The shifted power v^k, v = u - u_0, has valuation k, so its coefficients below order k vanish identically and are never formed: the recurrence runs over the (k, n) triangle n >= k only. That is a constant-factor saving (~N^3/6 instead of ~N^3/2 elementwise multiply-adds) returning bit-for-bit the same values as the dense form, **not** an asymptotic one -- see omnibias.jax.jet for why no O(N^2) claim is made for an arbitrary tower, and compose_jet_riccati for the quadratic fastpath that a Riccati activation does admit.

Two documented consequences of no longer forming products that were exactly zero: a coefficient that *is* exactly zero may come out with the other sign of zero (-0.0 == 0.0), and an infinite or NaN tower entry no longer contaminates lower orders through inf * 0.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
sigma_towerpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When sigma_tower.shape[0] != np1, raises ValueError(f'sigma_tower order {sigma_tower.shape[0] - 1} must match jet order {np1 - 1}').

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

Executed regression · surrogate jet parity

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_surrogate_jet_parity() -> None:
    tj = _np(tq.surrogate_jet(torch.as_tensor(Z), BETA, order=4))
    jj = _np(jq.surrogate_jet(jnp.asarray(Z), BETA, order=4))
    assert np.allclose(tj, jj, rtol=1e-12, atol=1e-13)

python -m pytest "packages/omnibias-binary/tests/test_surrogate.py::test_surrogate_jet_parity" -q

Regression example · riccati order zero and constant polynomial

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_riccati_order_zero_and_constant_polynomial() -> None:
    u_jet = torch.tensor([[0.3], [0.5], [0.25], [0.1]])
    # P constant: sigma' = 3 everywhere, so sigma(u) = sigma(u0) + 3 (u - u0).
    got = compose_jet_riccati(u_jet, torch.tensor([2.0]), (3.0,))
    assert torch.allclose(got, torch.tensor([[2.0], [1.5], [0.75], [0.3]]))
    scalar = compose_jet_riccati(u_jet[:1], torch.tensor([2.0]), (0.0, 1.0, -1.0))
    assert scalar.shape == (1, 1)
    assert torch.equal(scalar[0], torch.tensor([2.0]))

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_riccati_order_zero_and_constant_polynomial -q

Read the implementation

def compose_jet(u_jet: Tensor, sigma_tower: Tensor) -> Tensor:
    """Compose an elementwise activation onto a jet: ``b = sigma(u)``.

    ``u_jet`` is the pre-activation jet of shape ``(N+1, ...)``;
    ``sigma_tower[k] = sigma^(k)(u_jet[0])`` (same trailing shape). Returns the
    jet of ``sigma(u(t))``.

    The shifted power ``v^k``, ``v = u - u_0``, has valuation ``k``, so its
    coefficients below order ``k`` vanish identically and are never formed:
    the recurrence runs over the ``(k, n)`` triangle ``n >= k`` only. That is a
    constant-factor saving (``~N^3/6`` instead of ``~N^3/2`` elementwise
    multiply-adds) returning bit-for-bit the same values as the dense form,
    **not** an asymptotic one -- see :mod:`omnibias.jax.jet` for why no
    ``O(N^2)`` claim is made for an arbitrary tower, and
    :func:`compose_jet_riccati` for the quadratic fastpath that a Riccati
    activation does admit.

    Two documented consequences of no longer forming products that were exactly
    zero: a coefficient that *is* exactly zero may come out with the other sign
    of zero (``-0.0 == 0.0``), and an infinite or NaN tower entry no longer
    contaminates lower orders through ``inf * 0``.
    """
    u_jet = torch.as_tensor(u_jet)
    sigma_tower = torch.as_tensor(sigma_tower)
    np1 = u_jet.shape[0]
    if sigma_tower.shape[0] != np1:
        raise ValueError(
            f"sigma_tower order {sigma_tower.shape[0] - 1} must match jet order "
            f"{np1 - 1}"
        )
    zero = torch.zeros_like(u_jet[0])
    result = [sigma_tower[0] * torch.ones_like(u_jet[0])]
    result += [zero for _ in range(np1 - 1)]
    # p[j] is the order-(k + j) coefficient of v^k; v^1 == u_jet[1:] exactly.
    p = [u_jet[n] for n in range(1, np1)]
    fact = 1.0
    for k in range(1, np1):
        fact *= k
        dk = sigma_tower[k] / fact
        for n in range(k, np1):
            result[n] = result[n] + dk * p[n - k]
        if k + 1 < np1:
            # p <- truncated_conv(p, v), keeping orders k+1 .. N of v^(k+1).
            new_p = []
            for n in range(k + 1, np1):
                acc = zero
                for i in range(k, n):
                    acc = acc + p[i - k] * u_jet[n - i]
                new_p.append(acc)
            p = new_p
    return torch.stack(result, dim=0)

Other public import paths

• omnibias.torch.compose_jet
• omnibias.torch.jet.compose_jet

FUNCTION

compose_jet_riccati

omnibias.torch.compose_jet_riccatiSource ↗

compose_jet_riccati(u_jet: Tensor, sigma_u0: Tensor, riccati_polynomial: Sequence[float]) -> Tensor

Riccati-class composition b = sigma(u) in O(deg(P) * N^2).

Specialised fastpath for an activation in the **Riccati class**, i.e. one whose derivative is a polynomial in the activation itself, sigma'(z) = P(sigma(z)) -- the riccati_polynomial recorded on omnibias.core.spec.ActivationSpec (sigmoid: s - s^2, tanh: 1 - t^2, exp: E, tan: 1 + t^2, ...). The chain rule then closes on b alone,

b'(t) = P(b(t))\, u'(t),

so equating Taylor coefficients gives the forward recurrence

(n+1)\, b_{n+1} = \sum_{i=0}^{n} P(b)_i \, (n-i+1)\, u_{n+1-i},

where P(b) needs the truncated powers b^m, m <= deg(P), each advanced one order per step. Every order costs O(deg(P) * n), hence O(deg(P) * N^2) overall, and the derivative tower is never built: the only activation evaluation is sigma_u0.

Parameters

u_jet Pre-activation jet, shape (N+1, ...). sigma_u0 sigma(u_jet[0]), broadcastable to the trailing shape of u_jet. riccati_polynomial Ascending coefficients of P, i.e. P(s) = sum_m c_m s^m.

Returns

Tensor Jet of sigma(u(t)), shape (N+1, ...).

Notes

Mathematically identical to compose_jet fed the exact tower, but it is a *different* sequence of floating-point operations, so it agrees to rounding (~1e-13 relative in float64) rather than bit-for-bit; against its omnibias.jax.jet twin it agrees to a few ULP, the same footing compose_jet sits on across backends. compose_jet remains the default everywhere so the pinned cross-backend goldens keep their meaning.

It also reaches *further* than the tower path: tan, cot and coth have order-capped fastpath kernels, so layer_jet raises above the cap, while this recurrence needs only sigma(u_0) and P.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	TensorRequired	Pre-activation jet, shape (N+1, ...).
sigma_u0positional-or-keyword	TensorRequired	sigma(u_jet[0]), broadcastable to the trailing shape of u_jet.
riccati_polynomialpositional-or-keyword	Sequence[float]Required	Ascending coefficients of P, i.e. P(s) = sum_m c_m s^m.

ReturnsTensor

Explicit validation & failure cases

• When not poly, raises ValueError('riccati_polynomial must hold at least one coefficient (P of sigma)').

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

Executed regression · riccati error paths

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_riccati_error_paths() -> None:
    u_jet = torch.tensor([[0.3], [0.5]])
    with pytest.raises(ValueError, match="at least one coefficient"):
        compose_jet_riccati(u_jet, torch.tensor([1.0]), ())
    rng = np.random.default_rng(43)
    W = torch.as_tensor(rng.normal(size=(2, 1)))
    with pytest.raises(ValueError, match="not in the Riccati class"):
        layer_jet(u_jet, W, None, "relu", 1, riccati=True)

python -m pytest "packages/omnibias-torch/tests/test_jet.py::test_riccati_error_paths" -q

Regression example · riccati order zero and constant polynomial

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_riccati_order_zero_and_constant_polynomial() -> None:
    u_jet = torch.tensor([[0.3], [0.5], [0.25], [0.1]])
    # P constant: sigma' = 3 everywhere, so sigma(u) = sigma(u0) + 3 (u - u0).
    got = compose_jet_riccati(u_jet, torch.tensor([2.0]), (3.0,))
    assert torch.allclose(got, torch.tensor([[2.0], [1.5], [0.75], [0.3]]))
    scalar = compose_jet_riccati(u_jet[:1], torch.tensor([2.0]), (0.0, 1.0, -1.0))
    assert scalar.shape == (1, 1)
    assert torch.equal(scalar[0], torch.tensor([2.0]))

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_riccati_order_zero_and_constant_polynomial -q

Read the implementation

def compose_jet_riccati(
    u_jet: Tensor,
    sigma_u0: Tensor,
    riccati_polynomial: Sequence[float],
) -> Tensor:
    r"""Riccati-class composition ``b = sigma(u)`` in ``O(deg(P) * N^2)``.

    Specialised fastpath for an activation in the **Riccati class**, i.e. one
    whose derivative is a polynomial in the activation itself,
    ``sigma'(z) = P(sigma(z))`` -- the ``riccati_polynomial`` recorded on
    :class:`omnibias.core.spec.ActivationSpec` (``sigmoid``: ``s - s^2``,
    ``tanh``: ``1 - t^2``, ``exp``: ``E``, ``tan``: ``1 + t^2``, ...). The
    chain rule then closes on ``b`` alone,

    .. math::

        b'(t) = P(b(t))\, u'(t),

    so equating Taylor coefficients gives the forward recurrence

    .. math::

        (n+1)\, b_{n+1} = \sum_{i=0}^{n} P(b)_i \, (n-i+1)\, u_{n+1-i},

    where ``P(b)`` needs the truncated powers ``b^m``, ``m <= deg(P)``, each
    advanced one order per step. Every order costs ``O(deg(P) * n)``, hence
    ``O(deg(P) * N^2)`` overall, and the derivative tower is never built: the
    only activation evaluation is ``sigma_u0``.

    Parameters
    ----------
    u_jet
        Pre-activation jet, shape ``(N+1, ...)``.
    sigma_u0
        ``sigma(u_jet[0])``, broadcastable to the trailing shape of ``u_jet``.
    riccati_polynomial
        Ascending coefficients of ``P``, i.e. ``P(s) = sum_m c_m s^m``.

    Returns
    -------
    Tensor
        Jet of ``sigma(u(t))``, shape ``(N+1, ...)``.

    Notes
    -----
    Mathematically identical to :func:`compose_jet` fed the exact tower, but it
    is a *different* sequence of floating-point operations, so it agrees to
    rounding (``~1e-13`` relative in float64) rather than bit-for-bit; against
    its :mod:`omnibias.jax.jet` twin it agrees to a few ULP, the same footing
    :func:`compose_jet` sits on across backends. :func:`compose_jet` remains
    the default everywhere so the pinned cross-backend goldens keep their
    meaning.

    It also reaches *further* than the tower path: ``tan``, ``cot`` and ``coth``
    have order-capped fastpath kernels, so :func:`layer_jet` raises above the
    cap, while this recurrence needs only ``sigma(u_0)`` and ``P``.
    """
    u_jet = torch.as_tensor(u_jet)
    sigma_u0 = torch.as_tensor(sigma_u0)
    poly = tuple(float(c) for c in riccati_polynomial)
    if not poly:
        raise ValueError(
            "riccati_polynomial must hold at least one coefficient (P of sigma)"
        )
    np1 = u_jet.shape[0]
    b = [sigma_u0 * torch.ones_like(u_jet[0])]
    if np1 == 1:
        return torch.stack(b, dim=0)
    deg = len(poly) - 1
    du = derivative_jet(u_jet)  # du[m] = (m+1) * u_jet[m+1]
    zero = torch.zeros_like(b[0])
    # powers[m] holds the coefficients of b^m computed so far; index 1 is b.
    powers: list[list[Tensor]] = [[], b] + [[] for _ in range(max(deg - 1, 0))]
    q: list[Tensor] = []
    for n in range(np1 - 1):
        for m in range(2, deg + 1):
            acc = zero
            for j in range(n + 1):
                acc = acc + powers[m - 1][j] * b[n - j]
            powers[m].append(acc)
        qn = (zero + poly[0]) if n == 0 else zero
        for m in range(1, deg + 1):
            if poly[m] != 0.0:
                qn = qn + poly[m] * powers[m][n]
        q.append(qn)
        acc = zero
        for i in range(n + 1):
            acc = acc + q[i] * du[n - i]
        b.append(acc / float(n + 1))
    return torch.stack(b, dim=0)

Related APIsomnibias.torch.derivative_jet↗

Other public import paths

• omnibias.torch.compose_jet_riccati
• omnibias.torch.jet.compose_jet_riccati

FUNCTION

derivative_jet

omnibias.torch.derivative_jetSource ↗

derivative_jet(jet: Tensor) -> Tensor

Taylor jet of f' from the jet of f (differentiation; leading axis = order).

The derivative's Taylor coefficients are (k+1) * jet[k+1], so the jet shortens by one order. Works on scalar (N+1,) and vector (N+1, ...) jets (e.g. an mlp_jet output). Exact float inverse of antiderivative_jet; the differentiation half of the Fundamental Theorem of Calculus in the jet register.

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · derivative jet scalar matches manual

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_derivative_jet_scalar_matches_manual() -> None:
    jet = torch.tensor([2.0, 3.0, 5.0, 7.0])  # a_0 .. a_3
    d = derivative_jet(jet)  # (k+1) a_{k+1}
    assert torch.allclose(d, torch.tensor([3.0, 10.0, 21.0]))

python -m pytest "packages/omnibias-torch/tests/test_jet.py::test_derivative_jet_scalar_matches_manual" -q

Regression example · antiderivative jet ftc part1 roundtrip

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_antiderivative_jet_ftc_part1_roundtrip() -> None:
    """FTC part 1: d/dt of the antiderivative jet recovers the original vector jet."""
    layers, x0, v = _build_mlp(seed=4)
    order = 6
    jet = mlp_jet(x0, v, layers, order)  # (order+1, C)
    anti = antiderivative_jet(jet, constant=0.3)
    assert anti.shape[0] == jet.shape[0] + 1
    assert torch.allclose(anti[0], torch.full_like(jet[0], 0.3))
    back = derivative_jet(anti)
    assert back.shape == jet.shape
    assert torch.allclose(back, jet, rtol=1e-12, atol=1e-12)

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_antiderivative_jet_ftc_part1_roundtrip -q

Read the implementation

def derivative_jet(jet: Tensor) -> Tensor:
    r"""Taylor jet of ``f'`` from the jet of ``f`` (differentiation; leading axis = order).

    The derivative's Taylor coefficients are ``(k+1) * jet[k+1]``, so the jet
    shortens by one order. Works on scalar ``(N+1,)`` and vector ``(N+1, ...)``
    jets (e.g. an :func:`mlp_jet` output). Exact float inverse of
    :func:`antiderivative_jet`; the differentiation half of the Fundamental
    Theorem of Calculus in the jet register.
    """
    jet = torch.as_tensor(jet)
    np1 = jet.shape[0]
    factors = torch.arange(1, np1, dtype=jet.dtype, device=jet.device)
    return _scale_axis0(jet[1:], factors)

Other public import paths

• omnibias.torch.derivative_jet
• omnibias.torch.jet.derivative_jet

FUNCTION

jet_to_tower

omnibias.torch.jet_to_towerSource ↗

jet_to_tower(jet: Tensor) -> Tensor

Convert a Taylor jet to the derivative tower (f, f', f'', ...).

tower[k] = jet[k] * k!. Matches the scalar-curve Taylor-coefficient convention used by jax.experimental.jet in the validation suite.

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · deep mlp matches nested jacfwd

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_deep_mlp_matches_nested_jacfwd() -> None:
    layers, x0, v = _build_mlp(seed=2)
    f = _forward(layers)
    order = 6

    def g(t):
        return f(x0 + t * v)

    def kth(k):
        fn = g
        for _ in range(k):
            fn = jacfwd(fn)
        return fn(torch.tensor(0.0))

    tower = jet_to_tower(mlp_jet(x0, v, layers, order))
    for k in range(order + 1):
        assert torch.allclose(tower[k], kth(k), rtol=1e-9, atol=1e-9)

python -m pytest "packages/omnibias-torch/tests/test_jet.py::test_deep_mlp_matches_nested_jacfwd" -q

Regression example · compose jet does not propagate inf to lower orders

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_compose_jet_does_not_propagate_inf_to_lower_orders() -> None:
    """The one documented semantic change: no ``inf * 0`` contamination."""
    rng = np.random.default_rng(29)
    order = 4
    u_jet = torch.as_tensor(rng.normal(size=(order + 1, 1)))
    tower = torch.as_tensor(rng.normal(size=(order + 1, 1)))
    tower[order] = math.inf
    dense = _dense_compose_jet(u_jet, tower)
    lean = compose_jet(u_jet, tower)
    assert torch.isnan(dense[:order]).all()  # inf * 0 -> NaN in the dense form
    assert torch.isfinite(lean[:order]).all()
    assert torch.isinf(lean[order]).all()

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_compose_jet_does_not_propagate_inf_to_lower_orders -q

Read the implementation

def jet_to_tower(jet: Tensor) -> Tensor:
    """Convert a Taylor jet to the derivative tower ``(f, f', f'', ...)``.

    ``tower[k] = jet[k] * k!``. Matches the scalar-curve Taylor-coefficient
    convention used by :mod:`jax.experimental.jet` in the validation suite.
    """
    jet = torch.as_tensor(jet)
    facts = _factorials(jet.shape[0], jet.dtype, jet.device)
    return _scale_axis0(jet, facts)

Other public import paths

• omnibias.torch.jet_to_tower
• omnibias.torch.jet.jet_to_tower

FUNCTION

layer_jet

omnibias.torch.layer_jetSource ↗

layer_jet(z_jet: Tensor, W: Tensor, b: Tensor | None, spec: ActivationSpec[Tensor] | str, order: int | None=None, *, riccati: bool=False) -> Tensor

Push a jet through one sigma(W z + b) layer.

The activation derivative tower is built from spec.fastpath; passing an activation whose fastpath does not reach the jet order raises ValueError.

riccati=True opts into compose_jet_riccati, which skips the tower entirely and costs O(N^2) instead of O(N^3); it requires spec.riccati_polynomial and changes only the rounding, not the mathematical result. The default stays False so the pinned goldens keep their bit-for-bit meaning.

Parameters

Name	Type & default	Contract
z_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | NoneRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor] | strRequired	See the operation contract and implementation below.
orderpositional-or-keyword	int | NoneDefault: None	See the operation contract and implementation below.
riccatikeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order is not None and order != jet_order, raises ValueError(f'order {order} must equal the jet order {jet_order} carried by z_jet').
• When poly is None, raises ValueError(f'activation {resolved.name!r} is not in the Riccati class (no riccati_polynomial); drop riccati=True to use the general tower kernel').

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

Executed regression · arctan order cap raises for k2 estimator

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_arctan_order_cap_raises_for_k2_estimator() -> None:
    layers = _make_layers(3, 4, depth=1, activation="arctan", seed=25)
    x = torch.as_tensor(np.random.default_rng(26).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_arctan_order_cap_raises_for_k2_estimator" -q

Regression example · riccati error paths

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_riccati_error_paths() -> None:
    u_jet = torch.tensor([[0.3], [0.5]])
    with pytest.raises(ValueError, match="at least one coefficient"):
        compose_jet_riccati(u_jet, torch.tensor([1.0]), ())
    rng = np.random.default_rng(43)
    W = torch.as_tensor(rng.normal(size=(2, 1)))
    with pytest.raises(ValueError, match="not in the Riccati class"):
        layer_jet(u_jet, W, None, "relu", 1, riccati=True)

python -m pytest packages/omnibias-torch/tests/test_jet.py -k test_riccati_error_paths -q

Read the implementation

def layer_jet(
    z_jet: Tensor,
    W: Tensor,
    b: Tensor | None,
    spec: ActivationSpec[Tensor] | str,
    order: int | None = None,
    *,
    riccati: bool = False,
) -> Tensor:
    """Push a jet through one ``sigma(W z + b)`` layer.

    The activation derivative tower is built from ``spec.fastpath``; passing an
    activation whose fastpath does not reach the jet order raises ``ValueError``.

    ``riccati=True`` opts into :func:`compose_jet_riccati`, which skips the
    tower entirely and costs ``O(N^2)`` instead of ``O(N^3)``; it requires
    ``spec.riccati_polynomial`` and changes only the rounding, not the
    mathematical result. The default stays ``False`` so the pinned goldens keep
    their bit-for-bit meaning.
    """
    z_jet = torch.as_tensor(z_jet)
    jet_order = z_jet.shape[0] - 1
    if order is not None and order != jet_order:
        raise ValueError(
            f"order {order} must equal the jet order {jet_order} carried by z_jet"
        )
    resolved = get_activation(spec)
    u_jet = affine_jet(z_jet, W, b)
    if riccati:
        poly = resolved.riccati_polynomial
        if poly is None:
            raise ValueError(
                f"activation {resolved.name!r} is not in the Riccati class "
                "(no riccati_polynomial); drop riccati=True to use the general "
                "tower kernel"
            )
        return compose_jet_riccati(u_jet, resolved.forward(u_jet[0]), poly)
    sigma_tower = _sigma_tower(resolved, u_jet[0], jet_order)
    return compose_jet(u_jet, sigma_tower)

Related APIsomnibias.torch.get_activation↗omnibias.torch.affine_jet↗omnibias.torch.compose_jet↗omnibias.torch.compose_jet_riccati↗

Other public import paths

• omnibias.torch.layer_jet
• omnibias.torch.jet.layer_jet

FUNCTION

lhopital_ratio

omnibias.torch.lhopital_ratioSource ↗

lhopital_ratio(num_jet: Tensor, den_jet: Tensor, order: int=1) -> Tensor

Differentiable L'Hopital limit of f(t)/g(t) as t -> 0.

For a 0/0 form whose numerator and denominator both vanish to order order - 1 -- so f(t) = a_p t^p + O(t^{p+1}) and g(t) = b_p t^p + O(t^{p+1}) with p = order -- the limit is the ratio of the leading Taylor coefficients a_p / b_p, i.e. num_jet[order] / den_jet[order] in the Taylor convention jet[k] = f^(k)(0)/k!.

This is the *differentiable* entry point: a plain elementwise division, so the limit backpropagates. Compose it with mlp_jet to take the limit of a *learned* field along a ray. Use limit_of_ratio when the vanishing order is not known ahead of time.

Parameters

Name	Type & default	Contract
num_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
den_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intDefault: 1	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When order > max_order, raises ValueError(f'order {order} exceeds available jet order {max_order}').

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

Executed regression · lhopital sinc

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_lhopital_sinc() -> None:
    assert float(lhopital_ratio(_SIN, _X, order=1)) == pytest.approx(1.0, abs=1e-15)

python -m pytest "packages/omnibias-torch/tests/test_torch_limit.py::test_lhopital_sinc" -q

Regression example · lhopital one minus cos over x2

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_lhopital_one_minus_cos_over_x2() -> None:
    assert float(lhopital_ratio(_ONE_MINUS_COS, _X2, order=2)) == pytest.approx(0.5, abs=1e-15)

python -m pytest packages/omnibias-torch/tests/test_torch_limit.py -k test_lhopital_one_minus_cos_over_x2 -q

Read the implementation

def lhopital_ratio(num_jet: Tensor, den_jet: Tensor, order: int = 1) -> Tensor:
    r"""Differentiable L'Hopital limit of ``f(t)/g(t)`` as ``t -> 0``.

    For a ``0/0`` form whose numerator and denominator both vanish to order
    ``order - 1`` -- so ``f(t) = a_p t^p + O(t^{p+1})`` and
    ``g(t) = b_p t^p + O(t^{p+1})`` with ``p = order`` -- the limit is the ratio
    of the leading Taylor coefficients ``a_p / b_p``, i.e.
    ``num_jet[order] / den_jet[order]`` in the Taylor convention
    ``jet[k] = f^(k)(0)/k!``.

    This is the *differentiable* entry point: a plain elementwise division, so
    the limit backpropagates. Compose it with :func:`mlp_jet` to take the limit
    of a *learned* field along a ray. Use :func:`limit_of_ratio` when the
    vanishing order is not known ahead of time.
    """
    num_jet = torch.as_tensor(num_jet)
    den_jet = torch.as_tensor(den_jet)
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    max_order = min(num_jet.shape[0], den_jet.shape[0]) - 1
    if order > max_order:
        raise ValueError(f"order {order} exceeds available jet order {max_order}")
    return num_jet[order] / den_jet[order]

Other public import paths

• omnibias.torch.lhopital_ratio
• omnibias.torch.jet.lhopital_ratio

FUNCTION

limit_of_ratio

omnibias.torch.limit_of_ratioSource ↗

limit_of_ratio(num_jet: Tensor, den_jet: Tensor, *, atol: float=1e-12) -> Tensor

Auto-detect the leading order of lim_{t->0} f(t)/g(t) (forward-only).

Scans both scalar jets (shape (N+1,)) for their lowest non-vanishing Taylor order and returns

• num_jet[q] / den_jet[q] when numerator and denominator share leading order q (a genuine 0/0 resolved by L'Hopital),
• 0 when the numerator vanishes to strictly higher order,
• +/- inf when the denominator vanishes faster (a pole).

Uses Python control flow on the detected order, so it is a convenience that is **not** differentiable at order transitions; the differentiable counterpart is lhopital_ratio.

Parameters

Name	Type & default	Contract
num_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
den_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
atolkeyword-only	floatDefault: 1e-12	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When num_jet.ndim != 1 or den_jet.ndim != 1, raises ValueError('limit_of_ratio expects scalar (1-D) jets').
• When den_ord is None, raises ValueError('denominator jet vanishes to all computed orders; the limit is undefined at this truncation order').

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

Executed regression · limit of ratio pole is infinite

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_limit_of_ratio_pole_is_infinite() -> None:
    assert math.isinf(float(limit_of_ratio(_X, _X2)))

python -m pytest "packages/omnibias-torch/tests/test_torch_limit.py::test_limit_of_ratio_pole_is_infinite" -q

Regression example · limit of ratio numerator vanishes faster is zero

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_limit_of_ratio_numerator_vanishes_faster_is_zero() -> None:
    assert float(limit_of_ratio(_X2, _X)) == 0.0

python -m pytest packages/omnibias-torch/tests/test_torch_limit.py -k test_limit_of_ratio_numerator_vanishes_faster_is_zero -q

Read the implementation

def limit_of_ratio(num_jet: Tensor, den_jet: Tensor, *, atol: float = 1e-12) -> Tensor:
    r"""Auto-detect the leading order of ``lim_{t->0} f(t)/g(t)`` (forward-only).

    Scans both scalar jets (shape ``(N+1,)``) for their lowest non-vanishing
    Taylor order and returns

    * ``num_jet[q] / den_jet[q]`` when numerator and denominator share leading
      order ``q`` (a genuine ``0/0`` resolved by L'Hopital),
    * ``0`` when the numerator vanishes to strictly higher order,
    * ``+/- inf`` when the denominator vanishes faster (a pole).

    Uses Python control flow on the detected order, so it is a convenience that
    is **not** differentiable at order transitions; the differentiable
    counterpart is :func:`lhopital_ratio`.
    """
    num_jet = torch.as_tensor(num_jet)
    den_jet = torch.as_tensor(den_jet)
    if num_jet.ndim != 1 or den_jet.ndim != 1:
        raise ValueError("limit_of_ratio expects scalar (1-D) jets")
    den_ord = _leading_order_scalar(den_jet, atol)
    if den_ord is None:
        raise ValueError(
            "denominator jet vanishes to all computed orders; the limit is "
            "undefined at this truncation order"
        )
    num_ord = _leading_order_scalar(num_jet, atol)
    dtype = torch.result_type(num_jet, den_jet)
    if num_ord is None or num_ord > den_ord:
        return torch.zeros((), dtype=dtype)
    if num_ord < den_ord:
        sign = float(torch.sign(num_jet[num_ord]) * torch.sign(den_jet[den_ord]))
        return torch.tensor(sign * math.inf, dtype=dtype)
    return (num_jet[den_ord] / den_jet[den_ord]).to(dtype)

Other public import paths

• omnibias.torch.limit_of_ratio
• omnibias.torch.jet.limit_of_ratio

FUNCTION

mlp_jet

omnibias.torch.mlp_jetSource ↗

mlp_jet(x0: Tensor, v: Tensor, layers: Sequence[tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | str | None]], order: int, *, riccati: bool=False) -> Tensor

Exact directional Taylor jet of a deep MLP along x(t) = x0 + t v.

layers is a sequence of (W, b, spec); spec=None is a pure affine readout. Returns a jet of shape (N+1, C) whose jet_to_tower gives the directional derivatives d^k/dt^k f(x0 + t v) at t = 0.

riccati=True forwards to the compose_jet_riccati fastpath in every activation layer and requires all of them to be Riccati-class.

Use a directional jet when you need several derivatives along one chosen input direction. Row k stores the k-th Taylor coefficient, so multiply by k! or call jet_to_tower before interpreting it as a raw derivative. Parameter gradients remain available through the supported tensor operations. The supplied examples assert their shapes and finite weight gradients; floating-point accuracy still depends on order, activation, and input scale.

Parameters

Name	Type & default	Contract
x0positional-or-keyword	TensorRequired	Base input with final feature axis D. The tested batch example accepts (B, D) and returns a jet shaped (N+1, B, C); a single point uses (D,).
vpositional-or-keyword	TensorRequired	Direction with the same feature shape as x0. This defines the path x(t) = x0 + t v; scaling v scales an order-k directional derivative by that scale to the k-th power.
layerspositional-or-keyword	Sequence[tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | str | None]]Required	Ordered (W, b, activation) triples. W maps input features to output features, b is an optional output-width bias, and activation=None makes a purely affine layer. Consecutive feature dimensions must match.
orderpositional-or-keyword	intRequired	Nonnegative maximum Taylor order N. The returned leading axis has N+1 entries, including the value at order zero.
riccatikeyword-only	boolDefault: False	When true, use the specialized Riccati composition path. Every nonlinear layer must provide the required Riccati polynomial.

ReturnsTensor

Complete runnable example · high order torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""A fourth-order residual, with gradients into model parameters."""
import torch
from omnibias.torch.jet import jet_to_tower, mlp_jet

weight = torch.tensor([[0.7], [-0.4]], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([0.1, 0.2], dtype=torch.float64, requires_grad=True)
readout = torch.tensor([[1.0, -0.5]], dtype=torch.float64, requires_grad=True)
layers = [(weight, bias, "tanh"), (readout, None, None)]
x = torch.linspace(-1.0, 1.0, 32, dtype=torch.float64).reshape(-1, 1)

tower = jet_to_tower(mlp_jet(x, torch.ones_like(x), layers, order=4))
assert tower.shape == (5, 32, 1)
u, du, d2u, d3u, d4u = tower.unbind(0)
loss = (d4u - 24.0).square().mean()
loss.backward()
for parameter in (weight, bias, readout):
    assert parameter.grad is not None
    assert torch.isfinite(parameter.grad).all()
print("tower:", tuple(tower.shape), "finite parameter gradients: True")
# This checks a residual calculation, not a solved boundary-value problem.

python high_order_torch.py

Executed regression · arctan order cap raises for k2 estimator

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_arctan_order_cap_raises_for_k2_estimator() -> None:
    layers = _make_layers(3, 4, depth=1, activation="arctan", seed=25)
    x = torch.as_tensor(np.random.default_rng(26).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_arctan_order_cap_raises_for_k2_estimator" -q

Read the implementation

def mlp_jet(
    x0: Tensor,
    v: Tensor,
    layers: Sequence[tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | str | None]],
    order: int,
    *,
    riccati: bool = False,
) -> Tensor:
    """Exact directional Taylor jet of a deep MLP along ``x(t) = x0 + t v``.

    ``layers`` is a sequence of ``(W, b, spec)``; ``spec=None`` is a pure affine
    readout. Returns a jet of shape ``(N+1, C)`` whose :func:`jet_to_tower` gives
    the directional derivatives ``d^k/dt^k f(x0 + t v)`` at ``t = 0``.

    ``riccati=True`` forwards to the :func:`compose_jet_riccati` fastpath in
    every activation layer and requires all of them to be Riccati-class.
    """
    jet = _path_jet(x0, v, order)
    for W, b, spec in layers:
        if spec is None:
            jet = affine_jet(jet, W, b)
        else:
            jet = layer_jet(jet, W, b, spec, order, riccati=riccati)
    return jet

Related APIsomnibias.torch.affine_jet↗omnibias.torch.layer_jet↗

Other public import paths

• omnibias.torch.mlp_jet
• omnibias.torch.jet.mlp_jet

FUNCTION

mlp_joint_jet

omnibias.torch.jet.mlp_joint_jetSource ↗

mlp_joint_jet(input_jet: Tensor, parameter_jet: Tensor, spec: RealizationSpec, *, order: int, max_coefficients: int=1000000) -> Tensor

Propagate Taylor-normalized input and parameter jets in one direction.

Delegates to the shared realization adapter, including its coefficient budget and six operator roles. Inputs, weights and biases remain live.

Parameters

Name	Type & default	Contract
input_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
parameter_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	RealizationSpecRequired	See the operation contract and implementation below.
orderkeyword-only	intRequired	See the operation contract and implementation below.
max_coefficientskeyword-only	intDefault: 1000000	See the operation contract and implementation below.

ReturnsTensor

Regression example · preseeded joint mlp wrappers preserve values live gradients and dtype

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_preseeded_joint_mlp_wrappers_preserve_values_live_gradients_and_dtype(dtype, dim):
    spec = RealizationSpec.dense((2, 2, 1))
    order = 2
    count = num_multi_indices(dim, order)
    x = np.linspace(-0.2, 0.3, count * 2, dtype=dtype).reshape(count, 2)
    p = np.linspace(-0.3, 0.4, count * spec.layout.size, dtype=dtype).reshape(count, -1)
    tx = torch.tensor(x, requires_grad=True)
    tp = torch.tensor(p, requires_grad=True)
    options = {"order": order, "max_coefficients": 1000}
    if dim == 1:
        twrapper, jwrapper = tj.mlp_joint_jet, jj.mlp_joint_jet
    else:
        twrapper, jwrapper = tmv.mlp_joint_jet_mv, jmv.mlp_joint_jet_mv
        options["dim"] = dim
    actual = twrapper(tx, tp, spec, **options)
    expected = tr.realization_jet(tx, tp, spec, dim=dim, order=order, max_coefficients=1000)
    assert actual.dtype == tx.dtype
    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
    gradients = torch.autograd.grad(actual.square().sum(), (tx, tp), retain_graph=True)
    reference = torch.autograd.grad(expected.square().sum(), (tx, tp))
    for got, want in zip(gradients, reference, strict=True):
        assert torch.isfinite(got).all() and torch.count_nonzero(got)
        torch.testing.assert_close(got, want, rtol=0, atol=0)

    def wrapped(a, b):
        return jwrapper(a, b, spec, **options)

    def direct(a, b):
        return jr.realization_jet(a, b, spec, dim=dim, order=order, max_coefficients=1000)

    jx, jp = jnp.asarray(x), jnp.asarray(p)
    result = jax.jit(wrapped)(jx, jp)
    assert result.dtype == jx.dtype
    np.testing.assert_array_equal(result, jax.jit(direct)(jx, jp))
    jgrads = jax.jit(jax.grad(lambda a, b: jnp.sum(wrapped(a, b) ** 2), argnums=(0, 1)))(jx, jp)
    references = jax.jit(jax.grad(lambda a, b: jnp.sum(direct(a, b) ** 2), argnums=(0, 1)))(jx, jp)
    for got, want in zip(jgrads, references, strict=True):
        assert np.isfinite(got).all() and np.count_nonzero(got)
        np.testing.assert_array_equal(got, want)
    tolerance = 2e-6 if dtype is np.float32 else 2e-12
    np.testing.assert_allclose(result, actual.detach().numpy(), rtol=tolerance, atol=tolerance)
    for wrapper, a, b in ((twrapper, tx, tp), (jwrapper, jx, jp)):
        with pytest.raises(ValueError, match="budget"):
            wrapper(a, b, spec, **{**options, "max_coefficients": 1})

python -m pytest tests/test_joint_jet_entrypoints.py -k test_preseeded_joint_mlp_wrappers_preserve_values_live_gradients_and_dtype -q

Read the implementation

def mlp_joint_jet(
    input_jet: Tensor,
    parameter_jet: Tensor,
    spec: RealizationSpec,
    *,
    order: int,
    max_coefficients: int = 1_000_000,
) -> Tensor:
    """Propagate Taylor-normalized input and parameter jets in one direction.

    Delegates to the shared realization adapter, including its coefficient
    budget and six operator roles. Inputs, weights and biases remain live.
    """
    from omnibias.torch.realization import realization_jet

    return realization_jet(
        input_jet, parameter_jet, spec, dim=1, order=order,
        max_coefficients=max_coefficients,
    )

FUNCTION

removable_value

omnibias.torch.removable_valueSource ↗

removable_value(jet: Tensor) -> Tensor

Limit t -> 0 of a path jet -- its zeroth Taylor coefficient.

For a jet built along x(t) = x0 + t v (e.g. via mlp_jet) the value at the base point is jet[0]; when the jet was assembled to cancel a removable singularity this is the finite limit. Differentiable (no host synchronisation).

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · removable value is zeroth coefficient

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_removable_value_is_zeroth_coefficient() -> None:
    jet = _t([3.5, 1.0, -2.0])
    assert float(removable_value(jet)) == 3.5

python -m pytest "packages/omnibias-torch/tests/test_torch_limit.py::test_removable_value_is_zeroth_coefficient" -q

Regression example · limit of a learned field along a ray

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_limit_of_a_learned_field_along_a_ray() -> None:
    torch.manual_seed(0)
    W1 = torch.randn(4, 2, dtype=_F64)
    b1 = torch.zeros(4, dtype=_F64)
    W2 = torch.randn(1, 4, dtype=_F64)
    b2 = torch.zeros(1, dtype=_F64)
    x0 = _t([0.3, -0.7])
    v = _t([1.0, 0.5])
    layers = [(W1, b1, "tanh"), (W2, b2, None)]
    jet = mlp_jet(x0, v, layers, order=2)
    direct = W2 @ torch.tanh(W1 @ x0 + b1) + b2
    assert float(removable_value(jet)[0]) == pytest.approx(float(direct[0]), abs=1e-12)

python -m pytest packages/omnibias-torch/tests/test_torch_limit.py -k test_limit_of_a_learned_field_along_a_ray -q

Read the implementation

def removable_value(jet: Tensor) -> Tensor:
    r"""Limit ``t -> 0`` of a path jet -- its zeroth Taylor coefficient.

    For a jet built along ``x(t) = x0 + t v`` (e.g. via :func:`mlp_jet`) the
    value at the base point is ``jet[0]``; when the jet was assembled to cancel a
    removable singularity this is the finite limit. Differentiable (no host
    synchronisation).
    """
    return torch.as_tensor(jet)[0]

Other public import paths

• omnibias.torch.removable_value
• omnibias.torch.jet.removable_value

FUNCTION

tower_to_jet

omnibias.torch.tower_to_jetSource ↗

tower_to_jet(tower: Tensor) -> Tensor

Convert a derivative tower (f, f', f'', ...) to a Taylor jet.

jet[k] = tower[k] / k!. Leading axis is the derivative order.

Parameters

Name	Type & default	Contract
towerpositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def tower_to_jet(tower: Tensor) -> Tensor:
    """Convert a derivative tower ``(f, f', f'', ...)`` to a Taylor jet.

    ``jet[k] = tower[k] / k!``. Leading axis is the derivative order.
    """
    tower = torch.as_tensor(tower)
    facts = _factorials(tower.shape[0], tower.dtype, tower.device)
    return _scale_axis0(tower, 1.0 / facts)

Other public import paths

• omnibias.torch.tower_to_jet
• omnibias.torch.jet.tower_to_jet

← Package overviewRunnable learning examples →