omnibias.jax.realization

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.realization

Module contract & mathematical conventions

Live input/weight realization jets and the six activation operator roles.

All coefficient arrays use the shared multi-index ordering and contain Taylor coefficients, not raw derivatives. Only static shapes and descriptions are inspected on the host; returned arrays keep their parameter graphs intact.

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 →

CLASS

Realization

omnibias.jax.realization.RealizationSource ↗

class Realization

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
`spec`dataclass field	`RealizationSpec`Required	See the operation contract and implementation below.

Fields & methods

specfield

spec: RealizationSpec

Declared RealizationSpec field on Realization. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

evaluatemethod

evaluate(self, x: Array, theta: Array) -> Array

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

Returns: Array

Method implementation

    def evaluate(self, x: Array, theta: Array) -> Array:
        return evaluate(x, theta, self.spec)

View source ↗

jvpmethod

jvp(self, x: Array, theta: Array, direction: Array) -> Array

Closed-form directional parameter jet, with live coefficients.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`direction`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

Returns: Array

Method implementation

    def jvp(self, x: Array, theta: Array, direction: Array) -> Array:
        """Closed-form directional parameter jet, with live coefficients."""
        return parameter_jet(x, theta, self.spec, direction.reshape(-1, 1), 1)[1]

View source ↗

vjpmethod

vjp(self, x: Array, theta: Array, cotangent: Array) -> Array

First-order backend reverse-mode adjoint of the live evaluation.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`cotangent`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

Returns: Array

Method implementation

    def vjp(self, x: Array, theta: Array, cotangent: Array) -> Array:
        """First-order backend reverse-mode adjoint of the live evaluation."""
        _, pullback = jax.vjp(lambda p: evaluate(x, p, self.spec), theta)
        return cast(Array, pullback(cotangent)[0])

View source ↗

Read the implementation

class Realization:
    spec: RealizationSpec

    def evaluate(self, x: Array, theta: Array) -> Array:
        return evaluate(x, theta, self.spec)

    def jvp(self, x: Array, theta: Array, direction: Array) -> Array:
        """Closed-form directional parameter jet, with live coefficients."""
        return parameter_jet(x, theta, self.spec, direction.reshape(-1, 1), 1)[1]

    def vjp(self, x: Array, theta: Array, cotangent: Array) -> Array:
        """First-order backend reverse-mode adjoint of the live evaluation."""
        _, pullback = jax.vjp(lambda p: evaluate(x, p, self.spec), theta)
        return cast(Array, pullback(cotangent)[0])

FUNCTION

affine_joint_jet

omnibias.jax.realization.affine_joint_jetSource ↗

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

Directional restriction of affine_joint_jet_mv.

Parameters

Name	Type & default	Contract
`input_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`weight_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`bias_jet`positional-or-keyword	`Array \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def affine_joint_jet(input_jet: Array, weight_jet: Array, bias_jet: Array | None = None) -> Array:
    """Directional restriction of :func:`affine_joint_jet_mv`."""
    return affine_joint_jet_mv(input_jet, weight_jet, bias_jet, dim=1, order=input_jet.shape[0] - 1)

FUNCTION

affine_joint_jet_mv

omnibias.jax.realization.affine_joint_jet_mvSource ↗

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

Cauchy product for a layer whose input, weights and bias all vary.

Parameters

Name	Type & default	Contract
`input_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`weight_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`bias_jet`positional-or-keyword	`Array \| None`Default: None	See the operation contract and implementation below.
`dim`keyword-only	`int`Required	See the operation contract and implementation below.
`order`keyword-only	`int`Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When input_jet.shape[0] != count or weight_jet.shape[0] != count, raises ValueError('input and weight jets must have the specified coefficient count').
When weight_jet.ndim != 3 or input_jet.shape[-1] != weight_jet.shape[-1], raises ValueError('weight jet must be (M, out, in) with matching input width').
When bias_jet.shape != (count, weight_jet.shape[1]), raises ValueError('bias jet must be (M, out)').

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

Read the implementation

def affine_joint_jet_mv(
    input_jet: Array, weight_jet: Array, bias_jet: Array | None = None,
    *, dim: int, order: int,
) -> Array:
    """Cauchy product for a layer whose input, weights and bias all vary."""
    count = num_multi_indices(dim, order)
    if input_jet.shape[0] != count or weight_jet.shape[0] != count:
        raise ValueError("input and weight jets must have the specified coefficient count")
    if weight_jet.ndim != 3 or input_jet.shape[-1] != weight_jet.shape[-1]:
        raise ValueError("weight jet must be (M, out, in) with matching input width")
    sample = jnp.tensordot(input_jet[0], weight_jet[0], axes=([-1], [-1]))
    rows = [jnp.zeros_like(sample) for _ in range(count)]
    for out, a, b in multiply_table(dim, order):
        rows[out] = rows[out] + jnp.tensordot(input_jet[a], weight_jet[b], axes=([-1], [-1]))
    if bias_jet is not None:
        if bias_jet.shape != (count, weight_jet.shape[1]):
            raise ValueError("bias jet must be (M, out)")
        rows = [row + bias_jet[i] for i, row in enumerate(rows)]
    return jnp.stack(rows, axis=0)

FUNCTION

evaluate

omnibias.jax.realization.evaluateSource ↗

evaluate(x: Array, theta: Array, spec: RealizationSpec) -> Array

Evaluate the exact described architecture with live parameters.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`RealizationSpec`Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When theta.ndim != 1 or theta.shape[0] != spec.layout.size, raises ValueError('theta does not match realization layout').

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

Read the implementation

def evaluate(x: Array, theta: Array, spec: RealizationSpec) -> Array:
    """Evaluate the exact described architecture with live parameters."""
    if theta.ndim != 1 or theta.shape[0] != spec.layout.size:
        raise ValueError("theta does not match realization layout")
    needed = spec.layout.size + math.prod(x.shape[:-1]) * (spec.input_dim + max(layer.out_features for layer in spec.layers))
    return realization_jet(x[None, ...], theta[None, ...], spec, dim=1, order=0,
                           max_coefficients=needed)[0]

FUNCTION

observe

omnibias.jax.realization.observeSource ↗

observe(theta: Array, spec: RealizationSpec, observation: ObservationSpec, *, points: Array | None=None, weights: Array | None=None, residual: Callable[[Array, Callable[[Array], Array]], Array] | None=None, coefficient_map: Callable[[Array], Array] | None=None) -> Array

Observe values/derivatives/quadrature, or an explicitly supplied operator.

These are observations only. A declared domain or coefficient-map label is not independently checked as a global equality proof by this numerical API.

Parameters

Name	Type & default	Contract
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`RealizationSpec`Required	See the operation contract and implementation below.
`observation`positional-or-keyword	`ObservationSpec`Required	See the operation contract and implementation below.
`points`keyword-only	`Array \| None`Default: None	See the operation contract and implementation below.
`weights`keyword-only	`Array \| None`Default: None	See the operation contract and implementation below.
`residual`keyword-only	`Callable[[Array, Callable[[Array], Array]], Array] \| None`Default: None	See the operation contract and implementation below.
`coefficient_map`keyword-only	`Callable[[Array], Array] \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When coefficient_map is None, raises ValueError('coefficient observations require an explicit coefficient map').
When coords.ndim != 2 or coords.shape[1] != spec.input_dim, raises ValueError('observation points must be (N, input_dim)').
When residual is None, raises ValueError('residual observations require an explicit residual operator').
When not indices or any((len(a) != spec.input_dim for a in indices)), raises ValueError('derivative observations require input multi-indices').
When not any((layer.op == 'integral' for layer in spec.layers)), raises ValueError('activation_window observations require an integral-role realization').
When weights is not None or observation.weights, raises ValueError('activation windows do not apply quadrature weights').
When w.shape != (coords.shape[0],), raises ValueError('quadrature observations require one weight per point').

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

Read the implementation

def observe(
    theta: Array, spec: RealizationSpec, observation: ObservationSpec,
    *, points: Array | None = None, weights: Array | None = None,
    residual: Callable[[Array, Callable[[Array], Array]], Array] | None = None,
    coefficient_map: Callable[[Array], Array] | None = None,
) -> Array:
    """Observe values/derivatives/quadrature, or an explicitly supplied operator.

    These are observations only. A declared domain or coefficient-map label is
    not independently checked as a global equality proof by this numerical API.
    """
    if observation.kind == "coefficients":
        if coefficient_map is None:
            raise ValueError("coefficient observations require an explicit coefficient map")
        return coefficient_map(theta).reshape(-1)
    coords = jnp.asarray(observation.points, dtype=theta.dtype) if points is None else points
    if coords.ndim != 2 or coords.shape[1] != spec.input_dim:
        raise ValueError("observation points must be (N, input_dim)")
    if observation.kind == "residuals":
        if residual is None:
            raise ValueError("residual observations require an explicit residual operator")
        return residual(coords, lambda x: evaluate(x, theta, spec)).reshape(-1)
    if observation.kind == "derivatives":
        indices = observation.derivative_indices
        if not indices or any(len(a) != spec.input_dim for a in indices):
            raise ValueError("derivative observations require input multi-indices")
        order = max(map(sum, indices))
        directions = jnp.zeros((theta.shape[0], spec.input_dim), dtype=theta.dtype)
        ibasis = jnp.eye(spec.input_dim, dtype=theta.dtype)
        jet = parameter_jet(coords, theta, spec, directions, order, input_basis=ibasis)
        positions = index_position(spec.input_dim, order)
        return jnp.stack([jet[positions[a]] * math.prod(math.factorial(i) for i in a) for a in indices], axis=0).reshape(-1)
    values = evaluate(coords, theta, spec)
    if observation.kind == "integrals":
        if observation.integral_kind == "activation_window":
            if not any(layer.op == "integral" for layer in spec.layers):
                raise ValueError("activation_window observations require an integral-role realization")
            if weights is not None or observation.weights:
                raise ValueError("activation windows do not apply quadrature weights")
            return values.reshape(-1)
        # Domain/measure observations are the declared weighted numerical rule.
        w = jnp.asarray(observation.weights, dtype=theta.dtype) if weights is None else weights
        if w.shape != (coords.shape[0],):
            raise ValueError("quadrature observations require one weight per point")
        return (values * w[:, None]).sum(axis=0).reshape(-1)
    return values.reshape(-1)

FUNCTION

operator_jet

omnibias.jax.realization.operator_jetSource ↗

operator_jet(jet: Array, layer: LayerSpec, *, dim: int, order: int, windows: tuple[Array, Array] | None=None) -> Array

Apply one static activation-level role, including live window endpoints.

Parameters

Name	Type & default	Contract
`jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`layer`positional-or-keyword	`LayerSpec`Required	See the operation contract and implementation below.
`dim`keyword-only	`int`Required	See the operation contract and implementation below.
`order`keyword-only	`int`Required	See the operation contract and implementation below.
`windows`keyword-only	`tuple[Array, Array] \| None`Default: None	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When windows is None, raises ValueError('window roles require endpoint jets').

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

Read the implementation

def operator_jet(
    jet: Array, layer: LayerSpec, *, dim: int, order: int,
    windows: tuple[Array, Array] | None = None,
) -> Array:
    """Apply one static activation-level role, including live window endpoints."""
    name = layer.activation
    if name is None:
        return jet
    if layer.op not in {"band", "integral"}:
        return _activation_jet(jet, name, dim, order, layer.activation_order)
    if windows is None:
        raise ValueError("window roles require endpoint jets")
    lo, hi = windows
    if layer.window_parameterization == "center_softplus_width":
        width = _activation_jet(hi, "softplus", dim, order)
        lo, hi = lo - 0.5 * width, lo + 0.5 * width
    lower, upper = _shift(jet, lo), _shift(jet, hi)
    if layer.op == "band":
        return _activation_jet(upper, name, dim, order) - _activation_jet(lower, name, dim, order)
    return _antiderivative_jet(upper, name, dim, order) - _antiderivative_jet(lower, name, dim, order)

FUNCTION

operator_parameter_jet

omnibias.jax.realization.operator_parameter_jetSource ↗

operator_parameter_jet(jet: Array, biases: Array, signs: Array, activation: str | ActivationSpec[Array], op: str, *, dim: int, order: int, derivative_order: int=0, normalize_integral: bool=False, integral_small_width: float=0.0, use_small_width_taylor: bool=False) -> Array

All six OperatorBlock roles in the original bias/sign coordinates.

Bias/sign jets have shape (M, channels, K). Integral biases store center/raw-softplus-width. The optional midpoint branch reproduces the existing block's explicitly approximate small-width policy, including its branch derivatives; zero threshold selects the antiderivative window.

Parameters

Name	Type & default	Contract
`jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`biases`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`signs`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`activation`positional-or-keyword	`str \| ActivationSpec[Array]`Required	See the operation contract and implementation below.
`op`positional-or-keyword	`str`Required	See the operation contract and implementation below.
`dim`keyword-only	`int`Required	See the operation contract and implementation below.
`order`keyword-only	`int`Required	See the operation contract and implementation below.
`derivative_order`keyword-only	`int`Default: 0	See the operation contract and implementation below.
`normalize_integral`keyword-only	`bool`Default: False	See the operation contract and implementation below.
`integral_small_width`keyword-only	`float`Default: 0.0	See the operation contract and implementation below.
`use_small_width_taylor`keyword-only	`bool`Default: False	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When biases.ndim != 3 or biases.shape[0] != count or signs.shape != biases.shape, raises ValueError('bias/sign jets must have matching (M, channels, K) shapes').
When jet.shape[0] != count or jet.shape[-1] != biases.shape[1], raises ValueError('input jet coefficient count/channels must match bias jets').
When offset < 1, raises ValueError('derivative_order must be positive').
When op not in {'identity', 'band', 'integral'}, raises ValueError(f'unknown operator role {op!r}').
When arity != (1 if op == 'identity' else 2), raises ValueError('operator bias arity does not match the role').

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

Read the implementation

def operator_parameter_jet(
    jet: Array, biases: Array, signs: Array, activation: str | ActivationSpec[Array],
    op: str, *, dim: int, order: int, derivative_order: int = 0,
    normalize_integral: bool = False, integral_small_width: float = 0.0,
    use_small_width_taylor: bool = False,
) -> Array:
    """All six OperatorBlock roles in the original bias/sign coordinates.

    Bias/sign jets have shape ``(M, channels, K)``. Integral biases store
    center/raw-softplus-width. The optional midpoint branch reproduces the
    existing block's explicitly approximate small-width policy, including its
    branch derivatives; zero threshold selects the antiderivative window.
    """
    count = num_multi_indices(dim, order)
    if biases.ndim != 3 or biases.shape[0] != count or signs.shape != biases.shape:
        raise ValueError("bias/sign jets must have matching (M, channels, K) shapes")
    if jet.shape[0] != count or jet.shape[-1] != biases.shape[1]:
        raise ValueError("input jet coefficient count/channels must match bias jets")
    if op in {"grad", "laplacian", "derivative"}:
        offset = {"grad": 1, "laplacian": 2, "derivative": derivative_order}[op]
        if offset < 1:
            raise ValueError("derivative_order must be positive")
        return _activation_jet(_shift(jet, biases.mean(axis=-1)), activation, dim, order, offset)
    if op not in {"identity", "band", "integral"}:
        raise ValueError(f"unknown operator role {op!r}")
    arity = biases.shape[-1]
    if arity != (1 if op == "identity" else 2):
        raise ValueError("operator bias arity does not match the role")
    endpoints = biases
    width: Array | None = None
    if op == "integral":
        center = biases[..., 0]
        width = _activation_jet(biases[..., 1], "softplus", dim, order)
        endpoints = jnp.stack((center - 0.5 * width, center + 0.5 * width), axis=-1)
    pieces = []
    for k in range(arity):
        shifted = _shift(jet, endpoints[..., k])
        value = (_antiderivative_jet(shifted, activation, dim, order) if op == "integral"
                 else _activation_jet(shifted, activation, dim, order))
        pieces.append(_multiply(value, signs[..., k], dim, order))
    result = pieces[0]
    for piece in pieces[1:]:
        result = result + piece
    if width is not None:
        if normalize_integral:
            # Match the existing forward denominator floor without evaluating 1/0.
            floor = _seed(jnp.full_like(width[0], 1e-12), None, dim, order)
            denominator = jnp.where(width[0] > 1e-12, width, floor)
            reciprocal = jnp.stack([((-1) ** n) * math.factorial(n) / denominator[0] ** (n + 1)
                                      for n in range(order + 1)], axis=0)
            result = _multiply(result, compose_jet_mv(denominator, reciprocal, dim, order), dim, order)
        if use_small_width_taylor and integral_small_width > 0:
            midpoint = _activation_jet(_shift(jet, biases[..., 0]), activation, dim, order)
            approx = midpoint if normalize_integral else _multiply(midpoint, width, dim, order)
            result = jnp.where(width[0] < integral_small_width, approx, result)
    return result

FUNCTION

parameter_jet

omnibias.jax.realization.parameter_jetSource ↗

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

Live Taylor jet along explicit P-by-q parameter directions.

input_basis is x.shape + (q,) (a (D,q) basis also broadcasts across input batches). Use zero parameter directions for pure input jets. The basis is explicit so a full parameter tensor is never allocated implicitly.

Parameters

Name	Type & default	Contract
`x`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`theta`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`RealizationSpec`Required	See the operation contract and implementation below.
`basis`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`order`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`input_basis`keyword-only	`Array \| None`Default: None	See the operation contract and implementation below.
`max_coefficients`keyword-only	`int`Default: 1000000	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When theta.ndim != 1 or theta.shape[0] != spec.layout.size, raises ValueError('theta must be the flat parameter vector described by layout').
When basis.ndim != 2 or basis.shape[0] != theta.shape[0] or basis.shape[1] < 1, raises ValueError('basis must have shape (P, q), q >= 1').
When count * (theta.size + x.size) > max_coefficients, raises ValueError('joint-jet coefficient budget exceeded before seeding').

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

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:
    """Live Taylor jet along explicit P-by-q parameter directions.

    ``input_basis`` is ``x.shape + (q,)`` (a ``(D,q)`` basis also broadcasts
    across input batches). Use zero parameter directions for pure input jets.
    The basis is explicit so a full parameter tensor is never allocated implicitly.
    """
    if theta.ndim != 1 or theta.shape[0] != spec.layout.size:
        raise ValueError("theta must be the flat parameter vector described by layout")
    if basis.ndim != 2 or basis.shape[0] != theta.shape[0] or basis.shape[1] < 1:
        raise ValueError("basis must have shape (P, q), q >= 1")
    dim = basis.shape[1]
    count = num_multi_indices(dim, order)
    if count * (theta.size + x.size) > max_coefficients:
        raise ValueError("joint-jet coefficient budget exceeded before seeding")
    if input_basis is not None:
        input_basis = jnp.broadcast_to(input_basis, x.shape + (dim,))
    return realization_jet(_seed(x, input_basis, dim, order), _seed(theta, basis, dim, order),
                           spec, dim=dim, order=order, max_coefficients=max_coefficients)

FUNCTION

realization_jet

omnibias.jax.realization.realization_jetSource ↗

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

Propagate joint input/parameter jets through a described dense network.

Parameters

Name	Type & default	Contract
`input_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`parameter_jet`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`spec`positional-or-keyword	`RealizationSpec`Required	See the operation contract and implementation below.
`dim`keyword-only	`int`Required	See the operation contract and implementation below.
`order`keyword-only	`int`Required	See the operation contract and implementation below.
`max_coefficients`keyword-only	`int`Default: 1000000	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When parameter_jet.shape != (count, spec.layout.size), raises ValueError('parameter jet does not match the layout/coefficient count').
When input_jet.shape[0] != count or input_jet.shape[-1] != spec.input_dim, raises ValueError('input jet does not match the realization').
When max_coefficients < 1 or needed > max_coefficients, raises ValueError(f'joint jet needs {needed} coefficients; budget is {max_coefficients}').

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

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 realization_jet(
    input_jet: Array, parameter_jet: Array, spec: RealizationSpec,
    *, dim: int, order: int, max_coefficients: int = 1_000_000,
) -> Array:
    """Propagate joint input/parameter jets through a described dense network."""
    count = num_multi_indices(dim, order)
    if parameter_jet.shape != (count, spec.layout.size):
        raise ValueError("parameter jet does not match the layout/coefficient count")
    if input_jet.shape[0] != count or input_jet.shape[-1] != spec.input_dim:
        raise ValueError("input jet does not match the realization")
    batch = math.prod(input_jet.shape[1:-1])
    peak = max(layer.out_features for layer in spec.layers)
    needed = count * (spec.layout.size + batch * (spec.input_dim + peak))
    if max_coefficients < 1 or needed > max_coefficients:
        raise ValueError(f"joint jet needs {needed} coefficients; budget is {max_coefficients}")
    jet = input_jet
    for i, layer in enumerate(spec.layers):
        weight = _block(parameter_jet, spec.layout, layer.weight_name or f"W{i}")
        bias = None if layer.bias_name is None else _block(parameter_jet, spec.layout, layer.bias_name)
        jet = affine_joint_jet_mv(jet, weight, bias, dim=dim, order=order)
        windows = None
        if layer.window_names is not None:
            windows = (_block(parameter_jet, spec.layout, layer.window_names[0]),
                       _block(parameter_jet, spec.layout, layer.window_names[1]))
        jet = operator_jet(jet, layer, dim=dim, order=order, windows=windows)
    return jet

← Package overview Runnable learning examples →

omnibias.jax.realization

Realization​

Parameters

Fields & methods

Parameters

Parameters

Parameters

affine_joint_jet​

Parameters

affine_joint_jet_mv​

Parameters

evaluate​

Parameters

observe​

Parameters

operator_jet​

Parameters

operator_parameter_jet​

Parameters

parameter_jet​

Parameters

realization_jet​

Parameters

Realization

affine_joint_jet

affine_joint_jet_mv

evaluate

observe

operator_jet

operator_parameter_jet

parameter_jet

realization_jet