API REFERENCE/omnibias-jax

omnibias.jax.jet_mv

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_mv

Module contract & mathematical conventions

Exact multivariate (multi-index) jets via Faà di Bruno (jax).

Algebraically corresponding twin of omnibias.torch.jet_mv. This is the full D-variable generalisation of the directional jets in omnibias.jax.jet: instead of a single Taylor series in the path parameter t, we carry the *whole* truncated multivariate Taylor expansion of a deep network around a base point x0,

f(x_0 + \delta) = \sum_{|\alpha| \le N} c_\alpha\, \delta^\alpha, \qquad c_\alpha = \frac{D^\alpha f(x_0)}{\alpha!},

so a single forward pass yields *every* mixed partial derivative up to total order N (the gradient, the full Hessian, third-order tensors, ...).

Representation

A multivariate jet is a dense array of shape (M, ...) where M = num_multi_indices(dim, order) and row i holds the Taylor coefficient c_alpha for alpha = multi_indices(dim, order)[i]. The trailing axes are hidden units / output components, exactly as for the directional jet.

Kernel

Affine maps act per coefficient (the map is degree-1 in delta); the bias hits the constant coefficient only. Activations use the same shifted-power identity as the directional kernel,

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

but the series powers (u - u_0)^k are built with *multivariate* truncated polynomial multiplication (omnibias.core.multi_index.multiply_table) rather than 1-D convolution. The exact derivative tower sigma^(k)(u_0) comes from the omnibias activation fast paths, so the composition is exact and needs neither nested autodiff nor finite differences.

important:

**Bit-parity requires 64-bit JAX.** As with omnibias.jax.jet, the "algebraically corresponding twin" guarantee against omnibias.torch.jet_mv holds in double precision only. JAX silently truncates arrays to float32 unless 64-bit mode is enabled *before* the first JAX import::

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

(equivalently JAX_ENABLE_X64=1). Without it the jet is still internally consistent, but cross-backend comparisons match only 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_mv

omnibias.jax.affine_jet_mvSource ↗

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

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

Identical per-coefficient action to omnibias.jax.affine_jet: the map is degree-1 in the perturbation, so it acts row-wise with the bias added to the constant coefficient (row 0) 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 · affine jet mv linear per 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_affine_jet_mv_linear_per_coefficient() -> None:
    rng = np.random.default_rng(7)
    dim, order = 2, 3
    m = len(multi_indices(dim, order))
    z_jet = jnp.asarray(rng.normal(size=(m, 5)))
    W = jnp.asarray(rng.normal(size=(3, 5)))
    b = jnp.asarray(rng.normal(size=(3,)))
    out = affine_jet_mv(z_jet, W, b)
    assert jnp.allclose(out[0], W @ z_jet[0] + b)
    for k in range(1, m):
        assert jnp.allclose(out[k], W @ z_jet[k])

python -m pytest "packages/omnibias-jax/tests/test_jet_mv.py::test_affine_jet_mv_linear_per_coefficient" -q

Read the implementation

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

    Identical per-coefficient action to :func:`omnibias.jax.affine_jet`: the map
    is degree-1 in the perturbation, so it acts row-wise with the bias added to
    the constant coefficient (row 0) only.
    """
    return affine_jet(z_jet, W, b)

Related APIsomnibias.jax.affine_jet↗

Other public import paths

• omnibias.jax.affine_jet_mv
• omnibias.jax.jet_mv.affine_jet_mv

FUNCTION

affine_joint_jet_mv

omnibias.jax.jet_mv.affine_joint_jet_mvSource ↗

affine_joint_jet_mv(input_jet: Array, weight_jet: Array, bias_jet: Array | None=None, *, dim: int, order: int) -> Array

Live multivariate 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.
dimkeyword-only	intRequired	See the operation contract and implementation below.
orderkeyword-only	intRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def affine_joint_jet_mv(
    input_jet: Array, weight_jet: Array, bias_jet: Array | None = None,
    *, dim: int, order: int,
) -> Array:
    """Live multivariate product with jet-valued weights and biases."""
    from omnibias.jax.realization import affine_joint_jet_mv as impl

    return impl(input_jet, weight_jet, bias_jet, dim=dim, order=order)

FUNCTION

compose_jet_mv

omnibias.jax.compose_jet_mvSource ↗

compose_jet_mv(u_jet: Array, sigma_tower: Array, dim: int, order: int) -> Array

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

Parameters

u_jet Pre-activation jet, shape (M, ...) with M = num_multi_indices(dim, order) (rows in the canonical ~omnibias.core.multi_index.multi_indices order). sigma_tower Derivative tower of the activation at the constant coefficient u_jet[0]: sigma_tower[k] = sigma^(k)(u_jet[0]) for k = 0..order (shape (order + 1, ...)). dim, order Number of input variables and the truncation total order.

Returns

Array Jet of sigma(u) of shape (M, ...).

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	ArrayRequired	Pre-activation jet, shape (M, ...) with M = num_multi_indices(dim, order) (rows in the canonical ~omnibias.core.multi_index.multi_indices order).
sigma_towerpositional-or-keyword	ArrayRequired	Derivative tower of the activation at the constant coefficient u_jet[0]: sigma_tower[k] = sigma^(k)(u_jet[0]) for k = 0..order (shape (order + 1, ...)).
dimpositional-or-keyword	intRequired	Number of input variables and the truncation total order.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When u_jet.shape[0] != m, raises ValueError(f'u_jet has {u_jet.shape[0]} rows but dim={dim}, order={order} requires {m}').
• When sigma_tower.shape[0] != order + 1, raises ValueError(f'sigma_tower order {sigma_tower.shape[0] - 1} must equal jet order {order}').

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

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

Read the implementation

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

    Parameters
    ----------
    u_jet
        Pre-activation jet, shape ``(M, ...)`` with
        ``M = num_multi_indices(dim, order)`` (rows in the canonical
        :func:`~omnibias.core.multi_index.multi_indices` order).
    sigma_tower
        Derivative tower of the activation at the constant coefficient
        ``u_jet[0]``: ``sigma_tower[k] = sigma^(k)(u_jet[0])`` for
        ``k = 0..order`` (shape ``(order + 1, ...)``).
    dim, order
        Number of input variables and the truncation total order.

    Returns
    -------
    Array
        Jet of ``sigma(u)`` of shape ``(M, ...)``.
    """
    u_jet = jnp.asarray(u_jet)
    sigma_tower = jnp.asarray(sigma_tower)
    m = num_multi_indices(dim, order)
    if u_jet.shape[0] != m:
        raise ValueError(
            f"u_jet has {u_jet.shape[0]} rows but dim={dim}, order={order} "
            f"requires {m}"
        )
    if sigma_tower.shape[0] != order + 1:
        raise ValueError(
            f"sigma_tower order {sigma_tower.shape[0] - 1} must equal jet order "
            f"{order}"
        )

    table = multiply_table(dim, order)
    zero = jnp.zeros_like(u_jet[0])
    # w = u - u0 has no constant term (row 0 is zero).
    w = [zero] + [u_jet[g] for g in range(1, m)]
    # p holds the coefficients of w^k; start at w^0 = 1 (only the constant row).
    p = [jnp.ones_like(u_jet[0])] + [zero for _ in range(m - 1)]
    result = [sigma_tower[0] * p[0]] + [zero for _ in range(m - 1)]
    fact = 1.0
    for k in range(1, order + 1):
        fact *= k
        # p <- truncated multivariate product (p * w) via the Cauchy table.
        acc = [zero for _ in range(m)]
        for g, a, b in table:
            acc[g] = acc[g] + p[a] * w[b]
        p = acc
        dk = sigma_tower[k] / fact
        for g in range(m):
            result[g] = result[g] + dk * p[g]
    return jnp.stack(result, axis=0)

Related APIsomnibias.core.multiply_table↗omnibias.core.num_multi_indices↗

Other public import paths

• omnibias.jax.compose_jet_mv
• omnibias.jax.jet_mv.compose_jet_mv

FUNCTION

identity_jet

omnibias.jax.identity_jetSource ↗

identity_jet(x0: Array, order: int) -> Array

Multivariate jet of the identity map x(delta) = x0 + delta.

Returns shape (M, D) where D = x0.shape[-1]: row 0 (the constant coefficient) is x0 and each unit-multi-index row e_i is the standard basis vector e_i (since d x_j / d delta_i = delta_{ij}). All other rows are zero. This is the seed for mlp_jet_mv.

Parameters

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

ReturnsArray

Explicit validation & failure cases

• When x0.ndim != 1, raises ValueError(f'x0 must be a 1-D vector, got shape {x0.shape}').
• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

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

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 · softmax rejects a scalar 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_softmax_rejects_a_scalar_jet(memory) -> None:
    x0 = memory["x0"]
    with pytest.raises(ValueError, match="trailing softmax axis"):
        jet_softmax(identity_jet(x0, ORDER)[:, 0], DIM, ORDER)

python -m pytest packages/omnibias-jax/tests/test_jet_attention.py -k test_softmax_rejects_a_scalar_jet -q

Read the implementation

def identity_jet(x0: Array, order: int) -> Array:
    """Multivariate jet of the identity map ``x(delta) = x0 + delta``.

    Returns shape ``(M, D)`` where ``D = x0.shape[-1]``: row 0 (the constant
    coefficient) is ``x0`` and each unit-multi-index row ``e_i`` is the standard
    basis vector ``e_i`` (since ``d x_j / d delta_i = delta_{ij}``). All other
    rows are zero. This is the seed for :func:`mlp_jet_mv`.
    """
    x0 = jnp.asarray(x0)
    if x0.ndim != 1:
        raise ValueError(f"x0 must be a 1-D vector, got shape {x0.shape}")
    dim = int(x0.shape[0])
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    pos = index_position(dim, order)
    m = num_multi_indices(dim, order)
    rows = [jnp.zeros((dim,), dtype=x0.dtype) for _ in range(m)]
    rows[pos[(0,) * dim]] = x0
    if order >= 1:
        eye = jnp.eye(dim, dtype=x0.dtype)
        for i in range(dim):
            e = tuple(1 if j == i else 0 for j in range(dim))
            rows[pos[e]] = eye[i]
    return jnp.stack(rows, axis=0)

Related APIsomnibias.core.index_position↗omnibias.core.num_multi_indices↗

Other public import paths

• omnibias.jax.identity_jet
• omnibias.jax.jet_mv.identity_jet

FUNCTION

jet_attention

omnibias.jax.jet_attentionSource ↗

jet_attention(q_jet: Array, keys: Array, values: Array, dim: int, order: int, beta: float | Array=1.0) -> Array

Jet of dot-product attention softmax(beta q K^T) V w.r.t. the *inputs*.

mathematically equivalent twin of omnibias.torch.jet_mv.jet_attention. omnibias.hopfield already differentiates the log-sum-exp core in closed form, but with respect to the **scores**. That is the wrong variable for a PDE: a residual needs d/dx. Pushing the query jet through the block supplies the missing coordinate story -- the value agrees with omnibias.hopfield.jax.ops.attention and every mixed partial D^alpha of the attention output comes out of the same single jet.

Parameters

q_jet Jet of the query, shape (M, ..., d_key). keys, values Memory of shape (n, d_key) and (n, d_val); constant in the input (they are network parameters, not functions of x), so they enter as two affine maps around the softmax. dim, order Input-variable count and truncation total order. beta Inverse temperature. An array is allowed, which keeps a *trainable* temperature closed form in exactly the same way.

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
q_jetpositional-or-keyword	ArrayRequired	Jet of the query, shape (M, ..., d_key).
keyspositional-or-keyword	ArrayRequired	Memory of shape (n, d_key) and (n, d_val); constant in the input (they are network parameters, not functions of x), so they enter as two affine maps around the softmax.
valuespositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	Input-variable count and truncation total order.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
betapositional-or-keyword	float | ArrayDefault: 1.0	Inverse temperature. An array is allowed, which keeps a *trainable* temperature closed form in exactly the same way.

ReturnsArray

Explicit validation & failure cases

• When keys.ndim != 2 or values.ndim != 2, raises ValueError(f'keys and values must be 2-D, got {keys.shape} and {values.shape}').
• When keys.shape[0] != values.shape[0], raises ValueError(f'keys and values must share the memory axis, got {keys.shape[0]} keys and {values.shape[0]} values').
• When q_jet.shape[-1] != keys.shape[-1], raises ValueError(f'query width {q_jet.shape[-1]} != key width {keys.shape[-1]}').

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

Executed regression · attention is jittable

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_attention_is_jittable(memory) -> None:
    W, b, K, V, x0 = (memory[k] for k in ("W", "b", "K", "V", "x0"))

    def block(x):
        return jet_attention(
            mlp_jet_mv(x, [(W, b, "tanh")], ORDER), K, V, DIM, ORDER, beta=1.2
        )

    assert _close(jax.jit(block)(x0), block(x0), 1e-14)

python -m pytest "packages/omnibias-jax/tests/test_jet_attention.py::test_attention_is_jittable" -q

Regression example · attention temperature is differentiable

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_attention_temperature_is_differentiable(memory) -> None:
    W, b, K, V, x0 = (memory[k] for k in ("W", "b", "K", "V", "x0"))
    q_jet = mlp_jet_mv(x0, [(W, b, "tanh")], ORDER)

    def total(beta):
        return jnp.sum(jet_attention(q_jet, K, V, DIM, ORDER, beta=beta))

    assert abs(float(jax.grad(total)(1.3))) > 0.0

python -m pytest packages/omnibias-jax/tests/test_jet_attention.py -k test_attention_temperature_is_differentiable -q

Read the implementation

def jet_attention(
    q_jet: Array,
    keys: Array,
    values: Array,
    dim: int,
    order: int,
    beta: float | Array = 1.0,
) -> Array:
    r"""Jet of dot-product attention ``softmax(beta q K^T) V`` w.r.t. the *inputs*.

    Bit-identical twin of :func:`omnibias.torch.jet_mv.jet_attention`.
    :mod:`omnibias.hopfield` already differentiates the log-sum-exp core in closed
    form, but with respect to the **scores**. That is the wrong variable for a PDE:
    a residual needs ``d/dx``. Pushing the query jet through the block supplies the
    missing coordinate story -- the value agrees with
    :func:`omnibias.hopfield.jax.ops.attention` and every mixed partial ``D^alpha``
    of the attention output comes out of the same single jet.

    Parameters
    ----------
    q_jet
        Jet of the query, shape ``(M, ..., d_key)``.
    keys, values
        Memory of shape ``(n, d_key)`` and ``(n, d_val)``; constant in the input
        (they are network parameters, not functions of ``x``), so they enter as two
        affine maps around the softmax.
    dim, order
        Input-variable count and truncation total order.
    beta
        Inverse temperature. An array is allowed, which keeps a *trainable*
        temperature closed form in exactly the same way.
    """
    q_jet = jnp.asarray(q_jet)
    keys = jnp.asarray(keys)
    values = jnp.asarray(values)
    if keys.ndim != 2 or values.ndim != 2:
        raise ValueError(
            f"keys and values must be 2-D, got {keys.shape} and {values.shape}"
        )
    if keys.shape[0] != values.shape[0]:
        raise ValueError(
            f"keys and values must share the memory axis, got {keys.shape[0]} "
            f"keys and {values.shape[0]} values"
        )
    if q_jet.shape[-1] != keys.shape[-1]:
        raise ValueError(
            f"query width {q_jet.shape[-1]} != key width {keys.shape[-1]}"
        )
    scores = affine_jet(q_jet, keys * beta)
    weights = jet_softmax(scores, dim, order)
    return affine_jet(weights, jnp.swapaxes(values, -2, -1))

Related APIsomnibias.jax.affine_jet↗omnibias.jax.jet_softmax↗

Other public import paths

• omnibias.jax.jet_attention
• omnibias.jax.jet_mv.jet_attention

FUNCTION

jet_exp

omnibias.jax.jet_expSource ↗

jet_exp(u_jet: Array, dim: int, order: int) -> Array

Jet of exp(u) from the jet of u.

A thin wrapper over compose_jet_mv with the tower of the registered "exp" activation (exp^(k) = exp at every order), so the exponential is not special-cased -- it is one more entry of the shared dictionary.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · attention is jittable

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_attention_is_jittable(memory) -> None:
    W, b, K, V, x0 = (memory[k] for k in ("W", "b", "K", "V", "x0"))

    def block(x):
        return jet_attention(
            mlp_jet_mv(x, [(W, b, "tanh")], ORDER), K, V, DIM, ORDER, beta=1.2
        )

    assert _close(jax.jit(block)(x0), block(x0), 1e-14)

python -m pytest "packages/omnibias-jax/tests/test_jet_attention.py::test_attention_is_jittable" -q

Regression example · exp 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_exp_parity(shared, order: int) -> None:
    t, j = _jets(shared, order)
    assert _agree(torch_exp(t, DIM, order), jax_exp(j, DIM, order))

python -m pytest tests/test_jet_attention_parity.py -k test_exp_parity -q

Read the implementation

def jet_exp(u_jet: Array, dim: int, order: int) -> Array:
    """Jet of ``exp(u)`` from the jet of ``u``.

    A thin wrapper over :func:`compose_jet_mv` with the tower of the registered
    ``"exp"`` activation (``exp^(k) = exp`` at every order), so the exponential is
    not special-cased -- it is one more entry of the shared dictionary.
    """
    u_jet = jnp.asarray(u_jet)
    tower = _sigma_tower(get_activation("exp"), u_jet[0], order)
    return compose_jet_mv(u_jet, tower, dim, order)

Related APIsomnibias.jax.get_activation↗omnibias.jax.compose_jet_mv↗

Other public import paths

• omnibias.jax.jet_exp
• omnibias.jax.jet_mv.jet_exp

FUNCTION

jet_gradient

omnibias.jax.jet_gradientSource ↗

jet_gradient(jet: Array, dim: int, order: int) -> Array

Gradient d f / d x_i from a jet, shape (D, ...).

Requires order >= 1.

Parameters

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

ReturnsArray

Explicit validation & failure cases

• When order < 1, raises ValueError(f'gradient needs order >= 1, got {order}').

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

Executed regression · make residual fn 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_make_residual_fn_roundtrip() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    r = residual_fn(flat0)
    assert r.shape[0] == 40  # 38 interior + 2 boundary
    assert jnp.all(jnp.isfinite(r))

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_make_residual_fn_roundtrip" -q

Regression example · gradient requires order one

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

python -m pytest packages/omnibias-jax/tests/test_jet_mv.py -k test_gradient_requires_order_one -q

Read the implementation

def jet_gradient(jet: Array, dim: int, order: int) -> Array:
    """Gradient ``d f / d x_i`` from a jet, shape ``(D, ...)``.

    Requires ``order >= 1``.
    """
    if order < 1:
        raise ValueError(f"gradient needs order >= 1, got {order}")
    jet = jnp.asarray(jet)
    pos = index_position(dim, order)
    rows = []
    for i in range(dim):
        e = tuple(1 if j == i else 0 for j in range(dim))
        rows.append(jet[pos[e]])  # alpha! = 1 for a unit multi-index
    return jnp.stack(rows, axis=0)

Related APIsomnibias.core.index_position↗

Other public import paths

• omnibias.jax.jet_gradient
• omnibias.jax.jet_mv.jet_gradient

FUNCTION

jet_hessian

omnibias.jax.jet_hessianSource ↗

jet_hessian(jet: Array, dim: int, order: int) -> Array

Hessian d^2 f / d x_i d x_j from a jet, shape (D, D, ...).

Requires order >= 2. The symmetric entry H[i, j] is read off the coefficient of alpha = e_i + e_j and rescaled by alpha! (2 on the diagonal, 1 off-diagonal).

Parameters

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

ReturnsArray

Explicit validation & failure cases

• When order < 2, raises ValueError(f'hessian needs order >= 2, got {order}').

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

Executed regression · make residual fn 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_make_residual_fn_roundtrip() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    r = residual_fn(flat0)
    assert r.shape[0] == 40  # 38 interior + 2 boundary
    assert jnp.all(jnp.isfinite(r))

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_make_residual_fn_roundtrip" -q

Regression example · hessian requires order two

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-jax/tests/test_jet_mv.py -k test_hessian_requires_order_two -q

Read the implementation

def jet_hessian(jet: Array, dim: int, order: int) -> Array:
    """Hessian ``d^2 f / d x_i d x_j`` from a jet, shape ``(D, D, ...)``.

    Requires ``order >= 2``. The symmetric entry ``H[i, j]`` is read off the
    coefficient of ``alpha = e_i + e_j`` and rescaled by ``alpha!`` (2 on the
    diagonal, 1 off-diagonal).
    """
    if order < 2:
        raise ValueError(f"hessian needs order >= 2, got {order}")
    jet = jnp.asarray(jet)
    pos = index_position(dim, order)
    rows = []
    for i in range(dim):
        row = []
        for j in range(dim):
            alpha = tuple(
                (1 if k == i else 0) + (1 if k == j else 0) for k in range(dim)
            )
            row.append(jet[pos[alpha]] * multi_index_factorial(alpha))
        rows.append(jnp.stack(row, axis=0))
    return jnp.stack(rows, axis=0)

Related APIsomnibias.core.index_position↗omnibias.core.multi_index_factorial↗

Other public import paths

• omnibias.jax.jet_hessian
• omnibias.jax.jet_mv.jet_hessian

FUNCTION

jet_multiply

omnibias.jax.jet_multiplySource ↗

jet_multiply(a: Array, b: Array, dim: int, order: int) -> Array

Truncated product of two multivariate jets: jet(a) * jet(b).

mathematically equivalent twin of omnibias.torch.jet_mv.jet_multiply. Given the jets of two analytic functions about the *same* base point, returns the jet of their pointwise product via the multivariate Cauchy product

(a\,b)_\gamma = \sum_{\alpha + \beta = \gamma} a_\alpha\, b_\beta,

truncated at total order order (the same ~omnibias.core.multi_index.multiply_table that drives compose_jet_mv). Both a and b carry M = num_multi_indices(dim, order) leading rows; any trailing shape broadcasts, so a scalar mask (M,) multiplies a vector field (M, C) component-wise. This is the jet-level Leibniz rule and is what keeps a product ansatz such as u = g + b * net closed-form to arbitrary order.

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When a.shape[0] != m, raises ValueError(f'a has {a.shape[0]} rows but dim={dim}, order={order} requires {m}').
• When b.shape[0] != m, raises ValueError(f'b has {b.shape[0]} rows but dim={dim}, order={order} requires {m}').

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

Executed regression · jet multiply rejects wrong rows

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_jet_multiply_rejects_wrong_rows() -> None:
    x0 = jnp.array([0.1, 0.2])
    idj = identity_jet(x0, 2)
    with pytest.raises(ValueError, match="rows"):
        jet_multiply(idj[:, 0], jnp.zeros(2), 2, 2)

python -m pytest "packages/omnibias-jax/tests/test_hardbc.py::test_jet_multiply_rejects_wrong_rows" -q

Regression example · jet multiply is polynomial product

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_jet_multiply_is_polynomial_product() -> None:
    x0 = jnp.array([0.3, -0.7])
    idj = identity_jet(x0, 2)
    prod = jet_multiply(idj[:, 0], idj[:, 1], 2, 2)
    assert jnp.allclose(jet_gradient(prod, 2, 1), jnp.array([-0.7, 0.3]))
    assert jnp.allclose(jet_hessian(prod, 2, 2), jnp.array([[0.0, 1.0], [1.0, 0.0]]))

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_jet_multiply_is_polynomial_product -q

Read the implementation

def jet_multiply(a: Array, b: Array, dim: int, order: int) -> Array:
    r"""Truncated product of two multivariate jets: ``jet(a) * jet(b)``.

    Bit-identical twin of :func:`omnibias.torch.jet_mv.jet_multiply`. Given the
    jets of two analytic functions about the *same* base point, returns the jet of
    their pointwise product via the multivariate Cauchy product

    .. math::

        (a\,b)_\gamma = \sum_{\alpha + \beta = \gamma} a_\alpha\, b_\beta,

    truncated at total order ``order`` (the same
    :func:`~omnibias.core.multi_index.multiply_table` that drives
    :func:`compose_jet_mv`). Both ``a`` and ``b`` carry
    ``M = num_multi_indices(dim, order)`` leading rows; any trailing shape
    broadcasts, so a scalar mask ``(M,)`` multiplies a vector field ``(M, C)``
    component-wise. This is the jet-level Leibniz rule and is what keeps a
    product ansatz such as ``u = g + b * net`` closed-form to arbitrary order.
    """
    a = jnp.asarray(a)
    b = jnp.asarray(b)
    m = num_multi_indices(dim, order)
    if a.shape[0] != m:
        raise ValueError(
            f"a has {a.shape[0]} rows but dim={dim}, order={order} requires {m}"
        )
    if b.shape[0] != m:
        raise ValueError(
            f"b has {b.shape[0]} rows but dim={dim}, order={order} requires {m}"
        )
    table = multiply_table(dim, order)
    sample = a[0] * b[0]
    out = [jnp.zeros_like(sample) for _ in range(m)]
    for g, al, be in table:
        out[g] = out[g] + a[al] * b[be]
    return jnp.stack(out, axis=0)

Related APIsomnibias.core.multiply_table↗omnibias.core.num_multi_indices↗

Other public import paths

• omnibias.jax.jet_multiply
• omnibias.jax.jet_mv.jet_multiply

FUNCTION

jet_partials

omnibias.jax.jet_partialsSource ↗

jet_partials(jet: Array, dim: int, order: int) -> dict[tuple[int, ...], Array]

Convert a jet to raw partial derivatives {alpha: D^alpha f(x0)}.

Each coefficient row is scaled by alpha! (D^alpha f = alpha! c_alpha).

Parameters

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

Returnsdict[tuple[int, ...], Array]

Explicit validation & failure cases

• When jet.shape[0] != len(idx), raises ValueError(f'jet has {jet.shape[0]} rows but dim={dim}, order={order} requires {len(idx)}').

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

Executed regression · mobius 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_mobius_parity() -> None:
    for tt in _tts(21):
        t = _np(tb.mobius_coeffs(torch.tensor(tt, dtype=torch.float64)))
        j = _np(jb.mobius_coeffs(jnp.asarray(tt, dtype=jnp.float64)))
        assert np.allclose(t, j, rtol=1e-9, atol=1e-11)

python -m pytest "packages/omnibias-boolean/tests/test_cross_backend.py::test_mobius_parity" -q

Regression example · exp matches autodiff

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_exp_matches_autodiff(memory) -> None:
    W, b, x0 = memory["W"], memory["b"], memory["x0"]

    def f(x):
        return jnp.exp(jnp.tanh(W @ x + b))

    jet = mlp_jet_mv(x0, [(W, b, "tanh")], ORDER)
    got = jet_partials(jet_exp(jet, DIM, ORDER), DIM, ORDER)
    for alpha, value in _ad_partials(f, x0).items():
        assert _close(got[alpha], value), alpha

python -m pytest packages/omnibias-jax/tests/test_jet_attention.py -k test_exp_matches_autodiff -q

Read the implementation

def jet_partials(jet: Array, dim: int, order: int) -> dict[tuple[int, ...], Array]:
    """Convert a jet to raw partial derivatives ``{alpha: D^alpha f(x0)}``.

    Each coefficient row is scaled by ``alpha!`` (``D^alpha f = alpha! c_alpha``).
    """
    jet = jnp.asarray(jet)
    idx = multi_indices(dim, order)
    if jet.shape[0] != len(idx):
        raise ValueError(
            f"jet has {jet.shape[0]} rows but dim={dim}, order={order} "
            f"requires {len(idx)}"
        )
    return {alpha: jet[i] * multi_index_factorial(alpha) for i, alpha in enumerate(idx)}

Related APIsomnibias.core.multi_index_factorial↗omnibias.core.multi_indices↗

Other public import paths

• omnibias.jax.jet_partials
• omnibias.jax.jet_mv.jet_partials

FUNCTION

jet_reciprocal

omnibias.jax.jet_reciprocalSource ↗

jet_reciprocal(u_jet: Array, dim: int, order: int) -> Array

Jet of the reciprocal 1 / u from the jet of u.

mathematically equivalent twin of omnibias.torch.jet_mv.jet_reciprocal. 1/u is analytic wherever u(x_0) != 0 and its derivative tower is elementary,

\frac{d^k}{du^k}\,u^{-1} = (-1)^k\, k!\, u^{-(k+1)},

so the reciprocal composes through the *same* compose_jet_mv kernel as any activation -- it is simply a tower this module supplies itself rather than one the activation dictionary carries. This is what makes a *rational* map (a normalisation, a softmax denominator, a quotient ansatz) closed form at arbitrary order instead of a nested-autodiff fallback.

A vanishing u_jet[0] is not an error: the result is then infinite or nan in the usual floating-point way, matching division elsewhere.

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When u_jet.shape[0] != m, raises ValueError(f'u_jet has {u_jet.shape[0]} rows but dim={dim}, order={order} requires {m}').

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

Executed regression · row count is validated

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_row_count_is_validated(memory) -> None:
    jet = identity_jet(memory["x0"], ORDER)
    with pytest.raises(ValueError, match="requires"):
        jet_reciprocal(jet, DIM, ORDER + 1)
    with pytest.raises(ValueError, match="requires"):
        jet_softmax(jet, DIM, ORDER + 1)

python -m pytest "packages/omnibias-jax/tests/test_jet_attention.py::test_row_count_is_validated" -q

Regression example · reciprocal 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_reciprocal_parity(shared, order: int) -> None:
    t, j = _jets(shared, order)
    t_shift = torch.cat([(t[0] + 2.5).unsqueeze(0), t[1:]], dim=0)
    j_shift = jnp.concatenate([(j[0] + 2.5)[None], j[1:]], axis=0)
    assert _agree(
        torch_reciprocal(t_shift, DIM, order), jax_reciprocal(j_shift, DIM, order)
    )

python -m pytest tests/test_jet_attention_parity.py -k test_reciprocal_parity -q

Read the implementation

def jet_reciprocal(u_jet: Array, dim: int, order: int) -> Array:
    r"""Jet of the reciprocal ``1 / u`` from the jet of ``u``.

    Bit-identical twin of :func:`omnibias.torch.jet_mv.jet_reciprocal`. ``1/u`` is
    analytic wherever ``u(x_0) != 0`` and its derivative tower is elementary,

    .. math::

        \frac{d^k}{du^k}\,u^{-1} = (-1)^k\, k!\, u^{-(k+1)},

    so the reciprocal composes through the *same* :func:`compose_jet_mv` kernel as
    any activation -- it is simply a tower this module supplies itself rather than
    one the activation dictionary carries. This is what makes a *rational* map (a
    normalisation, a softmax denominator, a quotient ansatz) closed form at
    arbitrary order instead of a nested-autodiff fallback.

    A vanishing ``u_jet[0]`` is not an error: the result is then infinite or
    ``nan`` in the usual floating-point way, matching division elsewhere.
    """
    u_jet = jnp.asarray(u_jet)
    m = num_multi_indices(dim, order)
    if u_jet.shape[0] != m:
        raise ValueError(
            f"u_jet has {u_jet.shape[0]} rows but dim={dim}, order={order} "
            f"requires {m}"
        )
    inv = 1.0 / u_jet[0]
    rows = []
    power = inv
    coeff = 1.0
    for k in range(order + 1):
        if k > 0:
            power = power * inv
            coeff = -coeff * k
        rows.append(coeff * power)
    return compose_jet_mv(u_jet, jnp.stack(rows, axis=0), dim, order)

Related APIsomnibias.core.num_multi_indices↗omnibias.jax.compose_jet_mv↗

Other public import paths

• omnibias.jax.jet_reciprocal
• omnibias.jax.jet_mv.jet_reciprocal

FUNCTION

jet_softmax

omnibias.jax.jet_softmaxSource ↗

jet_softmax(s_jet: Array, dim: int, order: int) -> Array

Jet of softmax(s) over the trailing axis, from the jet of the scores.

mathematically equivalent twin of omnibias.torch.jet_mv.jet_softmax. Softmax is the first genuinely *non-elementwise* map in this module: every output couples to every score through the shared denominator. It still factors into primitives that are each exact --

\mathrm{softmax}(s)_j = e^{s_j} \cdot \Big(\sum_i e^{s_i}\Big)^{-1},

an jet_exp, a row-wise sum (linear, so it acts per coefficient), a jet_reciprocal, and a jet_multiply -- so the composite jet is exact to machine precision at arbitrary order, with no nested autodiff.

s_jet has shape (M, ..., n); the softmax is taken over the last axis. The constant coefficient is max-shifted before exponentiating, exactly as a stable softmax implementation does. The shift is held constant for the gradient and is identical across the axis, so it cancels in the mathematics and only removes overflow.

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
s_jetpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When s_jet.shape[0] != m, raises ValueError(f's_jet has {s_jet.shape[0]} rows but dim={dim}, order={order} requires {m}').
• When s_jet.ndim < 2, raises ValueError(f's_jet must carry a trailing softmax axis, got shape {s_jet.shape}').

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

Executed regression · softmax rejects a scalar jet

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_softmax_rejects_a_scalar_jet(memory) -> None:
    x0 = memory["x0"]
    with pytest.raises(ValueError, match="trailing softmax axis"):
        jet_softmax(identity_jet(x0, ORDER)[:, 0], DIM, ORDER)

python -m pytest "packages/omnibias-jax/tests/test_jet_attention.py::test_softmax_rejects_a_scalar_jet" -q

Regression example · row count is validated

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_row_count_is_validated(memory) -> None:
    jet = identity_jet(memory["x0"], ORDER)
    with pytest.raises(ValueError, match="requires"):
        jet_reciprocal(jet, DIM, ORDER + 1)
    with pytest.raises(ValueError, match="requires"):
        jet_softmax(jet, DIM, ORDER + 1)

python -m pytest packages/omnibias-jax/tests/test_jet_attention.py -k test_row_count_is_validated -q

Read the implementation

def jet_softmax(s_jet: Array, dim: int, order: int) -> Array:
    r"""Jet of ``softmax(s)`` over the trailing axis, from the jet of the scores.

    Bit-identical twin of :func:`omnibias.torch.jet_mv.jet_softmax`. Softmax is the
    first genuinely *non-elementwise* map in this module: every output couples to
    every score through the shared denominator. It still factors into primitives
    that are each exact --

    .. math::

        \mathrm{softmax}(s)_j = e^{s_j} \cdot \Big(\sum_i e^{s_i}\Big)^{-1},

    an :func:`jet_exp`, a row-wise sum (linear, so it acts per coefficient), a
    :func:`jet_reciprocal`, and a :func:`jet_multiply` -- so the composite jet is
    exact to machine precision at arbitrary order, with no nested autodiff.

    ``s_jet`` has shape ``(M, ..., n)``; the softmax is taken over the last axis.
    The constant coefficient is max-shifted before exponentiating, exactly as a
    stable softmax implementation does. The shift is held constant for the
    gradient and is identical across the axis, so it cancels in the mathematics
    and only removes overflow.
    """
    s_jet = jnp.asarray(s_jet)
    m = num_multi_indices(dim, order)
    if s_jet.shape[0] != m:
        raise ValueError(
            f"s_jet has {s_jet.shape[0]} rows but dim={dim}, order={order} "
            f"requires {m}"
        )
    if s_jet.ndim < 2:
        raise ValueError(
            f"s_jet must carry a trailing softmax axis, got shape {s_jet.shape}"
        )
    shift = jnp.max(jax.lax.stop_gradient(s_jet[0]), axis=-1, keepdims=True)
    centred = jnp.concatenate([(s_jet[0] - shift)[None, ...], s_jet[1:]], axis=0)
    e_jet = jet_exp(centred, dim, order)
    total = jnp.sum(e_jet, axis=-1, keepdims=True)  # linear: acts per coefficient
    return jet_multiply(e_jet, jet_reciprocal(total, dim, order), dim, order)

Related APIsomnibias.core.num_multi_indices↗omnibias.jax.jet_exp↗omnibias.jax.jet_multiply↗omnibias.jax.jet_reciprocal↗

Other public import paths

• omnibias.jax.jet_softmax
• omnibias.jax.jet_mv.jet_softmax

FUNCTION

layer_jet_mv

omnibias.jax.layer_jet_mvSource ↗

layer_jet_mv(z_jet: Array, W: Array, b: Array | None, spec: JaxActivationSpec | str, dim: int, order: int) -> Array

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

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.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	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

Read the implementation

def layer_jet_mv(
    z_jet: Array,
    W: Array,
    b: Array | None,
    spec: JaxActivationSpec | str,
    dim: int,
    order: int,
) -> Array:
    """Push a multivariate jet through one ``sigma(W z + b)`` layer."""
    z_jet = jnp.asarray(z_jet)
    resolved = get_activation(spec)
    u_jet = affine_jet(z_jet, W, b)
    sigma_tower = _sigma_tower(resolved, u_jet[0], order)
    return compose_jet_mv(u_jet, sigma_tower, dim, order)

Related APIsomnibias.jax.get_activation↗omnibias.jax.affine_jet↗omnibias.jax.compose_jet_mv↗

Other public import paths

• omnibias.jax.layer_jet_mv
• omnibias.jax.jet_mv.layer_jet_mv

FUNCTION

mlp_jet_mv

omnibias.jax.mlp_jet_mvSource ↗

mlp_jet_mv(x0: Array, layers: Sequence[tuple[Array, Array | None, JaxActivationSpec | str | None]], order: int) -> Array

Exact multivariate Taylor jet of a deep MLP around x0.

Parameters

x0 Base point, shape (D,). layers Sequence of (W, b, spec). spec=None is a pure affine readout; otherwise the layer computes sigma(W z + b). order Truncation total order N; the returned jet has num_multi_indices(D, N) rows.

Returns

Array Jet of shape (M, C) whose row i is the Taylor coefficient D^alpha f(x0) / alpha! for alpha = multi_indices(D, N)[i]. Use jet_partials, jet_gradient or jet_hessian to read out raw derivatives.

Use a multivariate jet for mixed partial derivatives in several coordinates. The row ordering follows multi_indices(D, order), and each coefficient is divided by the multi-index factorial. Read raw derivatives with the matching jet_partials, jet_gradient, or jet_hessian helpers. The number of coefficients grows combinatorially with dimension and order; consider a directional jet or direct contraction when you do not need the complete mixed-partial set.

Parameters

Name	Type & default	Contract
x0positional-or-keyword	ArrayRequired	One input point, shape (D,). Its coordinate order determines every multi-index in the result.
layerspositional-or-keyword	Sequence[tuple[Array, Array | None, JaxActivationSpec | str | None]]Required	Sequence of (W, b, activation) triples with compatible adjacent widths. A None activation is an affine readout.
orderpositional-or-keyword	intRequired	Maximum total multi-index degree N. The result has num_multi_indices(D, N) rows and one column per output component.

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 · exp 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_exp_parity(shared, order: int) -> None:
    t, j = _jets(shared, order)
    assert _agree(torch_exp(t, DIM, order), jax_exp(j, DIM, order))

python -m pytest tests/test_jet_attention_parity.py -k test_exp_parity -q

Read the implementation

def mlp_jet_mv(
    x0: Array,
    layers: Sequence[tuple[Array, Array | None, JaxActivationSpec | str | None]],
    order: int,
) -> Array:
    """Exact multivariate Taylor jet of a deep MLP around ``x0``.

    Parameters
    ----------
    x0
        Base point, shape ``(D,)``.
    layers
        Sequence of ``(W, b, spec)``. ``spec=None`` is a pure affine readout;
        otherwise the layer computes ``sigma(W z + b)``.
    order
        Truncation total order ``N``; the returned jet has
        ``num_multi_indices(D, N)`` rows.

    Returns
    -------
    Array
        Jet of shape ``(M, C)`` whose row ``i`` is the Taylor coefficient
        ``D^alpha f(x0) / alpha!`` for ``alpha = multi_indices(D, N)[i]``. Use
        :func:`jet_partials`, :func:`jet_gradient` or :func:`jet_hessian` to read
        out raw derivatives.
    """
    x0 = jnp.asarray(x0)
    dim = int(x0.shape[-1])
    jet = identity_jet(x0, order)
    for W, b, spec in layers:
        if spec is None:
            jet = affine_jet(jet, W, b)
        else:
            jet = layer_jet_mv(jet, W, b, spec, dim, order)
    return jet

Related APIsomnibias.jax.affine_jet↗omnibias.jax.identity_jet↗omnibias.jax.layer_jet_mv↗

Other public import paths

• omnibias.jax.mlp_jet_mv
• omnibias.jax.jet_mv.mlp_jet_mv

FUNCTION

mlp_joint_jet_mv

omnibias.jax.jet_mv.mlp_joint_jet_mvSource ↗

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

Propagate joint input/parameter Taylor jets in an explicit finite basis.

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.
dimkeyword-only	intRequired	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_mv(
    input_jet: Array,
    parameter_jet: Array,
    spec: RealizationSpec,
    *,
    dim: int,
    order: int,
    max_coefficients: int = 1_000_000,
) -> Array:
    """Propagate joint input/parameter Taylor jets in an explicit finite basis."""
    from omnibias.jax.realization import realization_jet

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

FUNCTION

parameter_jet

omnibias.jax.jet_mv.parameter_jetSource ↗

parameter_jet(x: Array, theta: Array, spec: RealizationSpec, basis: Array, order: int, *, input_basis: Array | None=None, max_coefficients: int=1000000) -> Array

Seed live joint Taylor jets from explicit parameter and input directions.

basis is P-by-q; input_basis broadcasts to x.shape + (q,). The coefficient budget is checked by the realization implementation.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
thetapositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
specpositional-or-keyword	RealizationSpecRequired	See the operation contract and implementation below.
basispositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
input_basiskeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
max_coefficientskeyword-only	intDefault: 1000000	See the operation contract and implementation below.

ReturnsArray

Regression example · existing multivariate parameter jet entrypoints preserve all live directions

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_existing_multivariate_parameter_jet_entrypoints_preserve_all_live_directions():
    spec = RealizationSpec.dense((2, 2, 1))
    x = torch.tensor([0.3, -0.2], dtype=torch.float64, requires_grad=True)
    theta = torch.linspace(-0.4, 0.6, spec.layout.size, dtype=torch.float64, requires_grad=True)
    basis = torch.full((spec.layout.size, 2), 0.15, dtype=torch.float64, requires_grad=True)
    input_basis = torch.tensor([[0.2, 0.1], [-0.1, 0.4]], dtype=torch.float64, requires_grad=True)
    args = (x, theta, basis, input_basis)

    def torch_call(fn):
        return fn(x, theta, spec, basis, 2, input_basis=input_basis, max_coefficients=1000)

    actual, expected = torch_call(tmv.parameter_jet), torch_call(tr.parameter_jet)
    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
    gradients = torch.autograd.grad(actual.square().sum(), args, retain_graph=True)
    reference = torch.autograd.grad(expected.square().sum(), args)
    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)

    jargs = tuple(jnp.asarray(arg.detach().numpy()) for arg in args)

    def loss(fn, a, b, directions, input_directions):
        return jnp.sum(
            fn(a, b, spec, directions, 2, input_basis=input_directions, max_coefficients=1000) ** 2
        )

    first = jax.jit(jax.grad(lambda *a: loss(jmv.parameter_jet, *a), argnums=(0, 1, 2, 3)))(*jargs)
    second = jax.jit(jax.grad(lambda *a: loss(jr.parameter_jet, *a), argnums=(0, 1, 2, 3)))(*jargs)
    for got, want in zip(first, second, strict=True):
        assert np.isfinite(got).all() and np.count_nonzero(got)
        np.testing.assert_array_equal(got, want)
    with pytest.raises(ValueError, match="budget"):
        tmv.parameter_jet(x, theta, spec, basis, 2, max_coefficients=1)
    with pytest.raises(ValueError, match="budget"):
        jmv.parameter_jet(jargs[0], jargs[1], spec, jargs[2], 2, max_coefficients=1)

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

Read the implementation

def parameter_jet(
    x: Array,
    theta: Array,
    spec: RealizationSpec,
    basis: Array,
    order: int,
    *,
    input_basis: Array | None = None,
    max_coefficients: int = 1_000_000,
) -> Array:
    """Seed live joint Taylor jets from explicit parameter and input directions.

    ``basis`` is P-by-q; ``input_basis`` broadcasts to x.shape + (q,).
    The coefficient budget is checked by the realization implementation.
    """
    from omnibias.jax.realization import parameter_jet as impl

    return impl(
        x, theta, spec, basis, order, input_basis=input_basis,
        max_coefficients=max_coefficients,
    )

← Package overviewRunnable learning examples →