API REFERENCE/omnibias-torch

omnibias.torch.jet_mv

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.jet_mv

Module contract & mathematical conventions

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

Bit-identical twin of omnibias.jax.jet_mv; see that module for the full mathematical exposition. A multivariate jet truncated at total order N over D variables stores every Taylor coefficient c_alpha = D^alpha f(x0)/alpha! densely along the leading axis (canonical ~omnibias.core.multi_index.multi_indices order), so a single forward pass yields all mixed partial derivatives up to order N. Affine maps act per coefficient; activations use the shifted-power identity with *multivariate* truncated polynomial multiplication driven by omnibias.core.multi_index.multiply_table.

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.torch.affine_jet_mvSource ↗

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

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

Identical per-coefficient action to omnibias.torch.affine_jet: the map is degree-1 in the perturbation, so it acts row-wise with the bias on the constant coefficient (row 0) only.

Parameters

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

ReturnsTensor

Executed regression · 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 = torch.as_tensor(rng.normal(size=(m, 5)))
    W = torch.as_tensor(rng.normal(size=(3, 5)))
    b = torch.as_tensor(rng.normal(size=(3,)))
    out = affine_jet_mv(z_jet, W, b)
    assert torch.allclose(out[0], W @ z_jet[0] + b)
    for k in range(1, m):
        assert torch.allclose(out[k], W @ z_jet[k])

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

Read the implementation

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

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

Related APIsomnibias.torch.affine_jet↗

Other public import paths

• omnibias.torch.affine_jet_mv
• omnibias.torch.jet_mv.affine_jet_mv

FUNCTION

affine_joint_jet_mv

omnibias.torch.jet_mv.affine_joint_jet_mvSource ↗

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

Live multivariate product with jet-valued weights and biases.

Parameters

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

ReturnsTensor

Read the implementation

def affine_joint_jet_mv(
    input_jet: Tensor, weight_jet: Tensor, bias_jet: Tensor | None = None,
    *, dim: int, order: int,
) -> Tensor:
    """Live multivariate product with jet-valued weights and biases."""
    from omnibias.torch.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.torch.compose_jet_mvSource ↗

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

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

u_jet has shape (M, ...) with M = num_multi_indices(dim, order); sigma_tower[k] = sigma^(k)(u_jet[0]) for k = 0..order. Returns the jet of sigma(u).

Parameters

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

ReturnsTensor

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 · hessian requires order two

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_hessian_requires_order_two() -> None:
    layers, x0 = _build_mlp(seed=6)
    jet = mlp_jet_mv(x0, layers, 1)
    with pytest.raises(ValueError, match="hessian needs order >= 2"):
        jet_hessian(jet, int(x0.shape[0]), 1)

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

Read the implementation

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

    ``u_jet`` has shape ``(M, ...)`` with ``M = num_multi_indices(dim, order)``;
    ``sigma_tower[k] = sigma^(k)(u_jet[0])`` for ``k = 0..order``. Returns the
    jet of ``sigma(u)``.
    """
    u_jet = torch.as_tensor(u_jet)
    sigma_tower = torch.as_tensor(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 = torch.zeros_like(u_jet[0])
    w = [zero] + [u_jet[g] for g in range(1, m)]
    p = [torch.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
        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 torch.stack(result, dim=0)

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

Other public import paths

• omnibias.torch.compose_jet_mv
• omnibias.torch.jet_mv.compose_jet_mv

FUNCTION

identity_jet

omnibias.torch.identity_jetSource ↗

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

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

Returns shape (M, D): row 0 is x0 and each unit-multi-index row e_i is the standard basis vector e_i; all other rows are zero.

Parameters

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

ReturnsTensor

Explicit validation & failure cases

• When x0.ndim != 1, raises ValueError(f'x0 must be a 1-D vector, got shape {tuple(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 · influence penalty is differentiable

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_influence_penalty_is_differentiable() -> None:
    vals = torch.rand(8, dtype=torch.float64, requires_grad=True)
    design.influence_penalty(vals).backward()
    assert vals.grad is not None

python -m pytest "packages/omnibias-boolean/tests/test_design.py::test_influence_penalty_is_differentiable" -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-torch/tests/test_jet_attention.py -k test_softmax_rejects_a_scalar_jet -q

Read the implementation

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

    Returns shape ``(M, D)``: row 0 is ``x0`` and each unit-multi-index row
    ``e_i`` is the standard basis vector ``e_i``; all other rows are zero.
    """
    x0 = torch.as_tensor(x0)
    if x0.ndim != 1:
        raise ValueError(f"x0 must be a 1-D vector, got shape {tuple(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 = [torch.zeros((dim,), dtype=x0.dtype, device=x0.device) for _ in range(m)]
    rows[pos[(0,) * dim]] = x0
    if order >= 1:
        eye = torch.eye(dim, dtype=x0.dtype, device=x0.device)
        for i in range(dim):
            e = tuple(1 if j == i else 0 for j in range(dim))
            rows[pos[e]] = eye[i]
    return torch.stack(rows, dim=0)

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

Other public import paths

• omnibias.torch.identity_jet
• omnibias.torch.jet_mv.identity_jet

FUNCTION

jet_attention

omnibias.torch.jet_attentionSource ↗

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

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

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.torch.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. A tensor is allowed, which keeps a *trainable* temperature closed form in exactly the same way.

Parameters

Name	Type & default	Contract
q_jetpositional-or-keyword	TensorRequired	Jet of the query, shape (M, ..., d_key).
keyspositional-or-keyword	TensorRequired	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	TensorRequired	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 | TensorDefault: 1.0	Inverse temperature. A tensor is allowed, which keeps a *trainable* temperature closed form in exactly the same way.

ReturnsTensor

Explicit validation & failure cases

• When keys.ndim != 2 or values.ndim != 2, raises ValueError(f'keys and values must be 2-D, got {tuple(keys.shape)} and {tuple(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 validates shapes

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_validates_shapes(memory, keys, values, match) -> None:
    W, b, x0 = memory["W"], memory["b"], memory["x0"]
    q_jet = mlp_jet_mv(x0, [(W, b, "tanh")], ORDER)
    with pytest.raises(ValueError, match=match):
        jet_attention(q_jet, keys.double(), values.double(), DIM, ORDER)

python -m pytest "packages/omnibias-torch/tests/test_jet_attention.py::test_attention_validates_shapes[keys0-values0-must be 2-D]" -q

Regression example · softmax survives scores that overflow a naive exp

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_survives_scores_that_overflow_a_naive_exp(memory) -> None:
    W, b, K, x0 = memory["W"], memory["b"], memory["K"], memory["x0"]
    scores = mlp_jet_mv(x0, [(W, b, "tanh"), (K * 400.0, None, None)], ORDER)
    p = jet_softmax(scores, DIM, ORDER)
    assert bool(torch.isfinite(p).all())
    assert abs(float(p[0].sum()) - 1.0) < 1e-14

python -m pytest packages/omnibias-torch/tests/test_jet_attention.py -k test_softmax_survives_scores_that_overflow_a_naive_exp -q

Read the implementation

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

    :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.torch.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. A tensor is allowed, which keeps a *trainable*
        temperature closed form in exactly the same way.
    """
    q_jet = torch.as_tensor(q_jet)
    keys = torch.as_tensor(keys)
    values = torch.as_tensor(values)
    if keys.ndim != 2 or values.ndim != 2:
        raise ValueError(
            f"keys and values must be 2-D, got {tuple(keys.shape)} and "
            f"{tuple(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, values.transpose(-2, -1))

Related APIsomnibias.torch.affine_jet↗omnibias.torch.jet_softmax↗

Other public import paths

• omnibias.torch.jet_attention
• omnibias.torch.jet_mv.jet_attention

FUNCTION

jet_exp

omnibias.torch.jet_expSource ↗

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

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

ReturnsTensor

Executed regression · softmax survives scores that overflow a naive exp

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_survives_scores_that_overflow_a_naive_exp(memory) -> None:
    W, b, K, x0 = memory["W"], memory["b"], memory["K"], memory["x0"]
    scores = mlp_jet_mv(x0, [(W, b, "tanh"), (K * 400.0, None, None)], ORDER)
    p = jet_softmax(scores, DIM, ORDER)
    assert bool(torch.isfinite(p).all())
    assert abs(float(p[0].sum()) - 1.0) < 1e-14

python -m pytest "packages/omnibias-torch/tests/test_jet_attention.py::test_softmax_survives_scores_that_overflow_a_naive_exp" -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: Tensor, dim: int, order: int) -> Tensor:
    """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 = torch.as_tensor(u_jet)
    tower = _sigma_tower(get_activation("exp"), u_jet[0], order)
    return compose_jet_mv(u_jet, tower, dim, order)

Related APIsomnibias.torch.get_activation↗omnibias.torch.compose_jet_mv↗

Other public import paths

• omnibias.torch.jet_exp
• omnibias.torch.jet_mv.jet_exp

FUNCTION

jet_gradient

omnibias.torch.jet_gradientSource ↗

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

Gradient d f / d x_i from a jet, shape (D, ...). Needs order >= 1.

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	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.

ReturnsTensor

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 · gradient requires order one

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_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-torch/tests/test_jet_mv.py::test_gradient_requires_order_one" -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:
    """jet(x) * jet(y) is the jet of x*y: grad = [y, x], Hessian = [[0,1],[1,0]]."""
    x0 = torch.tensor([0.3, -0.7], dtype=torch.float64)
    idj = identity_jet(x0, 2)
    prod = jet_multiply(idj[:, 0], idj[:, 1], 2, 2)
    assert torch.allclose(jet_gradient(prod, 2, 1), torch.tensor([-0.7, 0.3]).double())
    h = jet_hessian(prod, 2, 2)
    assert torch.allclose(h, torch.tensor([[0.0, 1.0], [1.0, 0.0]]).double())

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

Read the implementation

def jet_gradient(jet: Tensor, dim: int, order: int) -> Tensor:
    """Gradient ``d f / d x_i`` from a jet, shape ``(D, ...)``. Needs ``order >= 1``."""
    if order < 1:
        raise ValueError(f"gradient needs order >= 1, got {order}")
    jet = torch.as_tensor(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 torch.stack(rows, dim=0)

Related APIsomnibias.core.index_position↗

Other public import paths

• omnibias.torch.jet_gradient
• omnibias.torch.jet_mv.jet_gradient

FUNCTION

jet_hessian

omnibias.torch.jet_hessianSource ↗

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

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

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

Parameters

Name	Type & default	Contract
jetpositional-or-keyword	TensorRequired	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.

ReturnsTensor

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 · hessian requires order two

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_hessian_requires_order_two() -> None:
    layers, x0 = _build_mlp(seed=6)
    jet = mlp_jet_mv(x0, layers, 1)
    with pytest.raises(ValueError, match="hessian needs order >= 2"):
        jet_hessian(jet, int(x0.shape[0]), 1)

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

Regression example · value grad hessian consistent

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_grad_hessian_consistent() -> None:
    net = JetMLP(2, 10, 1, depth=2, base="tanh").double()
    field = homogeneous_box(net, [0.0, 0.0], [1.0, 1.0])
    x = torch.rand(5, 2, dtype=torch.float64)
    v, g, h = field.value_grad_hessian(x)
    assert torch.allclose(v, field.value(x), atol=1e-12)
    assert torch.allclose(g, field.gradient(x), atol=1e-12)
    assert torch.allclose(h, field.hessian(x), atol=1e-12)

python -m pytest packages/omnibias-torch/tests/test_hardbc.py -k test_value_grad_hessian_consistent -q

Read the implementation

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

    Needs ``order >= 2``. ``H[i, j]`` is the coefficient of ``alpha = e_i + e_j``
    rescaled by ``alpha!`` (2 on the diagonal, 1 off-diagonal).
    """
    if order < 2:
        raise ValueError(f"hessian needs order >= 2, got {order}")
    jet = torch.as_tensor(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(torch.stack(row, dim=0))
    return torch.stack(rows, dim=0)

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

Other public import paths

• omnibias.torch.jet_hessian
• omnibias.torch.jet_mv.jet_hessian

FUNCTION

jet_multiply

omnibias.torch.jet_multiplySource ↗

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

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

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.

Parameters

Name	Type & default	Contract
apositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	TensorRequired	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.

ReturnsTensor

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

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

Regression example · jet multiply bit 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_jet_multiply_bit_parity() -> None:
    xt, xj = _xpair(seed=0, b=1, d=2)
    it = torch_identity_jet(xt[0], 3)
    ij = jax_identity_jet(xj[0], 3)
    pt = torch_jet_multiply(it[:, 0], it[:, 1], 2, 3)
    pj = jax_jet_multiply(ij[:, 0], ij[:, 1], 2, 3)
    assert np.allclose(pt.detach().numpy(), np.asarray(pj), atol=1e-13)

python -m pytest tests/test_hardbc_parity.py -k test_jet_multiply_bit_parity -q

Read the implementation

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

    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 = torch.as_tensor(a)
    b = torch.as_tensor(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 = [torch.zeros_like(sample) for _ in range(m)]
    for g, al, be in table:
        out[g] = out[g] + a[al] * b[be]
    return torch.stack(out, dim=0)

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

Other public import paths

• omnibias.torch.jet_multiply
• omnibias.torch.jet_mv.jet_multiply

FUNCTION

jet_partials

omnibias.torch.jet_partialsSource ↗

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

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	TensorRequired	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, ...], Tensor]

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 · influence penalty is differentiable

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_influence_penalty_is_differentiable() -> None:
    vals = torch.rand(8, dtype=torch.float64, requires_grad=True)
    design.influence_penalty(vals).backward()
    assert vals.grad is not None

python -m pytest "packages/omnibias-boolean/tests/test_design.py::test_influence_penalty_is_differentiable" -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 torch.exp(torch.tanh(W @ x + b))

    jet = mlp_jet_mv(x0, [(W, b, "tanh")], ORDER)
    got = jet_partials(jet_exp(jet, DIM, ORDER), DIM, ORDER)
    want = _ad_partials(f, x0)
    for alpha, value in want.items():
        assert torch.allclose(got[alpha], value, rtol=1e-11, atol=1e-11), alpha

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

Read the implementation

def jet_partials(jet: Tensor, dim: int, order: int) -> dict[tuple[int, ...], Tensor]:
    """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 = torch.as_tensor(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.torch.jet_partials
• omnibias.torch.jet_mv.jet_partials

FUNCTION

jet_reciprocal

omnibias.torch.jet_reciprocalSource ↗

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

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

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.

Raises nothing on a vanishing u_jet[0]: the result is then infinite or nan in the usual floating-point way, matching division elsewhere.

Parameters

Name	Type & default	Contract
u_jetpositional-or-keyword	TensorRequired	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.

ReturnsTensor

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:
    x0 = memory["x0"]
    jet = identity_jet(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-torch/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: Tensor, dim: int, order: int) -> Tensor:
    r"""Jet of the reciprocal ``1 / u`` from the jet of ``u``.

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

    Raises nothing on a vanishing ``u_jet[0]``: the result is then infinite or
    ``nan`` in the usual floating-point way, matching division elsewhere.
    """
    u_jet = torch.as_tensor(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, torch.stack(rows, dim=0), dim, order)

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

Other public import paths

• omnibias.torch.jet_reciprocal
• omnibias.torch.jet_mv.jet_reciprocal

FUNCTION

jet_softmax

omnibias.torch.jet_softmaxSource ↗

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

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

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 detached and identical across the axis, so it cancels in the mathematics and only removes overflow.

Parameters

Name	Type & default	Contract
s_jetpositional-or-keyword	TensorRequired	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.

ReturnsTensor

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 {tuple(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-torch/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:
    x0 = memory["x0"]
    jet = identity_jet(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-torch/tests/test_jet_attention.py -k test_row_count_is_validated -q

Read the implementation

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

    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 detached and identical
    across the axis, so it cancels in the mathematics and only removes overflow.
    """
    s_jet = torch.as_tensor(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(
            "s_jet must carry a trailing softmax axis, got shape "
            f"{tuple(s_jet.shape)}"
        )
    shift = s_jet[0].detach().max(dim=-1, keepdim=True).values
    centred = torch.cat([(s_jet[0] - shift).unsqueeze(0), s_jet[1:]], dim=0)
    e_jet = jet_exp(centred, dim, order)
    total = e_jet.sum(dim=-1, keepdim=True)  # linear: acts per coefficient
    return jet_multiply(e_jet, jet_reciprocal(total, dim, order), dim, order)

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

Other public import paths

• omnibias.torch.jet_softmax
• omnibias.torch.jet_mv.jet_softmax

FUNCTION

layer_jet_mv

omnibias.torch.layer_jet_mvSource ↗

layer_jet_mv(z_jet: Tensor, W: Tensor, b: Tensor | None, spec: ActivationSpec[Tensor] | str, dim: int, order: int) -> Tensor

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

Parameters

Name	Type & default	Contract
z_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | NoneRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor] | strRequired	See the operation contract and implementation below.
dimpositional-or-keyword	intRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · hessian requires order two

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_hessian_requires_order_two() -> None:
    layers, x0 = _build_mlp(seed=6)
    jet = mlp_jet_mv(x0, layers, 1)
    with pytest.raises(ValueError, match="hessian needs order >= 2"):
        jet_hessian(jet, int(x0.shape[0]), 1)

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

Read the implementation

def layer_jet_mv(
    z_jet: Tensor,
    W: Tensor,
    b: Tensor | None,
    spec: ActivationSpec[Tensor] | str,
    dim: int,
    order: int,
) -> Tensor:
    """Push a multivariate jet through one ``sigma(W z + b)`` layer."""
    z_jet = torch.as_tensor(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.torch.get_activation↗omnibias.torch.affine_jet↗omnibias.torch.compose_jet_mv↗

Other public import paths

• omnibias.torch.layer_jet_mv
• omnibias.torch.jet_mv.layer_jet_mv

FUNCTION

mlp_jet_mv

omnibias.torch.mlp_jet_mvSource ↗

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

Exact multivariate Taylor jet of a deep MLP around x0.

layers is a sequence of (W, b, spec); spec=None is a pure affine readout. Returns a jet of shape (M, C) whose row i is D^alpha f(x0) / alpha! for alpha = multi_indices(D, order)[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	TensorRequired	One input point, shape (D,). Its coordinate order determines every multi-index in the result.
layerspositional-or-keyword	Sequence[tuple[Tensor, Tensor | None, ActivationSpec[Tensor] | 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.

ReturnsTensor

Complete runnable example · mixed heat

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

Download complete script ↓

"""Mixed Taylor coefficients for a heat-equation residual at one point."""
import torch
from omnibias.torch.jet_mv import jet_partials, mlp_jet_mv

weights = torch.tensor([[0.5, 0.7], [-0.2, 0.3]], dtype=torch.float64,
                       requires_grad=True)
bias = torch.tensor([0.1, -0.1], dtype=torch.float64)
readout = torch.tensor([[0.6, -0.4]], dtype=torch.float64)
layers = [(weights, bias, "tanh"), (readout, None, None)]
point = torch.tensor([0.2, 0.4], dtype=torch.float64)  # coordinates (t, x)
jet = mlp_jet_mv(point, layers, order=2)
partials = jet_partials(jet, dim=2, order=2)
residual = partials[(1, 0)] - 0.1 * partials[(0, 2)]
residual.square().mean().backward()
assert partials[(1, 1)].shape == (1,)
assert weights.grad is not None and torch.isfinite(weights.grad).all()
print("Mixed partials:", sorted(partials), "finite parameter gradients: True")

python mixed_heat.py

Executed regression · hessian requires order two

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_hessian_requires_order_two() -> None:
    layers, x0 = _build_mlp(seed=6)
    jet = mlp_jet_mv(x0, layers, 1)
    with pytest.raises(ValueError, match="hessian needs order >= 2"):
        jet_hessian(jet, int(x0.shape[0]), 1)

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

Read the implementation

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

    ``layers`` is a sequence of ``(W, b, spec)``; ``spec=None`` is a pure affine
    readout. Returns a jet of shape ``(M, C)`` whose row ``i`` is
    ``D^alpha f(x0) / alpha!`` for ``alpha = multi_indices(D, order)[i]``. Use
    :func:`jet_partials`, :func:`jet_gradient` or :func:`jet_hessian` to read
    out raw derivatives.
    """
    x0 = torch.as_tensor(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.torch.affine_jet↗omnibias.torch.identity_jet↗omnibias.torch.layer_jet_mv↗

Other public import paths

• omnibias.torch.mlp_jet_mv
• omnibias.torch.jet_mv.mlp_jet_mv

FUNCTION

mlp_joint_jet_mv

omnibias.torch.jet_mv.mlp_joint_jet_mvSource ↗

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

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

Parameters

Name	Type & default	Contract
input_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
parameter_jetpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	RealizationSpecRequired	See the operation contract and implementation below.
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.

ReturnsTensor

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

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

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

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

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

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

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

Read the implementation

def mlp_joint_jet_mv(
    input_jet: Tensor,
    parameter_jet: Tensor,
    spec: RealizationSpec,
    *,
    dim: int,
    order: int,
    max_coefficients: int = 1_000_000,
) -> Tensor:
    """Propagate joint input/parameter Taylor jets in an explicit finite basis."""
    from omnibias.torch.realization import realization_jet

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

FUNCTION

parameter_jet

omnibias.torch.jet_mv.parameter_jetSource ↗

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

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	TensorRequired	See the operation contract and implementation below.
thetapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
specpositional-or-keyword	RealizationSpecRequired	See the operation contract and implementation below.
basispositional-or-keyword	TensorRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
input_basiskeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.
max_coefficientskeyword-only	intDefault: 1000000	See the operation contract and implementation below.

ReturnsTensor

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: Tensor,
    theta: Tensor,
    spec: RealizationSpec,
    basis: Tensor,
    order: int,
    *,
    input_basis: Tensor | None = None,
    max_coefficients: int = 1_000_000,
) -> Tensor:
    """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.torch.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 →