Skip to main content
API REFERENCE/omnibias-jax

omnibias.jax.laplacian

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.laplacian
Module contract & mathematical conventions

Closed-form Laplacian for an omnibias-style one-layer scalar field.

The model we accelerate is the elementary multi-bias layer:

f(x) = b + sum_h c_h * sigma(W_h . x + beta_h), x in R^D

Its Laplacian on R^D is

nabla_x^2 f(x) = sum_h c_h * sigma''(W_h . x + beta_h) * ||W_h||^2.

Computing this with standard JAX would call jax.hessian or jax.jacfwd(jax.jacrev(f)) on the field, which builds a full D x D Hessian matrix and traces it. That is O(B * D * H) per backward pass and produces O(B * D^2) intermediate memory.

The omnibias kernel here is O(B * H) and O(1) in D for the *Laplacian* (the ||W_h||^2 term is a per-row scalar reduction, computed once per parameter update).

This module is the JAX-side workhorse a FermiNet/DeepQMC-style code would call to assemble the kinetic-energy term of a neural wavefunction without nesting two jax.jacrev calls.

The neural_field_value_and_laplacian function additionally returns the gradient, which most VMC loops need anyway.

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

deep_field_laplacian​

omnibias.jax.laplacian.deep_field_laplacianSource ↗
deep_field_laplacian(x: Array, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool=False) -> Array

nabla_x^2 f(x) for a deep MLP, no ceiling at any D (Tier A).

Thin wrapper over deep_field_value_grad_laplacian discarding the value and gradient. See that function for the recursion and the depth-1 collapse onto neural_field_laplacian.

Compute the Laplacian directly when your residual needs the Hessian trace rather than all Hessian entries. For a batch shaped (B, D), the supported deep-MLP path returns one Laplacian per output component, shaped (B, C). Dimension still affects arithmetic and storage even though no dense D-by-D Hessian is materialized by this contraction. The example compares the result with a dense Hessian trace and checks a parameter gradient.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredInput batch, shape (B, D), with a floating-point dtype compatible with the layer tensors.
layerspositional-or-keywordSequence[LayerSpec]RequiredOrdered (W, b, activation) triples. The first W accepts D features; the final width C is the number of output components. None activates an affine readout.
allow_almost_everywherekeyword-onlyboolDefault: FalseOpt in to documented almost-everywhere activation conventions. Such a derivative is not a classical derivative at a nonsmooth knot and omits distributional terms.
ReturnsArray
Complete runnable example · deep laplacian jax

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

Download complete script ↓
"""JAX Laplacian contraction against an independent Hessian trace."""
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import numpy as np
from omnibias.jax.laplacian import deep_field_laplacian

weight = jnp.array([[0.2, -0.3], [0.5, 0.1]])
bias = jnp.array([0.1, -0.1])
readout = jnp.array([[0.7, -0.4]])
point = jnp.array([[0.3, -0.2]])
def lap(W):
    return deep_field_laplacian(point, [(W, bias, "tanh"), (readout, None, None)])
def field(x):
    return (readout @ jnp.tanh(weight @ x + bias)).squeeze()
reference = jnp.trace(jax.hessian(field)(point[0]))
np.testing.assert_allclose(lap(weight).squeeze(), reference, rtol=1e-10, atol=1e-12)
gradient = jax.grad(lambda W: jnp.mean(lap(W)**2))(weight)
assert gradient.shape == weight.shape and bool(jnp.isfinite(gradient).all())
print("JAX Laplacian and parameter gradient verified.")
python deep_laplacian_jax.py
Executed regression · arctan k1 laplacian still works

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

Full file, imports & fixtures ↗
def test_arctan_k1_laplacian_still_works() -> None:
    """k=1 only ever needs fastpath order 2, which arctan does support."""
    layers = _make_layers(3, 4, depth=2, activation="arctan", seed=27)
    x = jnp.asarray(np.random.default_rng(28).normal(size=(3,)))
    lap = deep_field_laplacian(x, layers)
    assert bool(jnp.all(jnp.isfinite(lap)))
python -m pytest "packages/omnibias-jax/tests/test_deep_laplacian.py::test_arctan_k1_laplacian_still_works" -q
Read the implementation
def deep_field_laplacian(
    x: Array, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool = False
) -> Array:
    r"""``nabla_x^2 f(x)`` for a deep MLP, no ceiling at any ``D`` (Tier A).

    Thin wrapper over :func:`deep_field_value_grad_laplacian` discarding the
    value and gradient. See that function for the recursion and the
    depth-1 collapse onto :func:`neural_field_laplacian`.
    """
    _, _, lap = deep_field_value_grad_laplacian(
        x, layers, allow_almost_everywhere=allow_almost_everywhere
    )
    return lap
Related APIsomnibias.jax.laplacian.deep_field_value_grad_laplacian↗
FUNCTION

deep_field_polylaplacian​

omnibias.jax.laplacian.deep_field_polylaplacianSource ↗
deep_field_polylaplacian(x: Array, layers: Sequence[LayerSpec], k: int, *, mode: str='auto', n_directions: int=256, budget: int | None=None, seed: int=0, key: Array | None=None) -> Array

Delta^k f(x) for a deep MLP at any input dimension D.

k=1 always routes to deep_field_laplacian (Tier A: exact, no ceiling). For k >= 2, mode="auto" (the default) picks Tier B ("support", exact) while omnibias.core.contraction.select_mode finds the number of support sets within budget (omnibias.core.contraction.DEFAULT_SUPPORT_BUDGET when budget is None), else Tier C ("estimator", exact in expectation over n_directions sampled unit directions).

Parameters

x : Array of shape (D,) or (B, D). layers : sequence of (W, b, spec), as for deep_field_laplacian. k : int >= 1 Poly-Laplacian order. mode : {"auto", "support", "estimator"} Force a tier for k >= 2 instead of the budget-driven choice. n_directions : int Sample count for Tier C (ignored otherwise). budget : int, optional Override omnibias.core.contraction.DEFAULT_SUPPORT_BUDGET for the Tier B/C decision. seed : int PRNG seed for Tier C direction sampling when key is None. key : Array, optional JAX PRNG key for Tier C; when supplied, seed is ignored and the caller can vary the key per step without forcing a jit retrace on seed.

Returns

Array of shape (C,) or (B, C).

Notes

Tier B is exact to floating-point round-off; Tier C is exact only *in expectation*. Use deep_field_polylaplacian_with_report for the resolved tier, sample standard error, and a probabilistic enclosure. Never claim "O(1) at arbitrary order" for Tier C: its cost is O(n_directions) regardless of D, but it is a Monte Carlo estimator, not a closed form.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredArray of shape (D,) or (B, D).
layerspositional-or-keywordSequence[LayerSpec]Requiredsequence of (W, b, spec), as for deep_field_laplacian.
kpositional-or-keywordintRequiredint >= 1 Poly-Laplacian order.
modekeyword-onlystrDefault: 'auto'{"auto", "support", "estimator"} Force a tier for k >= 2 instead of the budget-driven choice.
n_directionskeyword-onlyintDefault: 256int Sample count for Tier C (ignored otherwise).
budgetkeyword-onlyint | NoneDefault: Noneint, optional Override omnibias.core.contraction.DEFAULT_SUPPORT_BUDGET for the Tier B/C decision.
seedkeyword-onlyintDefault: 0int PRNG seed for Tier C direction sampling when key is None.
keykeyword-onlyArray | NoneDefault: NoneArray, optional JAX PRNG key for Tier C; when supplied, seed is ignored and the caller can vary the key per step without forcing a jit retrace on seed.
ReturnsArray
Executed regression · polylaplacian k must be positive

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_polylaplacian_k_must_be_positive() -> None:
    layers = _make_layers(3, 4, depth=1, activation="tanh", seed=21)
    x = jnp.asarray(np.random.default_rng(22).normal(size=(3,)))
    with pytest.raises(ValueError, match="polylaplacian order k must be >= 1"):
        deep_field_polylaplacian(x, layers, 0)
python -m pytest "packages/omnibias-jax/tests/test_deep_laplacian.py::test_polylaplacian_k_must_be_positive" -q
Regression example · arctan order cap raises for k2 estimator

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

Full file, imports & fixtures ↗
def test_arctan_order_cap_raises_for_k2_estimator() -> None:
    layers = _make_layers(3, 4, depth=1, activation="arctan", seed=25)
    x = jnp.asarray(np.random.default_rng(26).normal(size=(3,)))
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")
python -m pytest packages/omnibias-jax/tests/test_deep_laplacian.py -k test_arctan_order_cap_raises_for_k2_estimator -q
Read the implementation
def deep_field_polylaplacian(
    x: Array,
    layers: Sequence[LayerSpec],
    k: int,
    *,
    mode: str = "auto",
    n_directions: int = 256,
    budget: int | None = None,
    seed: int = 0,
    key: Array | None = None,
) -> Array:
    r"""``Delta^k f(x)`` for a deep MLP at any input dimension ``D``.

    ``k=1`` always routes to :func:`deep_field_laplacian` (Tier A: exact, no
    ceiling). For ``k >= 2``, ``mode="auto"`` (the default) picks Tier B
    (``"support"``, exact) while :func:`omnibias.core.contraction.select_mode`
    finds the number of support sets within ``budget``
    (:data:`omnibias.core.contraction.DEFAULT_SUPPORT_BUDGET` when
    ``budget`` is ``None``), else Tier C (``"estimator"``, exact in
    expectation over ``n_directions`` sampled unit directions).

    Parameters
    ----------
    x : Array of shape ``(D,)`` or ``(B, D)``.
    layers : sequence of ``(W, b, spec)``, as for :func:`deep_field_laplacian`.
    k : int >= 1
        Poly-Laplacian order.
    mode : {"auto", "support", "estimator"}
        Force a tier for ``k >= 2`` instead of the budget-driven choice.
    n_directions : int
        Sample count for Tier C (ignored otherwise).
    budget : int, optional
        Override :data:`omnibias.core.contraction.DEFAULT_SUPPORT_BUDGET`
        for the Tier B/C decision.
    seed : int
        PRNG seed for Tier C direction sampling when ``key`` is ``None``.
    key : Array, optional
        JAX PRNG key for Tier C; when supplied, ``seed`` is ignored and the
        caller can vary the key per step without forcing a ``jit`` retrace on
        ``seed``.

    Returns
    -------
    Array of shape ``(C,)`` or ``(B, C)``.

    Notes
    -----
    Tier B is exact to floating-point round-off; Tier C is exact only *in
    expectation*. Use :func:`deep_field_polylaplacian_with_report` for the
    resolved tier, sample standard error, and a probabilistic enclosure.
    Never claim "O(1) at arbitrary order" for Tier C: its cost is
    ``O(n_directions)`` regardless of ``D``, but it is a Monte Carlo
    estimator, not a closed form.
    """
    out, _ = deep_field_polylaplacian_with_report(
        x,
        layers,
        k,
        mode=mode,
        n_directions=n_directions,
        budget=budget,
        seed=seed,
        key=key,
    )
    return out
Related APIsomnibias.jax.laplacian.deep_field_polylaplacian_with_report↗
FUNCTION

deep_field_polylaplacian_with_report​

omnibias.jax.laplacian.deep_field_polylaplacian_with_reportSource ↗
deep_field_polylaplacian_with_report(x: Array, layers: Sequence[LayerSpec], k: int, *, mode: str='auto', n_directions: int=256, budget: int | None=None, seed: int=0, key: Array | None=None) -> tuple[Array, PolylaplacianReport]

Delta^k f(x) plus metadata about which tier ran.

Returns the same point estimate as deep_field_polylaplacian and a ~omnibias.core.contraction.PolylaplacianReport carrying the resolved mode, D, k, the budget, and for Tier C the direction count, sample standard error, and a ~omnibias.core.verified.sampled.ConcentrationReport from ~omnibias.core.verified.sampled.hoeffding_enclosure.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
layerspositional-or-keywordSequence[LayerSpec]RequiredSee the operation contract and implementation below.
kpositional-or-keywordintRequiredSee the operation contract and implementation below.
modekeyword-onlystrDefault: 'auto'See the operation contract and implementation below.
n_directionskeyword-onlyintDefault: 256See the operation contract and implementation below.
budgetkeyword-onlyint | NoneDefault: NoneSee the operation contract and implementation below.
seedkeyword-onlyintDefault: 0See the operation contract and implementation below.
keykeyword-onlyArray | NoneDefault: NoneSee the operation contract and implementation below.
Returnstuple[Array, PolylaplacianReport]
Explicit validation & failure cases
  • When k < 1, raises ValueError(f'polylaplacian order k must be >= 1, got {k}').
  • When resolved_mode == 'forward', raises ValueError("mode='forward' only applies to k=1 (the Laplacian); use 'support' or 'estimator' for k >= 2").

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

Executed regression · polylaplacian k must be positive

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_polylaplacian_k_must_be_positive() -> None:
    layers = _make_layers(3, 4, depth=1, activation="tanh", seed=21)
    x = jnp.asarray(np.random.default_rng(22).normal(size=(3,)))
    with pytest.raises(ValueError, match="polylaplacian order k must be >= 1"):
        deep_field_polylaplacian(x, layers, 0)
python -m pytest "packages/omnibias-jax/tests/test_deep_laplacian.py::test_polylaplacian_k_must_be_positive" -q
Read the implementation
def deep_field_polylaplacian_with_report(
    x: Array,
    layers: Sequence[LayerSpec],
    k: int,
    *,
    mode: str = "auto",
    n_directions: int = 256,
    budget: int | None = None,
    seed: int = 0,
    key: Array | None = None,
) -> tuple[Array, PolylaplacianReport]:
    r"""``Delta^k f(x)`` plus metadata about which tier ran.

    Returns the same point estimate as :func:`deep_field_polylaplacian` and a
    :class:`~omnibias.core.contraction.PolylaplacianReport` carrying the
    resolved mode, ``D``, ``k``, the budget, and for Tier C the direction
    count, sample standard error, and a :class:`~omnibias.core.verified.sampled.ConcentrationReport`
    from :func:`~omnibias.core.verified.sampled.hoeffding_enclosure`.
    """
    if k < 1:
        raise ValueError(f"polylaplacian order k must be >= 1, got {k}")
    x = jnp.asarray(x)
    batched = x.ndim > 1
    x2 = x if batched else x[None, :]
    dim = x2.shape[-1]
    resolved_mode, resolved_budget = _resolve_polylaplacian_mode(dim, k, mode, budget)
    concentration = None
    if k == 1:
        out = deep_field_laplacian(x2, layers)
    elif resolved_mode == "forward":
        raise ValueError(
            "mode='forward' only applies to k=1 (the Laplacian); use "
            "'support' or 'estimator' for k >= 2"
        )
    elif resolved_mode == "support":
        out = _polylaplacian_support(x2, layers, k)
    elif resolved_mode == "estimator":
        out, scaled = _polylaplacian_estimator(
            x2,
            layers,
            k,
            n_directions=n_directions,
            seed=seed,
            key=key,
            return_samples=True,
        )
        concentration = _tier_c_concentration(scaled)
    else:
        raise ValueError(f"unknown mode {resolved_mode!r}")
    report = PolylaplacianReport(
        mode=resolved_mode,
        dim=dim,
        k=k,
        budget=resolved_budget,
        n_directions=n_directions if resolved_mode == "estimator" else None,
        concentration=concentration,
    )
    return (out if batched else out[0]), report
Related APIsomnibias.core.contraction.PolylaplacianReport↗omnibias.jax.laplacian.deep_field_laplacian↗
FUNCTION

deep_field_value_grad_laplacian​

omnibias.jax.laplacian.deep_field_value_grad_laplacianSource ↗
deep_field_value_grad_laplacian(x: Array, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool=False) -> tuple[Array, Array, Array]

(f(x), grad_x f(x), nabla_x^2 f(x)) for a deep MLP, no ceiling.

Forward-Laplacian recursion (Tier A): carries the Jacobian J of the running activation w.r.t. the *original* input and the scalar Laplacian L of each unit through every layer, using the closed-form sigma'/sigma'' fastpath (no autodiff through the activation). Memory is O(B * H * D) -- the same order as the forward pass -- with no comb(D + 2k, D) combinatorial term at any D.

Parameters

x : Array of shape (D,) or (B, D). layers : sequence of (W, b, spec) As consumed by omnibias.jax.jet.mlp_jet / omnibias.jax.jet_mv.mlp_jet_mv. spec=None is a pure affine layer (used for the readout); every other layer computes sigma(W a + b).

Returns

value : Array of shape (C,) or (B, C) grad : Array of shape (D, C) or (B, D, C) laplacian : Array of shape (C,) or (B, C)

At depth=1 (one hidden layer plus an affine readout) this reproduces neural_field_value_grad_laplacian to floating-point round-off -- the single-layer field is the D=1 collapse of the deep recursion.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredArray of shape (D,) or (B, D).
layerspositional-or-keywordSequence[LayerSpec]Requiredsequence of (W, b, spec) As consumed by omnibias.jax.jet.mlp_jet / omnibias.jax.jet_mv.mlp_jet_mv. spec=None is a pure affine layer (used for the readout); every other layer computes sigma(W a + b).
allow_almost_everywherekeyword-onlyboolDefault: FalseSee the operation contract and implementation below.
Returnstuple[Array, Array, Array]
Explicit validation & failure cases
  • When not layers, raises ValueError('layers must be non-empty').
  • When resolved.fastpath is None, raises ValueError(f'activation {resolved.name!r} has no fast-path kernel, cannot use the deep-field Laplacian').

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

Executed regression · arctan k1 laplacian still works

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

Full file, imports & fixtures ↗
def test_arctan_k1_laplacian_still_works() -> None:
    """k=1 only ever needs fastpath order 2, which arctan does support."""
    layers = _make_layers(3, 4, depth=2, activation="arctan", seed=27)
    x = jnp.asarray(np.random.default_rng(28).normal(size=(3,)))
    lap = deep_field_laplacian(x, layers)
    assert bool(jnp.all(jnp.isfinite(lap)))
python -m pytest "packages/omnibias-jax/tests/test_deep_laplacian.py::test_arctan_k1_laplacian_still_works" -q
Regression example · value grad laplacian matches the split calls

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_laplacian_matches_the_split_calls() -> None:
    layers = _make_layers(5, 6, depth=2, activation="tanh", seed=6)
    x = jnp.asarray(np.random.default_rng(7).normal(size=(3, 5)))
    value, grad, lap = deep_field_value_grad_laplacian(x, layers)
    lap_only = deep_field_laplacian(x, layers)
    assert jnp.allclose(lap, lap_only)
    assert value.shape == (3, 1)
    assert grad.shape == (3, 5, 1)
python -m pytest packages/omnibias-jax/tests/test_deep_laplacian.py -k test_value_grad_laplacian_matches_the_split_calls -q
Read the implementation
def deep_field_value_grad_laplacian(
    x: Array,
    layers: Sequence[LayerSpec],
    *,
    allow_almost_everywhere: bool = False,
) -> tuple[Array, Array, Array]:
    r"""``(f(x), grad_x f(x), nabla_x^2 f(x))`` for a deep MLP, no ceiling.

    Forward-Laplacian recursion (Tier A): carries the Jacobian ``J`` of the
    running activation w.r.t. the *original* input and the scalar Laplacian
    ``L`` of each unit through every layer, using the closed-form
    ``sigma'``/``sigma''`` fastpath (no autodiff through the activation).
    Memory is ``O(B * H * D)`` -- the same order as the forward pass -- with
    no ``comb(D + 2k, D)`` combinatorial term at any ``D``.

    Parameters
    ----------
    x : Array of shape ``(D,)`` or ``(B, D)``.
    layers : sequence of ``(W, b, spec)``
        As consumed by :func:`omnibias.jax.jet.mlp_jet` / :func:`omnibias.jax.jet_mv.mlp_jet_mv`.
        ``spec=None`` is a pure affine layer (used for the readout); every
        other layer computes ``sigma(W a + b)``.

    Returns
    -------
    value : Array of shape ``(C,)`` or ``(B, C)``
    grad : Array of shape ``(D, C)`` or ``(B, D, C)``
    laplacian : Array of shape ``(C,)`` or ``(B, C)``

    At ``depth=1`` (one hidden layer plus an affine readout) this reproduces
    :func:`neural_field_value_grad_laplacian` to floating-point round-off --
    the single-layer field is the ``D=1`` collapse of the deep recursion.
    """
    if not layers:
        raise ValueError("layers must be non-empty")
    x = jnp.asarray(x)
    batched = x.ndim > 1
    a = x if batched else x[None, :]
    B, dim = a.shape
    L = jnp.zeros_like(a)  # (B, D); Laplacian of the identity map is 0.
    J: Array | None = None  # sentinel for the identity Jacobian (never materialised)
    for W, b, spec in layers:
        W = jnp.asarray(W)
        u = jnp.matmul(a, W.T)
        if b is not None:
            b_arr = jnp.asarray(b)
            u = u + (b_arr if b_arr.ndim == 1 else b_arr)
        first_layer = J is None
        if first_layer:
            # First layer: (identity Jacobian) @ W.T collapses to W.T broadcast
            # over the batch, so the O(B*D^2) identity is never formed.
            J_u = jnp.broadcast_to(W.T[None, :, :], (B, dim, W.shape[0]))
        else:
            assert J is not None
            J_u = jnp.tensordot(J, W, axes=([-1], [-1]))  # (B, D, H_out)
        L_u = jnp.matmul(L, W.T)  # (B, H_out)
        if spec is None:
            a, J, L = u, J_u, L_u
            continue
        resolved = get_activation(spec)
        if resolved.fastpath is None:
            raise ValueError(
                f"activation {resolved.name!r} has no fast-path kernel, cannot "
                "use the deep-field Laplacian"
            )
        _check_piecewise_second_order(
            resolved.name, allow_almost_everywhere=allow_almost_everywhere
        )
        try:
            sigma_p = resolved.fastpath(u, 1)
            sigma_pp = resolved.fastpath(u, 2)
        except NotImplementedError as exc:
            raise ValueError(
                f"activation {resolved.name!r} fastpath does not support order "
                "2, required for the deep-field Laplacian"
            ) from exc
        a = resolved.forward(u)
        J = sigma_p[:, None, :] * J_u
        if first_layer:
            row_norm_sq = jnp.sum(W * W, axis=-1)  # (H,) == ||W_h||^2, batch-free
            lap_contrib = sigma_pp * row_norm_sq
        else:
            lap_contrib = sigma_pp * jnp.sum(J_u * J_u, axis=1)
        L = lap_contrib + sigma_p * L_u
    assert J is not None  # every loop iteration above assigns J; layers is non-empty
    value = a if batched else a[0]
    grad = J if batched else J[0]
    lap = L if batched else L[0]
    return value, grad, lap
Related APIsomnibias.jax.get_activation↗
FUNCTION

neural_field_hessian​

omnibias.jax.neural_field_hessianSource ↗
neural_field_hessian(x: Array, W: Array, beta: Array, c: Array, activation: str | JaxActivationSpec) -> Array

H_x f -- the full D x D Hessian of the one-layer field.

Returns

H : Array of shape (..., D, D) The dense Hessian. For a batched x of shape (B, D) the output has shape (B, D, D).

Notes

b does not enter the Hessian, so it is not a parameter here. The Hessian is *symmetric* by construction (it is a sum of symmetric rank-1 outer products sigma''(z_h) * W_h W_h^T).

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel, cannot use closed-form Hessian').

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

Executed regression · neural field hessian matches jax hessian

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_neural_field_hessian_matches_jax_hessian(name: str) -> None:
    rng = np.random.default_rng(6)
    D, H = 4, 6
    W = jnp.asarray(rng.normal(scale=0.3, size=(H, D)))
    beta = jnp.asarray(rng.normal(scale=0.2, size=(H,)))
    c = jnp.asarray(rng.normal(scale=0.4, size=(H,)))
    x = jnp.asarray(rng.normal(scale=0.7, size=(D,)))

    def f(x):
        return neural_field_value(x, W, beta, c, 0.0, name)

    hess_closed = neural_field_hessian(x, W, beta, c, name)
    hess_autograd = jax.hessian(f)(x)
    np.testing.assert_allclose(
        hess_closed, hess_autograd, rtol=1e-9, atol=1e-12
    )
    np.testing.assert_allclose(
        hess_closed, hess_closed.T, rtol=0.0, atol=1e-13
    )
python -m pytest "packages/omnibias-jax/tests/test_audit_regressions.py::test_neural_field_hessian_matches_jax_hessian[exp]" -q
Regression example · hessian trace equals laplacian

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_trace_equals_laplacian(activation: str, D: int) -> None:
    """``trace(H_x f) == nabla_x^2 f``. The two omnibias primitives
    must be consistent with each other."""
    W, beta, c, b, x = _rand_params(
        D,
        _H,
        seed=hash((activation, D, "trace")) & 0xFFFF,
    )

    H = neural_field_hessian(x, W, beta, c, activation)
    lap = neural_field_laplacian(x, W, beta, c, activation)

    np.testing.assert_allclose(
        float(jnp.trace(H)),
        float(lap),
        rtol=_RTOL,
        atol=_ATOL,
        err_msg=(
            f"trace(H) != Laplacian for {activation!r} at D={D}: "
            f"trace={float(jnp.trace(H))!r}, Laplacian={float(lap)!r}"
        ),
    )
python -m pytest packages/omnibias-jax/tests/test_laplacian_hessian.py -k test_hessian_trace_equals_laplacian -q
Read the implementation
def neural_field_hessian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    activation: str | JaxActivationSpec,
) -> Array:
    """``H_x f`` -- the full ``D x D`` Hessian of the one-layer field.

    Returns
    -------
    H : Array of shape ``(..., D, D)``
        The dense Hessian. For a batched ``x`` of shape ``(B, D)`` the
        output has shape ``(B, D, D)``.

    Notes
    -----
    ``b`` does not enter the Hessian, so it is not a parameter here.
    The Hessian is *symmetric* by construction (it is a sum of
    symmetric rank-1 outer products ``sigma''(z_h) * W_h W_h^T``).
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(
            f"activation {spec.name!r} has no fast-path kernel, cannot use closed-form Hessian"
        )
    z = jnp.matmul(x, W.T) + beta  # (..., H)
    sigma_pp = spec.fastpath(z, 2)  # (..., H)
    weights = sigma_pp * c  # (..., H)
    # H = W^T diag(weights) W  =  einsum over the hidden dimension.
    return jnp.einsum("...h,hi,hj->...ij", weights, W, W)
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_hessian
  • omnibias.jax.laplacian.neural_field_hessian
FUNCTION

neural_field_laplacian​

omnibias.jax.neural_field_laplacianSource ↗
neural_field_laplacian(x: Array, W: Array, beta: Array, c: Array, activation: str | JaxActivationSpec) -> Array

nabla_x^2 f(x) in closed form.

b does not enter the Laplacian, so it is not a parameter here.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel, cannot use closed-form Laplacian').

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

Executed regression · polylaplacian k eq 1 matches laplacian

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_polylaplacian_k_eq_1_matches_laplacian(activation: str) -> None:
    """For k=1, polylaplacian must equal ``neural_field_laplacian``."""
    W, beta, c, _, x = _rand_params(6, _H, seed=hash(activation) & 0xFFFF)
    lap_k1 = neural_field_polylaplacian(x, W, beta, c, activation, k=1)
    lap_ref = neural_field_laplacian(x, W, beta, c, activation)
    np.testing.assert_allclose(
        np.asarray(lap_k1),
        np.asarray(lap_ref),
        rtol=1e-14,
        atol=1e-14,
    )
python -m pytest "packages/omnibias-jax/tests/test_polylaplacian.py::test_polylaplacian_k_eq_1_matches_laplacian[exp]" -q
Regression example · neural field laplacian jit

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_neural_field_laplacian_jit(name: str) -> None:
    rng = np.random.default_rng(0)
    D, H = 3, 5
    W = jnp.asarray(rng.normal(size=(H, D)))
    beta = jnp.asarray(rng.normal(size=(H,)))
    c = jnp.asarray(rng.normal(size=(H,)))
    x = jnp.asarray(rng.normal(size=(D,)))

    f_jit = jax.jit(neural_field_laplacian, static_argnums=(4,))
    val_jit = f_jit(x, W, beta, c, name)
    val_eager = neural_field_laplacian(x, W, beta, c, name)
    np.testing.assert_allclose(val_jit, val_eager, rtol=0.0, atol=1e-12)
python -m pytest packages/omnibias-jax/tests/test_audit_regressions.py -k test_neural_field_laplacian_jit -q
Read the implementation
def neural_field_laplacian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    activation: str | JaxActivationSpec,
) -> Array:
    """``nabla_x^2 f(x)`` in closed form.

    ``b`` does not enter the Laplacian, so it is not a parameter here.
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(
            f"activation {spec.name!r} has no fast-path kernel, cannot use closed-form Laplacian"
        )
    z = jnp.matmul(x, W.T) + beta  # (B, H) or (H,)
    sigma_pp = spec.fastpath(z, 2)  # (B, H) or (H,)
    row_norm_sq = (W * W).sum(axis=-1)  # (H,)
    return (sigma_pp * (c * row_norm_sq)) @ jnp.ones_like(c)
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_laplacian
  • omnibias.jax.laplacian.neural_field_laplacian
FUNCTION

neural_field_local_p4_over_psi​

omnibias.jax.laplacian.neural_field_local_p4_over_psiSource ↗
neural_field_local_p4_over_psi(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec) -> Array

Compute the local p^4 operator ratio Delta^2 psi / psi for a wavefunction represented directly by the supported one-layer field.

This is a ratio of the biharmonic field value to the field value itself. It is singular at a node where psi=0, and near-zero denominators can magnify numerical error. The field projection and weight-norm computation depend on dimension; avoiding a dense derivative tensor does not make the complete computation dimension independent.

The corresponding mass-velocity Hamiltonian term additionally requires its physical prefactor and a consistent unit convention. A model representing log|psi| needs the appropriate chain rule; do not pass a log-amplitude field here and interpret its ratio as the wavefunction's local operator.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
ReturnsArray
Read the implementation
def neural_field_local_p4_over_psi(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array | float,
    activation: str | JaxActivationSpec,
) -> Array:
    """Relativistic mass-velocity local-energy operator p^4 psi / psi.

    For a wavefunction ``psi(x) = f(x)`` represented directly by the
    one-layer field (no log-magnitude transform), the mass-velocity
    local energy is

        L_rel(x) = p^4 psi(x) / psi(x) = Delta^2 f / f.

    This is the operator that enters the relativistic Foldy-Wouthuysen
    Hamiltonian as

        H_rel = - p^4 / (8 m^3 c^2)

    (in CGS-Gaussian / Hartree atomic units, where m = 1 and the
    speed of light c ~ 137.036).  The cost of this estimator on a
    K=H-neuron omnibias field is O(B * H), independent of D --
    compare to folx-nested which would be O(B * D * H).

    For wavefunctions parameterised as ``psi = exp(L)`` with L a
    neural log-magnitude (the standard FermiNet pattern), see
    ``neural_field_p4_chain_rule_log`` (roadmap to-do; the chain
    rule has six terms in 1D and twenty terms in 3D).
    """
    f, p4_psi = neural_field_value_and_polylaplacian(x, W, beta, c, b, activation, k=2)
    return p4_psi / f
Related APIsomnibias.jax.laplacian.neural_field_value_and_polylaplacian↗
FUNCTION

neural_field_polylaplacian​

omnibias.jax.laplacian.neural_field_polylaplacianSource ↗
neural_field_polylaplacian(x: Array, W: Array, beta: Array, c: Array, activation: str | JaxActivationSpec, k: int) -> Array

Delta^k f(x) = (nabla^2)^k f(x) in closed form.

Computes the k-th iterated Laplacian of the one-layer field f(x) = b + sum_h c_h sigma(W_h . x + beta_h). b cancels for k >= 1, so it is not a parameter here.

Parameters

x : (..., D) W : (H, D) beta : (H,) c : (H,) activation : str | JaxActivationSpec Must be a Riccati-class activation with fastpath(z, 2k) defined (i.e., sigmoid, tanh, softplus, gaussian, exp). k : int >= 1 Polylaplacian order. k=1 reproduces neural_field_laplacian; k=2 is the relativistic mass-velocity correction operator \hat p^4 = (\nabla^2)^2; k=3 enters the third-order Foldy-Wouthuysen expansion.

Returns

Array of shape (...,) -- the k-th iterated Laplacian.

Notes

Memory cost: O(B * H), *independent of k*. Time cost: one forward pass through sigma^{(2k)} plus one (B,H) reduction. Compare to folx-nested Delta^k which is O(B * D^{k-1} * H).

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequired(..., D)
Wpositional-or-keywordArrayRequired(H, D)
betapositional-or-keywordArrayRequired(H,)
cpositional-or-keywordArrayRequired(H,)
activationpositional-or-keywordstr | JaxActivationSpecRequiredstr | JaxActivationSpec Must be a Riccati-class activation with fastpath(z, 2k) defined (i.e., sigmoid, tanh, softplus, gaussian, exp).
kpositional-or-keywordintRequiredint >= 1 Polylaplacian order. k=1 reproduces neural_field_laplacian; k=2 is the relativistic mass-velocity correction operator \hat p^4 = (\nabla^2)^2; k=3 enters the third-order Foldy-Wouthuysen expansion.
ReturnsArray
Explicit validation & failure cases
  • When k < 1, raises ValueError(f'polylaplacian order k must be >= 1, got {k}').
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel, cannot use closed-form polylaplacian').

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

Executed regression · neural field polylaplacian vmap

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_neural_field_polylaplacian_vmap(name: str) -> None:
    rng = np.random.default_rng(4)
    B, D, H = 6, 3, 5
    W = jnp.asarray(rng.normal(size=(H, D)))
    beta = jnp.asarray(rng.normal(size=(H,)))
    c = jnp.asarray(rng.normal(size=(H,)))
    xs = jnp.asarray(rng.normal(size=(B, D)))

    f_v = jax.vmap(
        lambda x: neural_field_polylaplacian(x, W, beta, c, name, 2),
        in_axes=0,
    )
    out = f_v(xs)
    assert out.shape == (B,)
python -m pytest "packages/omnibias-jax/tests/test_audit_regressions.py::test_neural_field_polylaplacian_vmap[exp]" -q
Regression example · polylaplacian k eq 1 matches laplacian

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_polylaplacian_k_eq_1_matches_laplacian(activation: str) -> None:
    """For k=1, polylaplacian must equal ``neural_field_laplacian``."""
    W, beta, c, _, x = _rand_params(6, _H, seed=hash(activation) & 0xFFFF)
    lap_k1 = neural_field_polylaplacian(x, W, beta, c, activation, k=1)
    lap_ref = neural_field_laplacian(x, W, beta, c, activation)
    np.testing.assert_allclose(
        np.asarray(lap_k1),
        np.asarray(lap_ref),
        rtol=1e-14,
        atol=1e-14,
    )
python -m pytest packages/omnibias-jax/tests/test_polylaplacian.py -k test_polylaplacian_k_eq_1_matches_laplacian -q
Read the implementation
def neural_field_polylaplacian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    activation: str | JaxActivationSpec,
    k: int,
) -> Array:
    """``Delta^k f(x) = (nabla^2)^k f(x)`` in closed form.

    Computes the k-th iterated Laplacian of the one-layer field
    ``f(x) = b + sum_h c_h sigma(W_h . x + beta_h)``.  ``b`` cancels
    for k >= 1, so it is not a parameter here.

    Parameters
    ----------
    x : (..., D)
    W : (H, D)
    beta : (H,)
    c : (H,)
    activation : str | JaxActivationSpec
        Must be a Riccati-class activation with ``fastpath(z, 2k)``
        defined (i.e., ``sigmoid``, ``tanh``, ``softplus``, ``gaussian``,
        ``exp``).
    k : int >= 1
        Polylaplacian order.  k=1 reproduces ``neural_field_laplacian``;
        k=2 is the relativistic mass-velocity correction operator
        :math:`\\hat p^4 = (\\nabla^2)^2`; k=3 enters the third-order
        Foldy-Wouthuysen expansion.

    Returns
    -------
    Array of shape ``(...,)`` -- the k-th iterated Laplacian.

    Notes
    -----
    Memory cost: ``O(B * H)``, *independent of k*.  Time cost: one
    forward pass through ``sigma^{(2k)}`` plus one ``(B,H)`` reduction.
    Compare to folx-nested ``Delta^k`` which is ``O(B * D^{k-1} * H)``.
    """
    if k < 1:
        raise ValueError(f"polylaplacian order k must be >= 1, got {k}")
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(
            f"activation {spec.name!r} has no fast-path kernel, "
            "cannot use closed-form polylaplacian"
        )
    z = jnp.matmul(x, W.T) + beta  # (..., H)
    sigma_2k = spec.fastpath(z, 2 * k)  # (..., H)
    row_norm_sq = (W * W).sum(axis=-1)  # (H,)
    row_norm_2k = row_norm_sq**k  # (H,)
    return sigma_2k @ (c * row_norm_2k)
Related APIsomnibias.jax.get_activation↗
FUNCTION

neural_field_value​

omnibias.jax.neural_field_valueSource ↗
neural_field_value(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec) -> Array

f(x) = b + sum_h c_h sigma(W_h . x + beta_h).

Batched in x if it has shape (B, D); returns (B,). For a single point (D,) it returns a scalar.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · neural field laplacian matches jax hessian

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_neural_field_laplacian_matches_jax_hessian(name: str) -> None:
    rng = np.random.default_rng(5)
    D, H = 4, 6
    W = jnp.asarray(rng.normal(scale=0.3, size=(H, D)))
    beta = jnp.asarray(rng.normal(scale=0.2, size=(H,)))
    c = jnp.asarray(rng.normal(scale=0.4, size=(H,)))
    x = jnp.asarray(rng.normal(scale=0.7, size=(D,)))

    def f(x):
        return neural_field_value(x, W, beta, c, 0.0, name)

    lap_closed = neural_field_laplacian(x, W, beta, c, name)
    lap_autograd = jnp.trace(jax.hessian(f)(x))
    np.testing.assert_allclose(
        lap_closed, lap_autograd, rtol=1e-9, atol=1e-12
    )
python -m pytest "packages/omnibias-jax/tests/test_audit_regressions.py::test_neural_field_laplacian_matches_jax_hessian[exp]" -q
Regression example · neural field grad matches jax grad

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_neural_field_grad_matches_jax_grad(name: str) -> None:
    rng = np.random.default_rng(7)
    D, H = 4, 6
    W = jnp.asarray(rng.normal(scale=0.3, size=(H, D)))
    beta = jnp.asarray(rng.normal(scale=0.2, size=(H,)))
    c = jnp.asarray(rng.normal(scale=0.4, size=(H,)))
    x = jnp.asarray(rng.normal(scale=0.7, size=(D,)))

    def f(x):
        return neural_field_value(x, W, beta, c, 0.0, name)

    _, grad_closed, _ = neural_field_value_grad_laplacian(x, W, beta, c, 0.0, name)
    grad_autograd = jax.grad(f)(x)
    np.testing.assert_allclose(
        grad_closed, grad_autograd, rtol=1e-10, atol=1e-13
    )
python -m pytest packages/omnibias-jax/tests/test_audit_regressions.py -k test_neural_field_grad_matches_jax_grad -q
Read the implementation
def neural_field_value(
    x: Array,  # (B, D) or (D,)
    W: Array,  # (H, D)
    beta: Array,  # (H,)
    c: Array,  # (H,)
    b: Array | float,  # scalar
    activation: str | JaxActivationSpec,
) -> Array:
    """``f(x) = b + sum_h c_h sigma(W_h . x + beta_h)``.

    Batched in ``x`` if it has shape ``(B, D)``; returns ``(B,)``. For
    a single point ``(D,)`` it returns a scalar.
    """
    spec = get_activation(activation)
    z = jnp.matmul(x, W.T) + beta  # (B, H) or (H,)
    return b + spec.forward(z) @ c  # (B,) or scalar
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_value
  • omnibias.jax.laplacian.neural_field_value
FUNCTION

neural_field_value_and_laplacian​

omnibias.jax.neural_field_value_and_laplacianSource ↗
neural_field_value_and_laplacian(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec) -> tuple[Array, Array]

Returns (f(x), nabla_x^2 f(x)) in one forward pass.

Reuses the pre-activation z between the value and the Laplacian, so this is strictly cheaper than calling neural_field_value and neural_field_laplacian separately.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
Returnstuple[Array, Array]
Explicit validation & failure cases
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel').

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

Read the implementation
def neural_field_value_and_laplacian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array | float,
    activation: str | JaxActivationSpec,
) -> tuple[Array, Array]:
    """Returns ``(f(x), nabla_x^2 f(x))`` in one forward pass.

    Reuses the pre-activation ``z`` between the value and the
    Laplacian, so this is strictly cheaper than calling
    :func:`neural_field_value` and :func:`neural_field_laplacian`
    separately.
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(f"activation {spec.name!r} has no fast-path kernel")
    z = jnp.matmul(x, W.T) + beta
    sigma_z = spec.forward(z)
    sigma_pp = spec.fastpath(z, 2)
    f = b + sigma_z @ c
    row_norm_sq = (W * W).sum(axis=-1)
    lap = sigma_pp @ (c * row_norm_sq)
    return f, lap
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_value_and_laplacian
  • omnibias.jax.laplacian.neural_field_value_and_laplacian
FUNCTION

neural_field_value_and_polylaplacian​

omnibias.jax.laplacian.neural_field_value_and_polylaplacianSource ↗
neural_field_value_and_polylaplacian(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec, k: int) -> tuple[Array, Array]

Returns (f(x), Delta^k f(x)) in one fused pass.

Useful for VMC local-energy estimators that need both the wavefunction value (for log|psi| and importance sampling) and Delta^k psi / psi (for the local kinetic-energy term).

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
kpositional-or-keywordintRequiredSee the operation contract and implementation below.
Returnstuple[Array, Array]
Explicit validation & failure cases
  • When k < 1, raises ValueError(f'polylaplacian order k must be >= 1, got {k}').
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel').

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

Executed regression · value and polylaplacian matches separate

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

Full file, imports & fixtures ↗
def test_value_and_polylaplacian_matches_separate(activation: str) -> None:
    """Fused (value, Delta^k) call equals separate calls."""
    W, beta, c, b, x = _rand_params(6, _H, seed=hash(activation) & 0xFFFF)
    for k in (1, 2, 3):
        f_fused, pk_fused = neural_field_value_and_polylaplacian(x, W, beta, c, b, activation, k=k)
        f_sep = neural_field_value(x, W, beta, c, b, activation)
        pk_sep = neural_field_polylaplacian(x, W, beta, c, activation, k=k)
        np.testing.assert_allclose(np.asarray(f_fused), np.asarray(f_sep), rtol=1e-14, atol=1e-14)
        np.testing.assert_allclose(np.asarray(pk_fused), np.asarray(pk_sep), rtol=1e-14, atol=1e-14)
python -m pytest "packages/omnibias-jax/tests/test_polylaplacian.py::test_value_and_polylaplacian_matches_separate[exp]" -q
Read the implementation
def neural_field_value_and_polylaplacian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array | float,
    activation: str | JaxActivationSpec,
    k: int,
) -> tuple[Array, Array]:
    """Returns ``(f(x), Delta^k f(x))`` in one fused pass.

    Useful for VMC local-energy estimators that need both the
    wavefunction value (for ``log|psi|`` and importance sampling)
    and ``Delta^k psi / psi`` (for the local kinetic-energy term).
    """
    if k < 1:
        raise ValueError(f"polylaplacian order k must be >= 1, got {k}")
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(f"activation {spec.name!r} has no fast-path kernel")
    z = jnp.matmul(x, W.T) + beta
    sigma_z = spec.forward(z)
    sigma_2k = spec.fastpath(z, 2 * k)
    f = b + sigma_z @ c
    row_norm_sq = (W * W).sum(axis=-1)
    row_norm_2k = row_norm_sq**k
    poly_lap = sigma_2k @ (c * row_norm_2k)
    return f, poly_lap
Related APIsomnibias.jax.get_activation↗
FUNCTION

neural_field_value_grad_hessian​

omnibias.jax.neural_field_value_grad_hessianSource ↗
neural_field_value_grad_hessian(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec) -> tuple[Array, Array, Array]

Returns (f(x), grad_x f(x), H_x f(x)) in one fused pass.

Closed form for all three. Reuses the pre-activation z so the activation derivative tower is touched once per order rather than once per output. The Hessian is the building block for FermiNet Tier 2/3 integration -- it composes through the chain rule for backflow and equivariant-layer coordinate maps.

Parameters

x : (..., D) input point(s) W : (H, D) hidden-layer weights beta : (H,) hidden-layer biases c : (H,) output-layer weights b : scalar output-layer bias activation : str | JaxActivationSpec Must be one of the omnibias fast-path (Riccati-class) activations: tanh, sigmoid, softplus, gaussian, exp.

Returns

f : (..., ) scalar field value grad : (..., D) gradient hessian : (..., D, D) Symmetric Hessian matrix.

See Also

neural_field_value_grad_laplacian : same fused call but returning the *trace* of the Hessian (the Laplacian), for the single- electron Slater trace(A^{-1} L) path.

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequired(..., D) input point(s)
Wpositional-or-keywordArrayRequired(H, D) hidden-layer weights
betapositional-or-keywordArrayRequired(H,) hidden-layer biases
cpositional-or-keywordArrayRequired(H,) output-layer weights
bpositional-or-keywordArray | floatRequiredscalar output-layer bias
activationpositional-or-keywordstr | JaxActivationSpecRequiredstr | JaxActivationSpec Must be one of the omnibias fast-path (Riccati-class) activations: tanh, sigmoid, softplus, gaussian, exp.
Returnstuple[Array, Array, Array]
Explicit validation & failure cases
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel, cannot use closed-form Hessian').

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

Executed regression · neural field value grad hessian vmap

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_neural_field_value_grad_hessian_vmap(name: str) -> None:
    rng = np.random.default_rng(3)
    B, D, H = 7, 3, 5
    W = jnp.asarray(rng.normal(size=(H, D)))
    beta = jnp.asarray(rng.normal(size=(H,)))
    c = jnp.asarray(rng.normal(size=(H,)))
    xs = jnp.asarray(rng.normal(size=(B, D)))
    b = 0.0

    f_v = jax.vmap(
        lambda x: neural_field_value_grad_hessian(x, W, beta, c, b, name),
        in_axes=0,
    )
    fs, gs, hs = f_v(xs)
    assert fs.shape == (B,)
    assert gs.shape == (B, D)
    assert hs.shape == (B, D, D)
python -m pytest "packages/omnibias-jax/tests/test_audit_regressions.py::test_neural_field_value_grad_hessian_vmap[exp]" -q
Regression example · neural field hessian rejects non fastpath activation

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_neural_field_hessian_rejects_non_fastpath_activation() -> None:
    """A base without a fast-path kernel must raise.

    (``relu`` now carries an all-orders a.e. tower, so we use a deliberately
    fastpath-less spec to exercise the guard.)
    """
    from omnibias.core.spec import ActivationSpec

    no_fastpath = ActivationSpec(name="_no_fastpath_probe", forward=lambda z: z, fastpath=None)
    W, beta, c, b, x = _rand_params(3, 4, seed=0)
    with pytest.raises((ValueError, NotImplementedError), match=r"(fast-path|distributional)"):
        neural_field_hessian(x, W, beta, c, no_fastpath)
    with pytest.raises((ValueError, NotImplementedError), match=r"(fast-path|distributional)"):
        neural_field_value_grad_hessian(x, W, beta, c, b, no_fastpath)
python -m pytest packages/omnibias-jax/tests/test_laplacian_hessian.py -k test_neural_field_hessian_rejects_non_fastpath_activation -q
Read the implementation
def neural_field_value_grad_hessian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array | float,
    activation: str | JaxActivationSpec,
) -> tuple[Array, Array, Array]:
    """Returns ``(f(x), grad_x f(x), H_x f(x))`` in one fused pass.

    Closed form for all three. Reuses the pre-activation ``z`` so the
    activation derivative tower is touched once per order rather than
    once per output. The Hessian is the building block for
    FermiNet Tier 2/3 integration -- it composes through the
    chain rule for backflow and equivariant-layer coordinate maps.

    Parameters
    ----------
    x : (..., D)        input point(s)
    W : (H, D)          hidden-layer weights
    beta : (H,)         hidden-layer biases
    c : (H,)            output-layer weights
    b : scalar          output-layer bias
    activation : str | JaxActivationSpec
        Must be one of the omnibias fast-path (Riccati-class)
        activations: ``tanh``, ``sigmoid``, ``softplus``, ``gaussian``,
        ``exp``.

    Returns
    -------
    f : (..., )         scalar field value
    grad : (..., D)     gradient
    hessian : (..., D, D)
        Symmetric Hessian matrix.

    See Also
    --------
    neural_field_value_grad_laplacian : same fused call but returning
        the *trace* of the Hessian (the Laplacian), for the single-
        electron Slater ``trace(A^{-1} L)`` path.
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(
            f"activation {spec.name!r} has no fast-path kernel, cannot use closed-form Hessian"
        )
    z = jnp.matmul(x, W.T) + beta  # (..., H)
    sigma_z = spec.forward(z)  # (..., H)
    sigma_p = spec.fastpath(z, 1)  # (..., H)
    sigma_pp = spec.fastpath(z, 2)  # (..., H)

    f = b + sigma_z @ c  # (..., )
    grad = (sigma_p * c) @ W  # (..., D)
    weights = sigma_pp * c  # (..., H)
    hessian = jnp.einsum("...h,hi,hj->...ij", weights, W, W)
    return f, grad, hessian
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_value_grad_hessian
  • omnibias.jax.laplacian.neural_field_value_grad_hessian
FUNCTION

neural_field_value_grad_laplacian​

omnibias.jax.neural_field_value_grad_laplacianSource ↗
neural_field_value_grad_laplacian(x: Array, W: Array, beta: Array, c: Array, b: Array | float, activation: str | JaxActivationSpec) -> tuple[Array, Array, Array]

Returns (f(x), grad_x f(x), nabla_x^2 f(x)).

Closed form for all three. The gradient is needed by every VMC local-energy estimator (it shows up in the |grad log psi|^2 term).

Parameters

NameType & defaultContract
xpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
Wpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
betapositional-or-keywordArrayRequiredSee the operation contract and implementation below.
cpositional-or-keywordArrayRequiredSee the operation contract and implementation below.
bpositional-or-keywordArray | floatRequiredSee the operation contract and implementation below.
activationpositional-or-keywordstr | JaxActivationSpecRequiredSee the operation contract and implementation below.
Returnstuple[Array, Array, Array]
Explicit validation & failure cases
  • When spec.fastpath is None, raises ValueError(f'activation {spec.name!r} has no fast-path kernel').

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

Executed regression · neural field grad matches jax grad

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_neural_field_grad_matches_jax_grad(name: str) -> None:
    rng = np.random.default_rng(7)
    D, H = 4, 6
    W = jnp.asarray(rng.normal(scale=0.3, size=(H, D)))
    beta = jnp.asarray(rng.normal(scale=0.2, size=(H,)))
    c = jnp.asarray(rng.normal(scale=0.4, size=(H,)))
    x = jnp.asarray(rng.normal(scale=0.7, size=(D,)))

    def f(x):
        return neural_field_value(x, W, beta, c, 0.0, name)

    _, grad_closed, _ = neural_field_value_grad_laplacian(x, W, beta, c, 0.0, name)
    grad_autograd = jax.grad(f)(x)
    np.testing.assert_allclose(
        grad_closed, grad_autograd, rtol=1e-10, atol=1e-13
    )
python -m pytest "packages/omnibias-jax/tests/test_audit_regressions.py::test_neural_field_grad_matches_jax_grad[exp]" -q
Read the implementation
def neural_field_value_grad_laplacian(
    x: Array,
    W: Array,
    beta: Array,
    c: Array,
    b: Array | float,
    activation: str | JaxActivationSpec,
) -> tuple[Array, Array, Array]:
    """Returns ``(f(x), grad_x f(x), nabla_x^2 f(x))``.

    Closed form for all three. The gradient is needed by every VMC
    local-energy estimator (it shows up in the ``|grad log psi|^2``
    term).
    """
    spec = get_activation(activation)
    if spec.fastpath is None:
        raise ValueError(f"activation {spec.name!r} has no fast-path kernel")
    z = jnp.matmul(x, W.T) + beta  # (..., H)
    sigma_z = spec.forward(z)  # (..., H)
    sigma_p = spec.fastpath(z, 1)  # (..., H)
    sigma_pp = spec.fastpath(z, 2)  # (..., H)

    f = b + sigma_z @ c  # (..., )
    grad = (sigma_p * c) @ W  # (..., D)
    row_norm_sq = (W * W).sum(axis=-1)  # (H,)
    lap = sigma_pp @ (c * row_norm_sq)  # (..., )
    return f, grad, lap
Related APIsomnibias.jax.get_activation↗
Other public import paths
  • omnibias.jax.neural_field_value_grad_laplacian
  • omnibias.jax.laplacian.neural_field_value_grad_laplacian
FUNCTION

restrict_first_layer​

omnibias.jax.laplacian.restrict_first_layerSource ↗
restrict_first_layer(layers: Sequence[LayerSpec], x0: Array, support: tuple[int, ...]) -> list[LayerSpec]

Restrict layers to the input axes in support, folding the rest into the bias.

If x = x0 + E_support @ y (y living on the support axes only, everything else pinned at x0), then the first affine layer z = W x + b becomes z = W[:, support] @ y + (W @ x0 + b) -- exact, and the local weight matrix has only len(support) columns. Used by Tier B's per-support local jet and, in omnibias.pinn.jax.fields.jet_mlp, to isolate the spatial axes of a field that also carries a non-spatial (e.g. time) coordinate.

Parameters

NameType & defaultContract
layerspositional-or-keywordSequence[LayerSpec]RequiredSee the operation contract and implementation below.
x0positional-or-keywordArrayRequiredSee the operation contract and implementation below.
supportpositional-or-keywordtuple[int, ...]RequiredSee the operation contract and implementation below.
Returnslist[LayerSpec]
Executed regression · arctan order cap raises for k2 support

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

Full file, imports & fixtures ↗
def test_arctan_order_cap_raises_for_k2_support() -> None:
    """arctan's fastpath caps at order 2; Delta^2 needs order 4 and must raise,
    not silently mis-compute."""
    layers = _make_layers(3, 4, depth=1, activation="arctan", seed=23)
    x = jnp.asarray(np.random.default_rng(24).normal(size=(3,)))
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="support")
python -m pytest "packages/omnibias-jax/tests/test_deep_laplacian.py::test_arctan_order_cap_raises_for_k2_support" -q
Regression example · restrict first layer is exact

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_restrict_first_layer_is_exact() -> None:
    dim = 6
    layers = _make_layers(dim, 5, depth=2, activation="tanh", seed=10)
    x0 = jnp.asarray(np.random.default_rng(11).normal(size=(dim,)))
    support = (1, 3, 4)
    local = restrict_first_layer(layers, x0, support)
    y = jnp.asarray(np.random.default_rng(12).normal(size=(len(support),)) * 0.1)
    # x = x0 + E_support @ y
    x_full = x0.at[jnp.asarray(support)].add(y)

    def value(layer_list, xi):
        a = xi
        for W, b, spec in layer_list:
            u = a @ W.T + b
            if spec is None:
                a = u
            else:
                from omnibias.jax.activations import get_activation

                a = get_activation(spec).forward(u)
        return a

    got = value(local, y)
    want = value(layers, x_full)
    assert jnp.allclose(got, want, atol=1e-12)
python -m pytest packages/omnibias-jax/tests/test_deep_laplacian.py -k test_restrict_first_layer_is_exact -q
Read the implementation
def restrict_first_layer(
    layers: Sequence[LayerSpec], x0: Array, support: tuple[int, ...]
) -> list[LayerSpec]:
    """Restrict ``layers`` to the input axes in ``support``, folding the rest into the bias.

    If ``x = x0 + E_support @ y`` (``y`` living on the ``support`` axes only,
    everything else pinned at ``x0``), then the first affine layer
    ``z = W x + b`` becomes ``z = W[:, support] @ y + (W @ x0 + b)`` -- exact,
    and the local weight matrix has only ``len(support)`` columns. Used by
    Tier B's per-support local jet and, in
    :mod:`omnibias.pinn.jax.fields.jet_mlp`, to isolate the spatial axes of a
    field that also carries a non-spatial (e.g. time) coordinate.
    """
    W0, b0, spec0 = layers[0]
    W0 = jnp.asarray(W0)
    cols = jnp.asarray(support)
    W_local = jnp.take(W0, cols, axis=1)
    shift = jnp.matmul(W0, jnp.asarray(x0))
    b_local = shift if b0 is None else shift + jnp.asarray(b0)
    return [(W_local, b_local, spec0), *layers[1:]]
← Package overviewRunnable learning examples →