API REFERENCE/omnibias-jax

omnibias.jax.jet

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.jet

Module contract & mathematical conventions

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

Algebraically corresponding twin of omnibias.torch.jet.

A *jet* is the truncated Taylor expansion of a scalar-parametrised path. We store the Taylor coefficients a_k = f^(k)(0) / k! along the leading axis of an array of shape (N+1, ...) and broadcast over the trailing axes (hidden units, output components, batch).

For a deep network evaluated along a line x(t) = x0 + t v the pre-activation of every hidden unit is a power series in t. Propagating its jet through the elementwise activation is exactly Faà di Bruno's formula; because omnibias supplies the *exact* derivative tower sigma^(k) (Eulerian / Legendre / Hermite recurrences in omnibias.core.polynomials), the composition is exact and needs neither nested autodiff nor finite differences.

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,

truncated at order N; series powers are built by truncated convolution. The combinatorially-explicit Bell-polynomial form (omnibias.core.bell.faa_di_bruno_terms) is the cross-check used in the test-suite, not the production path.

Truncation-order cost

compose_jet takes an **arbitrary** derivative tower, so what it has to evaluate is the truncated composition of the polynomial D(x) = sum_k sigma^(k)(u_0) x^k / k! with the valuation-1 series v = u - u_0. The kernel knows nothing about D beyond its coefficients, and every output coefficient b_n really does depend on all sigma^(k)(u_0) and u_j with k, j <= n, which already forces Omega(N^2) work.

• Because v^k has valuation k, all coefficients of v^k below order k vanish identically. Restricting the recurrence to the triangle n >= k removes those structurally-zero products: ~N^3/6 elementwise multiply-adds instead of the dense ~N^3/2, returning the same values bit-for-bit in eager execution (see compose_jet). The exponent is unchanged -- this is a constant factor, measured at 3.3x fewer multiplies at N = 16 in benchmarks/jet_compose_cost.py.
• Sub-cubic composition is not ruled out in theory: baby-step / giant-step (Brent-Kung) reaches O(N^{5/2}) when truncated multiplication is schoolbook, and near-linear composition exists in the ring-operation model. Both buy the exponent by transforming or reordering the coefficient axis (FFT / Karatsuba style multiplication), which is *not* exact in floating point -- the Taylor coefficients here span many magnitudes, so a transform would trade the kernel's exactness (and the pinned cross-backend goldens) for the exponent. At the truncation orders this API is used at (N of order ten) N^{5/2} with its larger constant is not even a win. So **no O(N^2) claim is made for** compose_jet; it is cubic in N, with a smaller constant than the dense form.
• O(N^2) *is* available when the outer map is not arbitrary. An activation in the Riccati class satisfies sigma' = P(sigma) for a polynomial P (recorded as omnibias.core.spec.ActivationSpec.riccati_polynomial), which closes the chain rule on b itself and gives the O(deg(P) * N^2) recurrence in compose_jet_riccati -- also skipping tower construction. It is opt-in via riccati=True on layer_jet / mlp_jet because it rounds differently from compose_jet.

important:

**Bit-parity requires 64-bit JAX.** The "algebraically corresponding twin" guarantee against omnibias.torch.jet holds in double precision. PyTorch defaults to float64 for Python floats, whereas JAX silently truncates arrays to float32 unless 64-bit mode is enabled. Enable it *before* the first JAX import::

import jax jax.config.update("jax_enable_x64", True)

(equivalently set JAX_ENABLE_X64=1 / --jax_enable_x64 in the env). Without it the jets are still internally consistent, but cross-backend comparisons will only match to float32 tolerance, not bit-for-bit.

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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.jax.affine_jetSource ↗

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

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,). The map is 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	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · value equals jet row0

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_value_equals_jet_row0() -> None:
    net = make_jet_mlp(2, 16, 1, depth=3, seed=1)
    x = _x(2)
    jet = net.jet(x, 2)
    assert jnp.allclose(jet[:, 0, :], net.value(x), atol=1e-14)

python -m pytest "packages/omnibias-jax/tests/test_jet_pinn.py::test_value_equals_jet_row0" -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 = jnp.asarray(rng.normal(size=(order + 1, 5)))
    W = jnp.asarray(rng.normal(size=(3, 5)))
    b = jnp.asarray(rng.normal(size=(3,)))
    out = affine_jet(z_jet, W, b)
    assert jnp.allclose(out[0], W @ z_jet[0] + b)
    for k in range(1, order + 1):
        assert jnp.allclose(out[k], W @ z_jet[k])

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

Read the implementation

def affine_jet(z_jet: Array, W: Array, b: Array | None = None) -> Array:
    """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,)``. The map is linear in ``t`` so it
    acts per order, with the bias added to the zeroth coefficient only.
    """
    z_jet = jnp.asarray(z_jet)
    W = jnp.asarray(W)
    out = jnp.tensordot(z_jet, W, axes=([-1], [-1]))  # (N+1, ..., D_out)
    if b is not None:
        out = out.at[0].add(jnp.asarray(b))
    return out

Other public import paths

• omnibias.jax.affine_jet
• omnibias.jax.jet.affine_jet

FUNCTION

affine_joint_jet

omnibias.jax.jet.affine_joint_jetSource ↗

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

Live directional product with jet-valued weights and biases.

Parameters

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

ReturnsArray

Read the implementation

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

    return impl(input_jet, weight_jet, bias_jet)

FUNCTION

antiderivative_jet

omnibias.jax.antiderivative_jetSource ↗

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

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, JIT/vmap safe; 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	ArrayRequired	See the operation contract and implementation below.
constantpositional-or-keyword	floatDefault: 0.0	See the operation contract and implementation below.

ReturnsArray

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:
    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 jnp.allclose(anti[0], jnp.full_like(jet[0], 0.3))
    back = derivative_jet(anti)
    assert back.shape == jet.shape
    assert jnp.allclose(back, jet, rtol=1e-12, atol=1e-12)

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

Regression example · antiderivative derivative jit vmap safe

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_derivative_jit_vmap_safe() -> None:
    rng = np.random.default_rng(9)
    jet = jnp.asarray(rng.normal(size=(6, 3)))
    a1 = jax.jit(antiderivative_jet)(jet)
    a2 = antiderivative_jet(jet)
    assert jnp.array_equal(a1, a2)
    d1 = jax.jit(derivative_jet)(jet)
    assert jnp.array_equal(d1, derivative_jet(jet))
    batch = jnp.stack([jet, jet * 2.0])
    av = jax.vmap(antiderivative_jet)(batch)
    assert av.shape == (2, 7, 3)

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

Read the implementation

def antiderivative_jet(jet: Array, constant: float = 0.0) -> Array:
    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, JIT/vmap safe; 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 = jnp.asarray(jet)
    np1 = jet.shape[0]
    divisors = jnp.arange(1, np1 + 1, dtype=jet.dtype)
    scaled = _scale_axis0(jet, 1.0 / divisors)
    const_row = jnp.full_like(jet[0], constant)[None, ...]
    return jnp.concatenate([const_row, scaled], axis=0)

Other public import paths

• omnibias.jax.antiderivative_jet
• omnibias.jax.jet.antiderivative_jet

FUNCTION

compose_jet

omnibias.jax.compose_jetSource ↗

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

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

Parameters

u_jet Jet of the pre-activation, shape (N+1, ...) (Taylor coefficients). sigma_tower Derivative tower of the activation evaluated at u_jet[0]: sigma_tower[k] = sigma^(k)(u_jet[0]), same trailing shape as u_jet.

Returns

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

Notes

w(t) = u(t) - u_0 has valuation 1, so w^k cannot contribute below order k; the recurrence therefore runs over the triangle n >= k and never forms the structurally-zero coefficients. Values are unchanged -- bit-for-bit against the dense form in eager execution, for finite inputs -- so the win is a constant factor of about three, not a better exponent. Two documented consequences of dropping 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. See the module docstring for the cost discussion and compose_jet_riccati for the quadratic fastpath.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	ArrayRequired	Jet of the pre-activation, shape (N+1, ...) (Taylor coefficients).
sigma_towerpositional-or-keyword	ArrayRequired	Derivative tower of the activation evaluated at u_jet[0]: sigma_tower[k] = sigma^(k)(u_jet[0]), same trailing shape as u_jet.

ReturnsArray

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 · 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 = jnp.asarray(rng.normal(size=(order + 1, 1)))
    tower = jnp.asarray(rng.normal(size=(order + 1, 1))).at[order].set(jnp.inf)
    dense = _dense_compose_jet(u_jet, tower)
    lean = compose_jet(u_jet, tower)
    assert bool(jnp.all(jnp.isnan(dense[:order])))
    assert bool(jnp.all(jnp.isfinite(lean[:order])))
    assert bool(jnp.all(jnp.isinf(lean[order])))

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

Read the implementation

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

    Parameters
    ----------
    u_jet
        Jet of the pre-activation, shape ``(N+1, ...)`` (Taylor coefficients).
    sigma_tower
        Derivative tower of the activation evaluated at ``u_jet[0]``:
        ``sigma_tower[k] = sigma^(k)(u_jet[0])``, same trailing shape as
        ``u_jet``.

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

    Notes
    -----
    ``w(t) = u(t) - u_0`` has valuation ``1``, so ``w^k`` cannot contribute
    below order ``k``; the recurrence therefore runs over the triangle
    ``n >= k`` and never forms the structurally-zero coefficients. Values are
    unchanged -- bit-for-bit against the dense form in eager execution, for
    finite inputs -- so the win is a constant factor of about three, not a
    better exponent. Two documented consequences of dropping 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``. See the module
    docstring for the cost discussion and :func:`compose_jet_riccati` for the
    quadratic fastpath.
    """
    u_jet = jnp.asarray(u_jet)
    sigma_tower = jnp.asarray(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 = jnp.zeros_like(u_jet[0])
    result = [sigma_tower[0] * jnp.ones_like(u_jet[0])]
    result += [zero for _ in range(np1 - 1)]
    # p[j] is the order-(k + j) coefficient of w^k; w^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, w), keeping orders k+1 .. N of w^(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 jnp.stack(result, axis=0)

Other public import paths

• omnibias.jax.compose_jet
• omnibias.jax.jet.compose_jet

FUNCTION

compose_jet_riccati

omnibias.jax.compose_jet_riccatiSource ↗

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

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

Array 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.torch.jet twin it agrees to a few ULP, the same footing compose_jet sits on across backends. JIT/vmap safe. 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	ArrayRequired	Pre-activation jet, shape (N+1, ...).
sigma_u0positional-or-keyword	ArrayRequired	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.

ReturnsArray

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 = jnp.asarray([[0.3], [0.5]])
    with pytest.raises(ValueError, match="at least one coefficient"):
        compose_jet_riccati(u_jet, jnp.asarray([1.0]), ())
    rng = np.random.default_rng(43)
    W = jnp.asarray(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-jax/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 = jnp.asarray([[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, jnp.asarray([2.0]), (3.0,))
    assert jnp.allclose(got, jnp.asarray([[2.0], [1.5], [0.75], [0.3]]))
    scalar = compose_jet_riccati(u_jet[:1], jnp.asarray([2.0]), (0.0, 1.0, -1.0))
    assert scalar.shape == (1, 1)
    assert jnp.array_equal(scalar[0], jnp.asarray([2.0]))

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

Read the implementation

def compose_jet_riccati(
    u_jet: Array,
    sigma_u0: Array,
    riccati_polynomial: Sequence[float],
) -> Array:
    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
    -------
    Array
        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.torch.jet` twin it agrees to a few ULP, the same footing
    :func:`compose_jet` sits on across backends. JIT/vmap safe.
    :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 = jnp.asarray(u_jet)
    sigma_u0 = jnp.asarray(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 * jnp.ones_like(u_jet[0])]
    if np1 == 1:
        return jnp.stack(b, axis=0)
    deg = len(poly) - 1
    du = derivative_jet(u_jet)  # du[m] = (m+1) * u_jet[m+1]
    zero = jnp.zeros_like(b[0])
    # powers[m] holds the coefficients of b^m computed so far; index 1 is b.
    powers: list[list[Array]] = [[], b] + [[] for _ in range(max(deg - 1, 0))]
    q: list[Array] = []
    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 jnp.stack(b, axis=0)

Related APIsomnibias.jax.derivative_jet↗

Other public import paths

• omnibias.jax.compose_jet_riccati
• omnibias.jax.jet.compose_jet_riccati

FUNCTION

derivative_jet

omnibias.jax.derivative_jetSource ↗

derivative_jet(jet: Array) -> Array

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) and is JIT/vmap safe. 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	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

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:
    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 jnp.allclose(anti[0], jnp.full_like(jet[0], 0.3))
    back = derivative_jet(anti)
    assert back.shape == jet.shape
    assert jnp.allclose(back, jet, rtol=1e-12, atol=1e-12)

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

Regression example · antiderivative derivative jit vmap safe

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_derivative_jit_vmap_safe() -> None:
    rng = np.random.default_rng(9)
    jet = jnp.asarray(rng.normal(size=(6, 3)))
    a1 = jax.jit(antiderivative_jet)(jet)
    a2 = antiderivative_jet(jet)
    assert jnp.array_equal(a1, a2)
    d1 = jax.jit(derivative_jet)(jet)
    assert jnp.array_equal(d1, derivative_jet(jet))
    batch = jnp.stack([jet, jet * 2.0])
    av = jax.vmap(antiderivative_jet)(batch)
    assert av.shape == (2, 7, 3)

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

Read the implementation

def derivative_jet(jet: Array) -> Array:
    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) and is JIT/vmap safe. Exact float
    inverse of :func:`antiderivative_jet`; the differentiation half of the
    Fundamental Theorem of Calculus in the jet register.
    """
    jet = jnp.asarray(jet)
    np1 = jet.shape[0]
    factors = jnp.arange(1, np1, dtype=jet.dtype)
    return _scale_axis0(jet[1:], factors)

Other public import paths

• omnibias.jax.derivative_jet
• omnibias.jax.jet.derivative_jet

FUNCTION

jet_to_tower

omnibias.jax.jet_to_towerSource ↗

jet_to_tower(jet: Array) -> Array

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

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

Parameters

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

ReturnsArray

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 = jax.jacfwd(fn)
        return fn(0.0)

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

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

Regression example · deep mlp matches experimental jet

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_experimental_jet() -> None:
    layers, x0, v = _build_mlp(seed=1)
    f = _forward(layers)
    order = 6
    series = [v] + [jnp.zeros_like(v) for _ in range(order - 1)]
    y0, terms = jax_jet.jet(f, (x0,), (series,))
    expected = [y0, *terms]  # derivative tower d^k/dt^k f
    tower = jet_to_tower(mlp_jet(x0, v, layers, order))
    for k in range(order + 1):
        assert jnp.allclose(tower[k], expected[k], rtol=1e-10, atol=1e-10)

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

Read the implementation

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

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

Other public import paths

• omnibias.jax.jet_to_tower
• omnibias.jax.jet.jet_to_tower

FUNCTION

layer_jet

omnibias.jax.layer_jetSource ↗

layer_jet(z_jet: Array, W: Array, b: Array | None, spec: JaxActivationSpec | str, order: int | None=None, *, riccati: bool=False) -> Array

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	ArrayRequired	See the operation contract and implementation below.
Wpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	Array | NoneRequired	See the operation contract and implementation below.
specpositional-or-keyword	JaxActivationSpec | 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.

ReturnsArray

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 = jnp.asarray(np.random.default_rng(26).normal(size=(3,)))
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")

python -m pytest "packages/omnibias-jax/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 = jnp.asarray([[0.3], [0.5]])
    with pytest.raises(ValueError, match="at least one coefficient"):
        compose_jet_riccati(u_jet, jnp.asarray([1.0]), ())
    rng = np.random.default_rng(43)
    W = jnp.asarray(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-jax/tests/test_jet.py -k test_riccati_error_paths -q

Read the implementation

def layer_jet(
    z_jet: Array,
    W: Array,
    b: Array | None,
    spec: JaxActivationSpec | str,
    order: int | None = None,
    *,
    riccati: bool = False,
) -> Array:
    """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 = jnp.asarray(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.jax.get_activation↗omnibias.jax.affine_jet↗omnibias.jax.compose_jet↗omnibias.jax.compose_jet_riccati↗

Other public import paths

• omnibias.jax.layer_jet
• omnibias.jax.jet.layer_jet

FUNCTION

lhopital_ratio

omnibias.jax.lhopital_ratioSource ↗

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

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 and is vmap/jit compatible. 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	ArrayRequired	See the operation contract and implementation below.
den_jetpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intDefault: 1	See the operation contract and implementation below.

ReturnsArray

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-jax/tests/test_jax_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-jax/tests/test_jax_limit.py -k test_lhopital_one_minus_cos_over_x2 -q

Read the implementation

def lhopital_ratio(num_jet: Array, den_jet: Array, order: int = 1) -> Array:
    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 and is ``vmap``/``jit`` compatible. 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 = jnp.asarray(num_jet)
    den_jet = jnp.asarray(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.jax.lhopital_ratio
• omnibias.jax.jet.lhopital_ratio

FUNCTION

limit_of_ratio

omnibias.jax.limit_of_ratioSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
den_jetpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
atolkeyword-only	floatDefault: 1e-12	See the operation contract and implementation below.

ReturnsArray

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:
    # x / x^2 -> +inf
    assert math.isinf(float(limit_of_ratio(_X, _X2)))

python -m pytest "packages/omnibias-jax/tests/test_jax_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:
    # x^2 / x  -> 0
    assert float(limit_of_ratio(_X2, _X)) == 0.0

python -m pytest packages/omnibias-jax/tests/test_jax_limit.py -k test_limit_of_ratio_numerator_vanishes_faster_is_zero -q

Read the implementation

def limit_of_ratio(num_jet: Array, den_jet: Array, *, atol: float = 1e-12) -> Array:
    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 = jnp.asarray(num_jet)
    den_jet = jnp.asarray(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 = jnp.result_type(num_jet.dtype, den_jet.dtype)
    if num_ord is None or num_ord > den_ord:
        return jnp.asarray(0.0, dtype=dtype)
    if num_ord < den_ord:
        sign = float(jnp.sign(num_jet[num_ord]) * jnp.sign(den_jet[den_ord]))
        return jnp.asarray(sign * jnp.inf, dtype=dtype)
    return (num_jet[den_ord] / den_jet[den_ord]).astype(dtype)

Other public import paths

• omnibias.jax.limit_of_ratio
• omnibias.jax.jet.limit_of_ratio

FUNCTION

mlp_jet

omnibias.jax.mlp_jetSource ↗

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

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

Parameters

x0, v Base point and direction, shape (D,) each. layers Sequence of (W, b, spec). spec=None is a pure affine readout (no activation); otherwise the layer computes sigma(W z + b). order Truncation order N; the returned jet has N+1 coefficients. riccati Opt into the compose_jet_riccati fastpath in every activation layer; requires all of them to be Riccati-class.

Returns

Array 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.

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	ArrayRequired	Base point with shape (D,). For multiple points, use the backend batch API or its documented vectorization transform.
vpositional-or-keyword	ArrayRequired	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[Array, Array | None, JaxActivationSpec | 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.

ReturnsArray

Complete runnable example · high order jax

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

Download complete script ↓

"""JAX directional jets and ordinary parameter autodiff."""
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from omnibias.jax.jet import jet_to_tower, mlp_jet

def residual_loss(weight):
    layers = [(weight, jnp.array([0.1, 0.2]), "tanh"),
              (jnp.array([[0.6, -0.4]]), None, None)]
    tower = jet_to_tower(mlp_jet(jnp.array([0.3]), jnp.ones(1), layers, order=4))
    assert tower.shape == (5, 1)
    return jnp.mean((tower[4] - 24.0) ** 2)

weight = jnp.array([[0.5], [-0.3]])
loss, gradient = jax.value_and_grad(residual_loss)(weight)
assert gradient.shape == weight.shape
assert bool(jnp.all(jnp.isfinite(gradient)))
print("gradient:", gradient.shape, "finite:", bool(jnp.isfinite(loss)))

python high_order_jax.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 = jnp.asarray(np.random.default_rng(26).normal(size=(3,)))
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")

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

Read the implementation

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

    Parameters
    ----------
    x0, v
        Base point and direction, shape ``(D,)`` each.
    layers
        Sequence of ``(W, b, spec)``. ``spec=None`` is a pure affine readout
        (no activation); otherwise the layer computes ``sigma(W z + b)``.
    order
        Truncation order ``N``; the returned jet has ``N+1`` coefficients.
    riccati
        Opt into the :func:`compose_jet_riccati` fastpath in every activation
        layer; requires all of them to be Riccati-class.

    Returns
    -------
    Array
        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``.
    """
    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.jax.affine_jet↗omnibias.jax.layer_jet↗

Other public import paths

• omnibias.jax.mlp_jet
• omnibias.jax.jet.mlp_jet

FUNCTION

mlp_joint_jet

omnibias.jax.jet.mlp_joint_jetSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
parameter_jetpositional-or-keyword	ArrayRequired	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.

ReturnsArray

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: Array,
    parameter_jet: Array,
    spec: RealizationSpec,
    *,
    order: int,
    max_coefficients: int = 1_000_000,
) -> Array:
    """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.jax.realization import realization_jet

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

FUNCTION

removable_value

omnibias.jax.removable_valueSource ↗

removable_value(jet: Array) -> Array

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 and jit-friendly (no host synchronisation).

Parameters

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

ReturnsArray

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 = jnp.array([3.5, 1.0, -2.0])
    assert float(removable_value(jet)) == 3.5

python -m pytest "packages/omnibias-jax/tests/test_jax_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:
    # The limit t->0 of a deep tanh MLP along x0 + t v is the network value at x0.
    key = jax.random.PRNGKey(0)
    k1, k2 = jax.random.split(key)
    W1 = jax.random.normal(k1, (4, 2), dtype=jnp.float64)
    b1 = jnp.zeros(4)
    W2 = jax.random.normal(k2, (1, 4), dtype=jnp.float64)
    b2 = jnp.zeros(1)
    x0 = jnp.array([0.3, -0.7])
    v = jnp.array([1.0, 0.5])
    layers = [(W1, b1, "tanh"), (W2, b2, None)]
    jet = mlp_jet(x0, v, layers, order=2)
    direct = (W2 @ jnp.tanh(W1 @ x0 + b1) + b2)
    assert float(removable_value(jet)[0]) == pytest.approx(float(direct[0]), abs=1e-12)

python -m pytest packages/omnibias-jax/tests/test_jax_limit.py -k test_limit_of_a_learned_field_along_a_ray -q

Read the implementation

def removable_value(jet: Array) -> Array:
    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 and
    ``jit``-friendly (no host synchronisation).
    """
    return jnp.asarray(jet)[0]

Other public import paths

• omnibias.jax.removable_value
• omnibias.jax.jet.removable_value

FUNCTION

tower_to_jet

omnibias.jax.tower_to_jetSource ↗

tower_to_jet(tower: Array) -> Array

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	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

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

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

Other public import paths

• omnibias.jax.tower_to_jet
• omnibias.jax.jet.tower_to_jet

← Package overviewRunnable learning examples →