API REFERENCE/omnibias-torch

omnibias.torch.laplacian

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.torch.laplacian

Module contract & mathematical conventions

Closed-form deep-network Laplacian fast lane (torch).

Bit-identical twin of omnibias.jax.laplacian; see that module's deep_field_* section for the full derivation and honesty notes. This module additionally carries every neural_field_* one-layer kernel (torch previously had no laplacian module at all -- this closes that asymmetry) plus the deep-network family:

• deep_field_laplacian / deep_field_value_grad_laplacian (Tier A) -- the forward-Laplacian recursion, exact and ceiling-free at any input dimension D.
• deep_field_polylaplacian (Tiers B/C) -- exact via a support-grouped local multivariate jet while a budget is respected, else an unbiased sphere-average estimator via the directional jet.

The one-layer kernels below (neural_field_*) are ports of the JAX originals, useful on their own and as the depth-1 reference the deep kernels are pinned against in the test suite.

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.torch.laplacian.deep_field_laplacianSource ↗

deep_field_laplacian(x: Tensor, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool=False) -> Tensor

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

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	Input batch, shape (B, D), with a floating-point dtype compatible with the layer tensors.
layerspositional-or-keyword	Sequence[LayerSpec]Required	Ordered (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-only	boolDefault: False	Opt in to documented almost-everywhere activation conventions. Such a derivative is not a classical derivative at a nonsmooth knot and omits distributional terms.

ReturnsTensor

Complete runnable example · deep laplacian

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

Download complete script ↓

"""A direct deep Laplacian, independently checked against a dense Hessian."""
import torch
from omnibias.torch.laplacian import deep_field_laplacian

weight = torch.tensor([[0.2, -0.3], [0.5, 0.1]], dtype=torch.float64,
                      requires_grad=True)
bias = torch.tensor([0.1, -0.1], dtype=torch.float64)
readout = torch.tensor([[0.7, -0.4]], dtype=torch.float64)
layers = [(weight, bias, "tanh"), (readout, None, None)]
point = torch.tensor([[0.3, -0.2]], dtype=torch.float64)
laplacian = deep_field_laplacian(point, layers)

def field(x):
    return (readout @ torch.tanh(weight @ x + bias)).squeeze()

reference = torch.func.hessian(field)(point[0]).trace()
assert laplacian.shape == (1, 1)
torch.testing.assert_close(laplacian.squeeze(), reference, rtol=1e-10, atol=1e-12)
laplacian.square().mean().backward()
assert weight.grad is not None and torch.isfinite(weight.grad).all()
print("Laplacian agrees with Hessian trace; parameter gradients retained.")

python deep_laplacian.py

Executed regression · relu laplacian opt in

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_relu_laplacian_opt_in() -> None:
    layers = _make_layers(3, 4, depth=1, activation="relu", seed=31)
    x = torch.as_tensor(np.random.default_rng(32).normal(size=(3,)), dtype=torch.float64)
    lap = deep_field_laplacian(x, layers, allow_almost_everywhere=True)
    assert bool(torch.all(torch.isfinite(lap)))

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_relu_laplacian_opt_in" -q

Read the implementation

def deep_field_laplacian(
    x: Tensor, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool = False
) -> Tensor:
    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.torch.laplacian.deep_field_value_grad_laplacian↗

FUNCTION

deep_field_polylaplacian

omnibias.torch.laplacian.deep_field_polylaplacianSource ↗

deep_field_polylaplacian(x: Tensor, layers: Sequence[LayerSpec], k: int, *, mode: str='auto', n_directions: int=256, budget: int | None=None, seed: int=0) -> Tensor

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

See deep_field_polylaplacian_with_report and the JAX twin for the full parameter documentation and honesty notes.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
layerspositional-or-keyword	Sequence[LayerSpec]Required	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.
modekeyword-only	strDefault: 'auto'	See the operation contract and implementation below.
n_directionskeyword-only	intDefault: 256	See the operation contract and implementation below.
budgetkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
seedkeyword-only	intDefault: 0	See the operation contract and implementation below.

ReturnsTensor

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 = torch.as_tensor(np.random.default_rng(22).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="polylaplacian order k must be >= 1"):
        deep_field_polylaplacian(x, layers, 0)

python -m pytest "packages/omnibias-torch/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 = torch.as_tensor(np.random.default_rng(26).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="estimator")

python -m pytest packages/omnibias-torch/tests/test_deep_laplacian.py -k test_arctan_order_cap_raises_for_k2_estimator -q

Read the implementation

def deep_field_polylaplacian(
    x: Tensor,
    layers: Sequence[LayerSpec],
    k: int,
    *,
    mode: str = "auto",
    n_directions: int = 256,
    budget: int | None = None,
    seed: int = 0,
) -> Tensor:
    r"""``Delta^k f(x)`` for a deep MLP at any input dimension ``D``.

    See :func:`deep_field_polylaplacian_with_report` and the JAX twin for the
    full parameter documentation and honesty notes.
    """
    out, _ = deep_field_polylaplacian_with_report(
        x,
        layers,
        k,
        mode=mode,
        n_directions=n_directions,
        budget=budget,
        seed=seed,
    )
    return out

Related APIsomnibias.torch.laplacian.deep_field_polylaplacian_with_report↗

FUNCTION

deep_field_polylaplacian_with_report

omnibias.torch.laplacian.deep_field_polylaplacian_with_reportSource ↗

deep_field_polylaplacian_with_report(x: Tensor, layers: Sequence[LayerSpec], k: int, *, mode: str='auto', n_directions: int=256, budget: int | None=None, seed: int=0) -> tuple[Tensor, 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.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
layerspositional-or-keyword	Sequence[LayerSpec]Required	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.
modekeyword-only	strDefault: 'auto'	See the operation contract and implementation below.
n_directionskeyword-only	intDefault: 256	See the operation contract and implementation below.
budgetkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
seedkeyword-only	intDefault: 0	See the operation contract and implementation below.

Returnstuple[Tensor, 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 = torch.as_tensor(np.random.default_rng(22).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="polylaplacian order k must be >= 1"):
        deep_field_polylaplacian(x, layers, 0)

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_polylaplacian_k_must_be_positive" -q

Read the implementation

def deep_field_polylaplacian_with_report(
    x: Tensor,
    layers: Sequence[LayerSpec],
    k: int,
    *,
    mode: str = "auto",
    n_directions: int = 256,
    budget: int | None = None,
    seed: int = 0,
) -> tuple[Tensor, 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`.
    """
    if k < 1:
        raise ValueError(f"polylaplacian order k must be >= 1, got {k}")
    x = torch.as_tensor(x)
    batched = x.ndim > 1
    x2 = x if batched else x.unsqueeze(0)
    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,
            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.torch.laplacian.deep_field_laplacian↗

FUNCTION

deep_field_value_grad_laplacian

omnibias.torch.laplacian.deep_field_value_grad_laplacianSource ↗

deep_field_value_grad_laplacian(x: Tensor, layers: Sequence[LayerSpec], *, allow_almost_everywhere: bool=False) -> tuple[Tensor, Tensor, Tensor]

(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 : Tensor of shape (D,) or (B, D). layers : sequence of (W, b, spec) As consumed by omnibias.torch.jet.mlp_jet / omnibias.torch.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 : Tensor of shape (C,) or (B, C) grad : Tensor of shape (D, C) or (B, D, C) laplacian : Tensor 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.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	Tensor of shape (D,) or (B, D).
layerspositional-or-keyword	Sequence[LayerSpec]Required	sequence of (W, b, spec) As consumed by omnibias.torch.jet.mlp_jet / omnibias.torch.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-only	boolDefault: False	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor, Tensor]

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 · relu laplacian opt in

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_relu_laplacian_opt_in() -> None:
    layers = _make_layers(3, 4, depth=1, activation="relu", seed=31)
    x = torch.as_tensor(np.random.default_rng(32).normal(size=(3,)), dtype=torch.float64)
    lap = deep_field_laplacian(x, layers, allow_almost_everywhere=True)
    assert bool(torch.all(torch.isfinite(lap)))

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_relu_laplacian_opt_in" -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 = torch.as_tensor(np.random.default_rng(7).normal(size=(3, 5)), dtype=torch.float64)
    value, grad, lap = deep_field_value_grad_laplacian(x, layers)
    lap_only = deep_field_laplacian(x, layers)
    assert torch.allclose(lap, lap_only)
    assert value.shape == (3, 1)
    assert grad.shape == (3, 5, 1)

python -m pytest packages/omnibias-torch/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: Tensor,
    layers: Sequence[LayerSpec],
    *,
    allow_almost_everywhere: bool = False,
) -> tuple[Tensor, Tensor, Tensor]:
    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 : Tensor of shape ``(D,)`` or ``(B, D)``.
    layers : sequence of ``(W, b, spec)``
        As consumed by :func:`omnibias.torch.jet.mlp_jet` /
        :func:`omnibias.torch.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 : Tensor of shape ``(C,)`` or ``(B, C)``
    grad : Tensor of shape ``(D, C)`` or ``(B, D, C)``
    laplacian : Tensor 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.
    """
    if not layers:
        raise ValueError("layers must be non-empty")
    x = torch.as_tensor(x)
    batched = x.ndim > 1
    a = x if batched else x.unsqueeze(0)
    B, dim = a.shape
    L = torch.zeros_like(a)  # (B, D); Laplacian of the identity map is 0.
    J: Tensor | None = None  # sentinel for the identity Jacobian (never materialised)
    for W, b, spec in layers:
        W = torch.as_tensor(W)
        u = torch.matmul(a, W.t())
        if b is not None:
            b_t = torch.as_tensor(b)
            u = u + (b_t if b_t.ndim == 1 else b_t)
        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 = W.t().unsqueeze(0).expand(B, dim, W.shape[0])
        else:
            J_u = torch.tensordot(J, W, dims=([-1], [-1]))  # (B, D, H_out)
        L_u = torch.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.unsqueeze(1) * J_u
        if first_layer:
            row_norm_sq = torch.sum(W * W, dim=-1)  # (H,) == ||W_h||^2
            lap_contrib = sigma_pp * row_norm_sq
        else:
            lap_contrib = sigma_pp * torch.sum(J_u * J_u, dim=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.torch.get_activation↗

FUNCTION

neural_field_hessian

omnibias.torch.laplacian.neural_field_hessianSource ↗

neural_field_hessian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, activation: str | ActivationSpec[Tensor]) -> Tensor

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

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

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

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.

Read the implementation

def neural_field_hessian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    activation: str | ActivationSpec[Tensor],
) -> Tensor:
    """``H_x f`` -- the full ``D x D`` Hessian of the one-layer field.

    ``b`` does not enter the Hessian, so it is not a parameter here. The
    Hessian is symmetric by construction (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 = torch.matmul(x, W.t()) + beta  # (..., H)
    sigma_pp = spec.fastpath(z, 2)  # (..., H)
    weights = sigma_pp * c  # (..., H)
    return torch.einsum("...h,hi,hj->...ij", weights, W, W)

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_laplacian

omnibias.torch.laplacian.neural_field_laplacianSource ↗

neural_field_laplacian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, activation: str | ActivationSpec[Tensor]) -> Tensor

nabla_x^2 f(x) in closed form.

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

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

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 · depth 1 collapses onto neural field 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_depth_1_collapses_onto_neural_field_laplacian() -> None:
    for activation in _RICCATI:
        dim, hidden = 6, 5
        layers = _make_layers(dim, hidden, depth=1, activation=activation, seed=4)
        (W1, b1, _), (W2, b2, _) = layers
        x = torch.as_tensor(
            np.random.default_rng(5).normal(scale=0.3, size=(dim,)), dtype=torch.float64
        )
        deep = deep_field_laplacian(x, layers)[0]
        shallow = neural_field_laplacian(x, W1, b1, c=W2[0, :], activation=activation)
        assert float(torch.abs(deep - shallow)) < 1e-13, activation

python -m pytest "packages/omnibias-torch/tests/test_deep_laplacian.py::test_depth_1_collapses_onto_neural_field_laplacian" -q

Read the implementation

def neural_field_laplacian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    activation: str | ActivationSpec[Tensor],
) -> Tensor:
    """``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 = torch.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(dim=-1)  # (H,)
    return (sigma_pp * (c * row_norm_sq)) @ torch.ones_like(c)

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_local_p4_over_psi

omnibias.torch.laplacian.neural_field_local_p4_over_psiSource ↗

neural_field_local_p4_over_psi(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor]) -> Tensor

Relativistic mass-velocity local-energy operator p^4 psi / psi.

L_rel(x) = p^4 psi(x) / psi(x) = Delta^2 f / f for a wavefunction psi = f represented directly by the one-layer field.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def neural_field_local_p4_over_psi(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    b: Tensor | float,
    activation: str | ActivationSpec[Tensor],
) -> Tensor:
    """Relativistic mass-velocity local-energy operator ``p^4 psi / psi``.

    ``L_rel(x) = p^4 psi(x) / psi(x) = Delta^2 f / f`` for a wavefunction
    ``psi = f`` represented directly by the one-layer field.
    """
    f, p4_psi = neural_field_value_and_polylaplacian(x, W, beta, c, b, activation, k=2)
    return p4_psi / f

Related APIsomnibias.torch.laplacian.neural_field_value_and_polylaplacian↗

FUNCTION

neural_field_polylaplacian

omnibias.torch.laplacian.neural_field_polylaplacianSource ↗

neural_field_polylaplacian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, activation: str | ActivationSpec[Tensor], k: int) -> Tensor

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. Memory cost: O(B * H), independent of k.

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.

ReturnsTensor

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.

Read the implementation

def neural_field_polylaplacian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    activation: str | ActivationSpec[Tensor],
    k: int,
) -> Tensor:
    """``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. Memory cost: ``O(B * H)``,
    independent of ``k``.
    """
    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 = torch.matmul(x, W.t()) + beta  # (..., H)
    sigma_2k = spec.fastpath(z, 2 * k)  # (..., H)
    row_norm_sq = (W * W).sum(dim=-1)  # (H,)
    row_norm_2k = row_norm_sq**k  # (H,)
    return sigma_2k @ (c * row_norm_2k)

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_value

omnibias.torch.laplacian.neural_field_valueSource ↗

neural_field_value(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor]) -> Tensor

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

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def neural_field_value(
    x: Tensor,  # (B, D) or (D,)
    W: Tensor,  # (H, D)
    beta: Tensor,  # (H,)
    c: Tensor,  # (H,)
    b: Tensor | float,  # scalar
    activation: str | ActivationSpec[Tensor],
) -> Tensor:
    """``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 = torch.matmul(x, W.t()) + beta  # (B, H) or (H,)
    return b + spec.forward(z) @ c  # (B,) or scalar

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_value_and_laplacian

omnibias.torch.laplacian.neural_field_value_and_laplacianSource ↗

neural_field_value_and_laplacian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor]) -> tuple[Tensor, Tensor]

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

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

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: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    b: Tensor | float,
    activation: str | ActivationSpec[Tensor],
) -> tuple[Tensor, Tensor]:
    """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 = torch.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(dim=-1)
    lap = sigma_pp @ (c * row_norm_sq)
    return f, lap

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_value_and_polylaplacian

omnibias.torch.laplacian.neural_field_value_and_polylaplacianSource ↗

neural_field_value_and_polylaplacian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor], k: int) -> tuple[Tensor, Tensor]

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

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

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.

Read the implementation

def neural_field_value_and_polylaplacian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    b: Tensor | float,
    activation: str | ActivationSpec[Tensor],
    k: int,
) -> tuple[Tensor, Tensor]:
    """Returns ``(f(x), Delta^k f(x))`` in one fused pass."""
    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 = torch.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(dim=-1)
    row_norm_2k = row_norm_sq**k
    poly_lap = sigma_2k @ (c * row_norm_2k)
    return f, poly_lap

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_value_grad_hessian

omnibias.torch.laplacian.neural_field_value_grad_hessianSource ↗

neural_field_value_grad_hessian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor]) -> tuple[Tensor, Tensor, Tensor]

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

See omnibias.jax.laplacian.neural_field_value_grad_hessian for the full exposition; this is its mathematically equivalent torch twin.

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

Parameters

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor, Tensor]

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.

Read the implementation

def neural_field_value_grad_hessian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    b: Tensor | float,
    activation: str | ActivationSpec[Tensor],
) -> tuple[Tensor, Tensor, Tensor]:
    """Returns ``(f(x), grad_x f(x), H_x f(x))`` in one fused pass.

    See :func:`omnibias.jax.laplacian.neural_field_value_grad_hessian` for
    the full exposition; this is its bit-identical torch twin.
    """
    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 = torch.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 = torch.einsum("...h,hi,hj->...ij", weights, W, W)
    return f, grad, hessian

Related APIsomnibias.torch.get_activation↗

FUNCTION

neural_field_value_grad_laplacian

omnibias.torch.laplacian.neural_field_value_grad_laplacianSource ↗

neural_field_value_grad_laplacian(x: Tensor, W: Tensor, beta: Tensor, c: Tensor, b: Tensor | float, activation: str | ActivationSpec[Tensor]) -> tuple[Tensor, Tensor, Tensor]

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

Name	Type & default	Contract
xpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
Wpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
cpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	Tensor | floatRequired	See the operation contract and implementation below.
activationpositional-or-keyword	str | ActivationSpec[Tensor]Required	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor, Tensor]

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_grad_laplacian(
    x: Tensor,
    W: Tensor,
    beta: Tensor,
    c: Tensor,
    b: Tensor | float,
    activation: str | ActivationSpec[Tensor],
) -> tuple[Tensor, Tensor, Tensor]:
    """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 = torch.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(dim=-1)  # (H,)
    lap = sigma_pp @ (c * row_norm_sq)  # (..., )
    return f, grad, lap

Related APIsomnibias.torch.get_activation↗

FUNCTION

restrict_first_layer

omnibias.torch.laplacian.restrict_first_layerSource ↗

restrict_first_layer(layers: Sequence[LayerSpec], x0: Tensor, 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.torch.fields.jet_mlp, to isolate the spatial axes of a field that also carries a non-spatial (e.g. time) coordinate.

Parameters

Name	Type & default	Contract
layerspositional-or-keyword	Sequence[LayerSpec]Required	See the operation contract and implementation below.
x0positional-or-keyword	TensorRequired	See the operation contract and implementation below.
supportpositional-or-keyword	tuple[int, ...]Required	See 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 = torch.as_tensor(np.random.default_rng(24).normal(size=(3,)), dtype=torch.float64)
    with pytest.raises(ValueError, match="does not support order"):
        deep_field_polylaplacian(x, layers, 2, mode="support")

python -m pytest "packages/omnibias-torch/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 = torch.as_tensor(np.random.default_rng(11).normal(size=(dim,)), dtype=torch.float64)
    support = (1, 3, 4)
    local = restrict_first_layer(layers, x0, support)
    y = torch.as_tensor(
        np.random.default_rng(12).normal(size=(len(support),)) * 0.1, dtype=torch.float64
    )
    x_full = x0.clone()
    x_full[list(support)] = x_full[list(support)] + y

    def value(layer_list, xi):
        a = xi
        for W, b, spec in layer_list:
            u = a @ W.t() + b
            a = u if spec is None else get_activation(spec).forward(u)
        return a

    got = value(local, y)
    want = value(layers, x_full)
    assert torch.allclose(got, want, atol=1e-12)

python -m pytest packages/omnibias-torch/tests/test_deep_laplacian.py -k test_restrict_first_layer_is_exact -q

Read the implementation

def restrict_first_layer(
    layers: Sequence[LayerSpec], x0: Tensor, 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.torch.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 = torch.as_tensor(W0)
    cols = torch.as_tensor(support, dtype=torch.long, device=W0.device)
    W_local = W0.index_select(1, cols)
    shift = torch.matmul(W0, torch.as_tensor(x0))
    b_local = shift if b0 is None else shift + torch.as_tensor(b0)
    return [(W_local, b_local, spec0), *layers[1:]]

← Package overviewRunnable learning examples →