omnibias.jax.architectures.pinn

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.architectures.pinn

Module contract & mathematical conventions

Deep, arbitrary-order omnibias PINNs for JAX -- algebraically corresponding twins of torch.

JetMLP is the JAX counterpart of omnibias.torch.architectures.pinn.JetMLP: a deep fully-connected network whose *every* mixed input partial D^alpha u(x) up to total order N comes **exactly** from one forward pass of omnibias.jax.jet_mv.mlp_jet_mv -- no jax.grad stacking, no finite differences. Because both backends call the same Faa di Bruno multivariate-jet kernel, the derivatives are *algebraically corresponding* across torch and JAX in double precision (enable jax_enable_x64).

Spectral-bias mitigation mirrors the torch surface:

FourierFeatureMLP lifts the input through a random Fourier-feature encoding gamma(x) = [cos(B x), sin(B x)] (Tancik et al. 2020). Because cos(z) = sin(z + pi/2) the encoding is a single omnibias sin layer, so the composite still yields every closed-form derivative through mlp_jet_mv. Build one with make_fourier_feature_mlp.
make_siren builds a SIREN (Sitzmann et al. 2020): a JetMLP with sin activations and the SIREN initialisation; its derivative tower is exact for every order (sin^{(n)}(z) = sin(z + n pi/2)).

Each network is a frozen dataclass registered as a JAX pytree (its weight/bias arrays are the leaves), so it flows through jax.grad / jax.jit for training while omnibias supplies the spatial differential operator.

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

CLASS

FourierFeatureMLP

omnibias.jax.architectures.FourierFeatureMLPSource ↗

class FourierFeatureMLP

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

Spectral-bias-mitigating MLP with a sin-encoded random Fourier front end (JAX).

The input is lifted by gamma(x) = [cos(B x), sin(B x)] (Tancik et al. 2020) and a standard MLP body acts on it. Because cos(z) = sin(z + pi/2) the encoding is a single omnibias sin layer (gamma(x) = sin([B; B] x + [pi/2; 0])), so D^alpha u(x) stays exactly closed form to arbitrary order through omnibias.jax.jet_mv.mlp_jet_mv. mathematically equivalent twin of omnibias.torch.architectures.pinn.FourierFeatureMLP.

Parameters

w_ff, b_ff The sin-encoding layer: w_ff has shape (2 F_total, in_dim) and equals [B; B]; b_ff is [pi/2 * 1_{F_total}; 0_{F_total}]. weights, biases Body parameters (hidden layers + affine readout). base_spec Body activation with a closed-form derivative fast path. in_dim, out_dim Input coordinate count and output component count. num_features Fourier features per band (F); the encoding width is 2 * F * len(scales). scales Frequency bands used to build B (kept for introspection).

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
`w_ff`dataclass field	`Array`Required	The sin-encoding layer: `w_ff` has shape `(2 F_total, in_dim)` and equals `[B; B]`; `b_ff` is `[pi/2 * 1_{F_total}; 0_{F_total}]`.
`b_ff`dataclass field	`Array`Required	See the operation contract and implementation below.
`weights`dataclass field	`tuple[Array, ...]`Required	Body parameters (hidden layers + affine readout).
`biases`dataclass field	`tuple[Array, ...]`Required	See the operation contract and implementation below.
`base_spec`dataclass field	`JaxActivationSpec`Required	Body activation with a closed-form derivative fast path.
`in_dim`dataclass field	`int`Required	Input coordinate count and output component count.
`out_dim`dataclass field	`int`Required	See the operation contract and implementation below.
`num_features`dataclass field	`int`Required	Fourier features per band (`F`); the encoding width is `2 * F * len(scales)`.
`scales`dataclass field	`tuple[float, ...]`Required	Frequency bands used to build `B` (kept for introspection).

Executed regression · fourier value equals jet row0

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

Full file, imports & fixtures ↗

def test_fourier_value_equals_jet_row0() -> None:
    net = make_fourier_feature_mlp(2, num_features=8, hidden=16, depth=2, seed=1)
    x = _x(2)
    jet = net.jet(x, 2)
    assert jnp.allclose(jet[:, 0, :], net.value(x), atol=1e-13)

python -m pytest "packages/omnibias-jax/tests/test_spectral_pinn.py::test_fourier_value_equals_jet_row0" -q

Regression example · fourier pytree grad flows

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_fourier_pytree_grad_flows() -> None:
    net = make_fourier_feature_mlp(2, num_features=6, hidden=10, depth=2, seed=9)
    x = _x(10, b=4)

    def loss(n: FourierFeatureMLP) -> jax.Array:
        v, g, h = n.value_grad_hessian(x)
        return jnp.mean((g[:, 1, 0] - 0.1 * h[:, 0, 0, 0]) ** 2)

    grads = jax.grad(loss)(net)
    assert jnp.all(jnp.isfinite(grads.w_ff))
    for w in grads.weights:
        assert jnp.all(jnp.isfinite(w))

python -m pytest packages/omnibias-jax/tests/test_spectral_pinn.py -k test_fourier_pytree_grad_flows -q

Fields & methods

w_fffield

w_ff: Array

The sin-encoding layer: w_ff has shape (2 F_total, in_dim) and equals [B; B]; b_ff is [pi/2 * 1_{F_total}; 0_{F_total}].

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

Name	Type & default	Contract
`weights`dataclass field	`tuple[Array, ...]`Required	Per-linear parameters; `len == depth + 1` (the last is the affine readout). `weights[i]` has shape `(out_i, in_i)` and `biases[i]` shape `(out_i,)`.
`biases`dataclass field	`tuple[Array, ...]`Required	See the operation contract and implementation below.
`spec`dataclass field	`JaxActivationSpec`Required	Base activation with a closed-form derivative fast path.
`in_dim`dataclass field	`int`Required	Input coordinate count and output component count.
`out_dim`dataclass field	`int`Required	See the operation contract and implementation below.

Name	Type & default	Contract
`in_dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`num_features`positional-or-keyword	`int`Default: 64	See the operation contract and implementation below.
`hidden`positional-or-keyword	`int`Default: 64	See the operation contract and implementation below.
`out_dim`positional-or-keyword	`int`Default: 1	See the operation contract and implementation below.
`depth`positional-or-keyword	`int`Default: 3	See the operation contract and implementation below.
`base`positional-or-keyword	`str \| JaxActivationSpec`Default: 'tanh'	See the operation contract and implementation below.
`frequency_scale`keyword-only	`float \| Sequence[float]`Default: 1.0	See the operation contract and implementation below.
`seed`keyword-only	`int`Default: 0	See the operation contract and implementation below.
`weight_init_scale`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.
`dtype`keyword-only	`Any`Default: jnp.float64	See the operation contract and implementation below.

omnibias.jax.architectures.pinn

FourierFeatureMLP​

Parameters

Fields & methods

Parameters

Parameters

Parameters

Parameters

Parameters

Parameters

Parameters

JetMLP​

Parameters

Fields & methods

Parameters

Parameters

Parameters

Parameters

Parameters

Parameters

Parameters

make_fourier_feature_mlp​

Parameters

make_jet_mlp​

Parameters

make_siren​

Parameters

FourierFeatureMLP

JetMLP

make_fourier_feature_mlp

make_jet_mlp

make_siren