NN API

The pyrox_nn subpackage ships uncertainty-aware neural network layers in four families:

1. Geographic / spherical encoders (re-exported from geonnax) — degree/radian, lon/lat, cyclic, spherical-harmonic, and Slepian preprocessing for geophysical inputs.
2. Dense / Bayesian-linear layers (pyrox_nn._dense) — reparameterization, Flipout, hierarchical, NCP, DVI, rank-1 ensemble, and variational-dropout variants of Wx + b.
3. Spectral / GP-flavoured layers — random-feature kernel maps, SNGP and VSSGP heads, deep random-feature expansions.
4. Ensembles & output heads — BatchEnsemble layers, heteroscedastic Monte-Carlo output heads, the Gaussian regression head.
5. Bayesian Neural Field stack (pyrox_nn._bnf) — five layers that together implement the BNF architecture (Saad et al., Nat. Comms. 2024); BayesianNeuralField wraps the deterministic geonnax.NeuralFieldMLP and takes an optional coordinate encoder (for example spherical harmonics on lon/lat).
6. Pure-JAX feature helpers (re-exported from geonnax.basis) — pandas-free building blocks the BNF layers wrap.

See also: Geo encoders for the longitude/latitude and spherical-harmonic API surface.

Dense / Bayesian-linear layers

DenseReparameterization

Bases: PyroxModule

Bayesian dense layer via the reparameterization trick.

Samples weight and bias from learned Gaussian posteriors at every forward pass. Registers NumPyro sample sites so the KL between the variational posterior and the prior is tracked by the ELBO.

\[ W \sim \mathcal{N}(\mu_W, \sigma_W^2), \quad b \sim \mathcal{N}(\mu_b, \sigma_b^2), \quad y = x W + b. \]

Attributes:

Name	Type	Description
in_features	int	Input dimension.
out_features	int	Output dimension.
bias	bool	Whether to include a bias term.
prior_scale	float	Scale of the isotropic Gaussian prior on weights and bias. The prior mean is zero.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseReparameterization(PyroxModule):
    r"""Bayesian dense layer via the reparameterization trick.

    Samples weight and bias from learned Gaussian posteriors at every
    forward pass. Registers NumPyro sample sites so the KL between the
    variational posterior and the prior is tracked by the ELBO.

    $$
    W \sim \mathcal{N}(\mu_W, \sigma_W^2), \quad
    b \sim \mathcal{N}(\mu_b, \sigma_b^2), \quad
    y = x W + b.
    $$

    Attributes:
        in_features: Input dimension.
        out_features: Output dimension.
        bias: Whether to include a bias term.
        prior_scale: Scale of the isotropic Gaussian prior on weights
            and bias. The prior mean is zero.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    prior_scale: float = eqx.field(static=True, default=1.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        prior_w = dist.Normal(
            jnp.zeros((self.in_features, self.out_features)),
            self.prior_scale,
        ).to_event(2)
        W = self.pyrox_sample("weight", prior_w)
        out = einx.dot("... din, din dout -> ... dout", x, W)
        if self.bias:
            prior_b = dist.Normal(
                jnp.zeros(self.out_features), self.prior_scale
            ).to_event(1)
            b = self.pyrox_sample("bias", prior_b)
            out = out + b
        return out

DenseFlipout

Bases: PyroxModule

Bayesian dense layer with Flipout sign-flip structure.

Samples weight from the prior and applies per-example Rademacher sign flips to the weight perturbation (Wen et al., 2018). Under a NumPyro guide that learns the posterior mean, the sign flips decorrelate gradient estimates across minibatch examples.

In model mode (no guide) this is equivalent to DenseReparameterization — the Flipout variance reduction activates when a guide provides a posterior centered at a learned mean.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
out_features	int	Output dimension.
bias	bool	Whether to include a bias term.
prior_scale	float	Scale of the isotropic Gaussian prior.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseFlipout(PyroxModule):
    r"""Bayesian dense layer with Flipout sign-flip structure.

    Samples weight from the prior and applies per-example Rademacher
    sign flips to the weight perturbation (Wen et al., 2018). Under a
    NumPyro guide that learns the posterior mean, the sign flips
    decorrelate gradient estimates across minibatch examples.

    In model mode (no guide) this is equivalent to
    `DenseReparameterization` — the Flipout variance reduction
    activates when a guide provides a posterior centered at a learned
    mean.

    Attributes:
        in_features: Input dimension.
        out_features: Output dimension.
        bias: Whether to include a bias term.
        prior_scale: Scale of the isotropic Gaussian prior.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    prior_scale: float = eqx.field(static=True, default=1.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        prior_w = dist.Normal(
            jnp.zeros((self.in_features, self.out_features)),
            self.prior_scale,
        ).to_event(2)
        W = self.pyrox_sample("weight", prior_w)
        out = einx.dot("... din, din dout -> ... dout", x, W)

        if self.bias:
            prior_b = dist.Normal(
                jnp.zeros(self.out_features), self.prior_scale
            ).to_event(1)
            b = self.pyrox_sample("bias", prior_b)
            out = out + b
        return out

DenseVariational

Bases: PyroxModule

Dense layer with a user-supplied prior factory.

Provides flexibility over the weight prior by accepting a callable that builds the prior distribution given the layer shape. The model samples from the prior; the posterior is handled by a NumPyro guide (e.g., AutoNormal).

Attributes:

Name	Type	Description
in_features	int	Input dimension.
out_features	int	Output dimension.
make_prior	Callable[..., Any]	Callable (in_features, out_features) -> Distribution.
bias	bool	Whether to include a bias term.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseVariational(PyroxModule):
    r"""Dense layer with a user-supplied prior factory.

    Provides flexibility over the weight prior by accepting a callable
    that builds the prior distribution given the layer shape. The
    model samples from the prior; the posterior is handled by a NumPyro
    guide (e.g., ``AutoNormal``).

    Attributes:
        in_features: Input dimension.
        out_features: Output dimension.
        make_prior: Callable ``(in_features, out_features) -> Distribution``.
        bias: Whether to include a bias term.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    make_prior: Callable[..., Any] = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        prior = self.make_prior(self.in_features, self.out_features)
        W = self.pyrox_sample("weight", prior)
        out = einx.dot("... din, din dout -> ... dout", x, W)
        if self.bias:
            b = self.pyrox_sample(
                "bias",
                dist.Normal(jnp.zeros(self.out_features), 1.0).to_event(1),
            )
            out = out + b
        return out

DenseDVI

Bases: PyroxModule

Deterministic Variational Inference dense layer (Wu et al., 2018).

Propagates a Gaussian distribution through the linear layer analytically — there is no Monte Carlo sampling. The input is a diagonal-covariance Gaussian \((\mu_x, \sigma_x^2)\), the output is the (still-diagonal) Gaussian \((\mu_y, \sigma_y^2)\) induced by an independent-Gaussian variational posterior \(q(W) = \mathcal{N}(M, S)\) over the weights and a separate diagonal posterior on the bias.

With weight posterior mean \(M\) (shape \(D_\mathrm{in}\times D_\mathrm{out}\)) and per-element posterior variance \(S\) (same shape):

\[ \mu_y = \mu_x M, \qquad \sigma_y^2 = \sigma_x^2 (M \circ M) + (\mu_x^{\circ 2} + \sigma_x^2)\,S, \]

plus the bias mean / variance if enabled. Compared to MC estimators, DVI gives zero-variance gradients of the ELBO at the cost of propagating second-order statistics layer by layer (so it only really pays off when all dense layers in a block are DVI; a single DVI layer in a sampling stack just adds bookkeeping).

The KL between the diagonal-Gaussian variational posterior and a fixed isotropic Gaussian prior \(p(W) = \mathcal{N}(0, \pi^2)\) is closed-form and is registered with numpyro.factor so SVI's Trace_ELBO picks it up:

\[ \mathrm{KL}\!\bigl[\mathcal{N}(M, S) \,\big\|\, \mathcal{N}(0, \pi^2)\bigr] = \sum_{ij}\Bigl[ \log \pi - \tfrac12\log S_{ij} + \frac{S_{ij} + M_{ij}^2}{2\pi^2} - \tfrac12 \Bigr]. \]

Plate semantics

Same as the rest of the pyrox Bayesian dense family — call this layer outside numpyro.plate("data", ..., subsample_size=...). The KL is a weight-prior term: it sums over the weight and bias matrices, not over the batch, so it's a single scalar per layer. numpyro.factor is still a sample-type site, though, and putting it inside a subsampled plate would broadcast the scalar to the plate dim and apply scale = N/B — the same over-counting trap that affects every per-layer numpyro.factor. Keep this layer at the top of the model (or outside any data plate) and only plate the observation likelihood:

def model(x, y=None):
    mean, var = dvi(x_mean, x_var)        # KL emitted here
    with numpyro.plate("data", x.shape[0]):
        numpyro.sample("obs",
            dist.Normal(mean, jnp.sqrt(var)), obs=y)

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_\mathrm{in}\).
out_features	int	Output dimension \(D_\mathrm{out}\).
bias	bool	Whether to include a diagonal-Gaussian bias.
prior_scale	float	Std \(\pi\) of the isotropic Gaussian prior.
init_log_var	float	Initial value for the log posterior variance (a small negative number keeps initial draws tight).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> dvi = DenseDVI(in_features=3, out_features=2, pyrox_name="dvi")
>>> mean = jnp.ones((4, 3))
>>> var = 0.1 * jnp.ones((4, 3))
>>> with handlers.seed(rng_seed=0):
...     out_mean, out_var = dvi(mean, var)
>>> out_mean.shape, out_var.shape
((4, 2), (4, 2))

References

Wu, A., Nowozin, S., Meeds, E., Turner, R. E., Hernández-Lobato, J. M., & Gaunt, A. L. (2018). Deterministic Variational Inference for Robust Bayesian Neural Networks. ICLR.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseDVI(PyroxModule):
    r"""Deterministic Variational Inference dense layer (Wu et al., 2018).

    Propagates a *Gaussian distribution* through the linear layer
    analytically — there is no Monte Carlo sampling. The input is a
    diagonal-covariance Gaussian $(\mu_x, \sigma_x^2)$, the
    output is the (still-diagonal) Gaussian $(\mu_y, \sigma_y^2)$
    induced by an independent-Gaussian variational posterior
    $q(W) = \mathcal{N}(M, S)$ over the weights and a separate
    diagonal posterior on the bias.

    With weight posterior mean $M$
    (shape $D_\mathrm{in}\times D_\mathrm{out}$) and per-element
    posterior variance $S$ (same shape):

    $$
    \mu_y = \mu_x M, \qquad
    \sigma_y^2 = \sigma_x^2 (M \circ M)
               + (\mu_x^{\circ 2} + \sigma_x^2)\,S,
    $$

    plus the bias mean / variance if enabled. Compared to MC
    estimators, DVI gives zero-variance gradients of the ELBO at the
    cost of propagating second-order statistics layer by layer (so it
    only really pays off when *all* dense layers in a block are DVI;
    a single DVI layer in a sampling stack just adds bookkeeping).

    The KL between the diagonal-Gaussian variational posterior and a
    fixed isotropic Gaussian prior $p(W) = \mathcal{N}(0, \pi^2)$
    is closed-form and is registered with `numpyro.factor` so
    SVI's ``Trace_ELBO`` picks it up:

    $$
    \mathrm{KL}\!\bigl[\mathcal{N}(M, S) \,\big\|\, \mathcal{N}(0, \pi^2)\bigr]
    = \sum_{ij}\Bigl[
        \log \pi - \tfrac12\log S_{ij}
        + \frac{S_{ij} + M_{ij}^2}{2\pi^2} - \tfrac12
      \Bigr].
    $$

    Plate semantics:
        Same as the rest of the pyrox Bayesian dense family — call
        this layer **outside** ``numpyro.plate("data", ..., subsample_size=...)``.
        The KL is a *weight-prior* term: it sums over the weight and
        bias matrices, not over the batch, so it's a single scalar
        per layer. ``numpyro.factor`` is still a sample-type site,
        though, and putting it inside a subsampled plate would broadcast
        the scalar to the plate dim and apply ``scale = N/B`` — the
        same over-counting trap that affects every per-layer
        ``numpyro.factor``. Keep this layer at the top of the model
        (or outside any data plate) and only plate the observation
        likelihood:

            def model(x, y=None):
                mean, var = dvi(x_mean, x_var)        # KL emitted here
                with numpyro.plate("data", x.shape[0]):
                    numpyro.sample("obs",
                        dist.Normal(mean, jnp.sqrt(var)), obs=y)

    Attributes:
        in_features: Input dimension $D_\mathrm{in}$.
        out_features: Output dimension $D_\mathrm{out}$.
        bias: Whether to include a diagonal-Gaussian bias.
        prior_scale: Std $\pi$ of the isotropic Gaussian prior.
        init_log_var: Initial value for the log posterior variance
            (a small negative number keeps initial draws tight).
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> dvi = DenseDVI(in_features=3, out_features=2, pyrox_name="dvi")
        >>> mean = jnp.ones((4, 3))
        >>> var = 0.1 * jnp.ones((4, 3))
        >>> with handlers.seed(rng_seed=0):
        ...     out_mean, out_var = dvi(mean, var)
        >>> out_mean.shape, out_var.shape
        ((4, 2), (4, 2))

    References:
        Wu, A., Nowozin, S., Meeds, E., Turner, R. E., Hernández-Lobato,
        J. M., & Gaunt, A. L. (2018). *Deterministic Variational
        Inference for Robust Bayesian Neural Networks.* ICLR.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    prior_scale: float = eqx.field(static=True, default=1.0)
    init_log_var: float = eqx.field(static=True, default=-3.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    def __post_init__(self) -> None:
        if self.in_features <= 0 or self.out_features <= 0:
            raise ValueError(
                "in_features and out_features must be > 0; "
                f"got {self.in_features=}, {self.out_features=}."
            )
        if self.prior_scale <= 0:
            raise ValueError(f"prior_scale must be > 0; got {self.prior_scale}.")

    @pyrox_method
    def __call__(
        self,
        mean: Float[Array, "*batch D_in"],
        var: Float[Array, "*batch D_in"],
    ) -> tuple[Float[Array, "*batch D_out"], Float[Array, "*batch D_out"]]:
        if mean.shape != var.shape:
            raise ValueError(f"mean.shape {mean.shape} != var.shape {var.shape}.")
        if mean.shape[-1] != self.in_features:
            raise ValueError(
                f"mean.shape[-1] = {mean.shape[-1]} does not match "
                f"in_features = {self.in_features}."
            )

        W_mean = self.pyrox_param(
            "weight_mean", jnp.zeros((self.in_features, self.out_features))
        )
        W_log_var = self.pyrox_param(
            "weight_log_var",
            jnp.full(
                (self.in_features, self.out_features),
                float(self.init_log_var),
            ),
        )
        W_var = jnp.exp(W_log_var)

        out_mean = einx.dot("... din, din dout -> ... dout", mean, W_mean)
        out_var = einx.dot("... din, din dout -> ... dout", var, W_mean**2) + einx.dot(
            "... din, din dout -> ... dout", mean**2 + var, W_var
        )

        if self.bias:
            b_mean = self.pyrox_param("bias_mean", jnp.zeros(self.out_features))
            b_log_var = self.pyrox_param(
                "bias_log_var",
                jnp.full((self.out_features,), float(self.init_log_var)),
            )
            out_mean = out_mean + b_mean
            out_var = out_var + jnp.exp(b_log_var)

        # Closed-form KL[N(M, S) || N(0, prior_scale^2)] over weights and bias.
        # Use math.log + jnp.asarray cast so prior constants pick up the
        # same dtype as the params (avoid silent float64 promotion under
        # jax_enable_x64 + float32 params).
        log_prior_scale = jnp.asarray(math.log(self.prior_scale), dtype=W_mean.dtype)
        prior_var = jnp.asarray(self.prior_scale**2, dtype=W_mean.dtype)
        kl = jnp.sum(
            _diag_gaussian_kl(
                W_mean,
                W_var,
                W_log_var,
                prior_mean=0.0,
                log_prior_scale=log_prior_scale,
                prior_var=prior_var,
            )
        )
        if self.bias:
            b_var = jnp.exp(b_log_var)
            kl = kl + jnp.sum(
                _diag_gaussian_kl(
                    b_mean,
                    b_var,
                    b_log_var,
                    prior_mean=0.0,
                    log_prior_scale=log_prior_scale,
                    prior_var=prior_var,
                )
            )
        # Add -KL to the model log density. This is a per-layer scalar
        # (sums over the weight / bias matrices, not the batch) — its
        # emission is independent of any data plate the layer lives in.
        numpyro.factor(self._pyrox_fullname("kl"), -kl)
        return out_mean, out_var

DenseHierarchical

Bases: PyroxModule

Hierarchical Bayesian dense layer with multiplicative shrinkage.

Decomposes the effective weight matrix into a deterministic base \(\theta \in \mathbb{R}^{D_\mathrm{in} \times D_\mathrm{out}}\) multiplied row-wise by a per-input-unit local scale \(z^{(\mathrm{loc})} \in \mathbb{R}^{D_\mathrm{in}}\) and an overall global scale \(z^{(\mathrm{glob})} \in \mathbb{R}\),

\[ W_{ij} = \theta_{ij} \cdot z_i^{(\mathrm{loc})} \cdot z^{(\mathrm{glob})}, \]

with isotropic Gaussian priors centred at one,

\[ z_i^{(\mathrm{loc})} \sim \mathcal{N}(1, \sigma_\mathrm{loc}^2), \qquad z^{(\mathrm{glob})} \sim \mathcal{N}(1, \sigma_\mathrm{glob}^2). \]

The local scale prunes individual input units (a column of \(\theta\) whose z_loc posterior concentrates near zero is effectively switched off) while the global scale modulates the overall layer activation — the same hierarchical-shrinkage structure used by horseshoe-style BNNs (Louizos et al., 2017). Both scales are pyrox_sample sites so any standard NumPyro guide (AutoNormal, etc.) drives the variational posterior; the deterministic base \(\theta\) and bias are pyrox_param.

Plate semantics

Same as the rest of pyrox_nn's Bayesian dense layers — call outside numpyro.plate("data", ..., subsample_size=...) and only plate the observation likelihood, otherwise the per-layer prior log-probabilities of z_loc and z_glob get scaled by the subsample ratio.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_\mathrm{in}\).
out_features	int	Output dimension \(D_\mathrm{out}\).
bias	bool	Whether to include a deterministic bias term.
prior_local_scale	float	Std \(\sigma_\mathrm{loc}\) of the local scale prior.
prior_global_scale	float	Std \(\sigma_\mathrm{glob}\) of the global scale prior.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> layer = DenseHierarchical(
...     in_features=4, out_features=2, pyrox_name="hier"
... )
>>> x = jnp.ones((3, 4))
>>> with handlers.seed(rng_seed=0):
...     y = layer(x)
>>> y.shape
(3, 2)

References

Louizos, C., Ullrich, K., & Welling, M. (2017). Bayesian Compression for Deep Learning. NeurIPS.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseHierarchical(PyroxModule):
    r"""Hierarchical Bayesian dense layer with multiplicative shrinkage.

    Decomposes the effective weight matrix into a deterministic base
    $\theta \in \mathbb{R}^{D_\mathrm{in} \times D_\mathrm{out}}$
    multiplied row-wise by a per-input-unit local scale
    $z^{(\mathrm{loc})} \in \mathbb{R}^{D_\mathrm{in}}$ and an
    overall global scale $z^{(\mathrm{glob})} \in \mathbb{R}$,

    $$
    W_{ij} = \theta_{ij} \cdot z_i^{(\mathrm{loc})}
             \cdot z^{(\mathrm{glob})},
    $$

    with isotropic Gaussian priors centred at one,

    $$
    z_i^{(\mathrm{loc})} \sim \mathcal{N}(1, \sigma_\mathrm{loc}^2),
    \qquad
    z^{(\mathrm{glob})} \sim \mathcal{N}(1, \sigma_\mathrm{glob}^2).
    $$

    The local scale prunes individual input units (a column of
    $\theta$ whose ``z_loc`` posterior concentrates near zero is
    effectively switched off) while the global scale modulates the
    overall layer activation — the same hierarchical-shrinkage
    structure used by horseshoe-style BNNs (Louizos et al., 2017).
    Both scales are ``pyrox_sample`` sites so any standard NumPyro
    guide (``AutoNormal``, etc.) drives the variational posterior; the
    deterministic base $\theta$ and bias are ``pyrox_param``.

    Plate semantics:
        Same as the rest of ``pyrox_nn``'s Bayesian dense layers — call
        outside ``numpyro.plate("data", ..., subsample_size=...)`` and
        only plate the observation likelihood, otherwise the
        per-layer prior log-probabilities of ``z_loc`` and ``z_glob``
        get scaled by the subsample ratio.

    Attributes:
        in_features: Input dimension $D_\mathrm{in}$.
        out_features: Output dimension $D_\mathrm{out}$.
        bias: Whether to include a deterministic bias term.
        prior_local_scale: Std $\sigma_\mathrm{loc}$ of the local
            scale prior.
        prior_global_scale: Std $\sigma_\mathrm{glob}$ of the
            global scale prior.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> layer = DenseHierarchical(
        ...     in_features=4, out_features=2, pyrox_name="hier"
        ... )
        >>> x = jnp.ones((3, 4))
        >>> with handlers.seed(rng_seed=0):
        ...     y = layer(x)
        >>> y.shape
        (3, 2)

    References:
        Louizos, C., Ullrich, K., & Welling, M. (2017). *Bayesian
        Compression for Deep Learning.* NeurIPS.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    prior_local_scale: float = 0.1
    prior_global_scale: float = 0.1
    pyrox_name: str | None = eqx.field(static=True, default=None)

    def __post_init__(self) -> None:
        if self.prior_local_scale <= 0:
            raise ValueError(
                f"prior_local_scale must be > 0; got {self.prior_local_scale}."
            )
        if self.prior_global_scale <= 0:
            raise ValueError(
                f"prior_global_scale must be > 0; got {self.prior_global_scale}."
            )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        theta = self.pyrox_param(
            "theta", jnp.zeros((self.in_features, self.out_features))
        )
        z_loc = self.pyrox_sample(
            "z_local",
            dist.Normal(jnp.ones(self.in_features), self.prior_local_scale).to_event(1),
        )
        z_glob = self.pyrox_sample(
            "z_global", dist.Normal(1.0, self.prior_global_scale)
        )
        # Effective weight uses the broadcasted multiplicative scaling.
        # Equivalent (and slightly cheaper) to scaling x first then matmul:
        #   y = ((x * z_loc) @ theta) * z_glob.
        out = einx.dot("... din, din dout -> ... dout", x * z_loc, theta) * z_glob
        if self.bias:
            b = self.pyrox_param("b", jnp.zeros(self.out_features))
            out = out + b
        return out

DenseVariationalDropout

Bases: PyroxModule

Sparse variational dropout dense layer.

Implements variational dropout (Kingma et al., 2015) extended by Molchanov et al. (2017) to a log-uniform prior that enables automatic sparsification via per-weight learnable dropout rates. The variational posterior on weights is

\[ q(W_{ij} \mid \theta_{ij}, \alpha_{ij}) = \mathcal{N}\!\bigl(\theta_{ij},\; \alpha_{ij}\,\theta_{ij}^2\bigr). \]

Forward passes use the local reparameterization trick — the pre-activation distribution is closed-form and the noise is sampled once per output unit per batch element rather than once per weight:

\[ \gamma = X\theta, \quad \delta = X^{\circ 2}\,(\alpha \circ \theta^{\circ 2}), \quad Y = \gamma + \sqrt{\delta} \circ \epsilon, \quad \epsilon \sim \mathcal{N}(0, I). \]

The KL between the posterior and the log-uniform prior is approximated analytically (Molchanov et al., 2017) and added to the NumPyro trace via numpyro.factor. SVI then optimizes

\[ \mathcal{L} = \mathbb{E}_q[\log p(y \mid f)] - \mathrm{KL}\bigl[q\,\|\,p\bigr]. \]

Weights with log_alpha > threshold (default 3.0, dropout rate ~0.95) are effectively pruned; inspect the trained pattern via sparsity.

Plate semantics

The KL contribution is registered via numpyro.factor, which is itself a sample-type site and therefore subject to numpyro.plate scaling. To keep the per-layer KL counted once (not scaled by the data-plate's subsample ratio), call the layer outside any plate("data", ..., subsample_size=...) block — the standard pyrox / NumPyro convention for global Bayesian parameters. Plate only the observation likelihood.

Correct (forward outside the data plate):

def model(x, y=None):
    layer = DenseVariationalDropout(in_features=D, out_features=1)
    f = layer(x).squeeze(-1)               # KL emitted here
    with numpyro.plate("data", x.shape[0]):
        numpyro.sample("obs", dist.Normal(f, 0.5), obs=y)

Incorrect (forward inside a subsampled data plate scales KL by N / subsample_size):

def model(x, y=None):
    with numpyro.plate("data", N, subsample_size=B) as idx:
        f = layer(x[idx]).squeeze(-1)      # ⚠ scales KL
        numpyro.sample("obs", ...)

Attributes:

Name	Type	Description
in_features	int	Input dimension.
out_features	int	Output dimension.
bias	bool	Whether to include a bias term.
log_alpha_init	float	Initial value for log_alpha (typically a small negative number, e.g., -5.0).
threshold	float	log_alpha threshold for declaring a weight pruned.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> layer = DenseVariationalDropout(
...     in_features=4, out_features=2, pyrox_name="vd"
... )
>>> x = jnp.ones((3, 4))
>>> with handlers.seed(rng_seed=0):
...     y = layer(x)
>>> y.shape
(3, 2)

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseVariationalDropout(PyroxModule):
    r"""Sparse variational dropout dense layer.

    Implements variational dropout (Kingma et al., 2015) extended by
    Molchanov et al. (2017) to a log-uniform prior that enables
    automatic sparsification via per-weight learnable dropout rates.
    The variational posterior on weights is

    $$
    q(W_{ij} \mid \theta_{ij}, \alpha_{ij}) =
    \mathcal{N}\!\bigl(\theta_{ij},\; \alpha_{ij}\,\theta_{ij}^2\bigr).
    $$

    Forward passes use the *local reparameterization trick* — the
    pre-activation distribution is closed-form and the noise is sampled
    once per output unit per batch element rather than once per weight:

    $$
    \gamma = X\theta, \quad
    \delta = X^{\circ 2}\,(\alpha \circ \theta^{\circ 2}), \quad
    Y = \gamma + \sqrt{\delta} \circ \epsilon, \quad
    \epsilon \sim \mathcal{N}(0, I).
    $$

    The KL between the posterior and the log-uniform prior is
    approximated analytically (Molchanov et al., 2017) and added to the
    NumPyro trace via `numpyro.factor`. SVI then optimizes

    $$
    \mathcal{L} = \mathbb{E}_q[\log p(y \mid f)] - \mathrm{KL}\bigl[q\,\|\,p\bigr].
    $$

    Weights with ``log_alpha > threshold`` (default 3.0, dropout rate
    ~0.95) are effectively pruned; inspect the trained pattern via
    `sparsity`.

    Plate semantics:
        The KL contribution is registered via `numpyro.factor`,
        which is itself a sample-type site and therefore subject to
        ``numpyro.plate`` scaling. To keep the per-layer KL counted
        once (not scaled by the data-plate's subsample ratio), call
        the layer **outside** any ``plate("data", ..., subsample_size=...)``
        block — the standard pyrox / NumPyro convention for global
        Bayesian parameters. Plate only the observation likelihood.

        Correct (forward outside the data plate):

            def model(x, y=None):
                layer = DenseVariationalDropout(in_features=D, out_features=1)
                f = layer(x).squeeze(-1)               # KL emitted here
                with numpyro.plate("data", x.shape[0]):
                    numpyro.sample("obs", dist.Normal(f, 0.5), obs=y)

        Incorrect (forward inside a subsampled data plate scales KL by
        ``N / subsample_size``):

            def model(x, y=None):
                with numpyro.plate("data", N, subsample_size=B) as idx:
                    f = layer(x[idx]).squeeze(-1)      # ⚠ scales KL
                    numpyro.sample("obs", ...)

    Attributes:
        in_features: Input dimension.
        out_features: Output dimension.
        bias: Whether to include a bias term.
        log_alpha_init: Initial value for ``log_alpha`` (typically a
            small negative number, e.g., ``-5.0``).
        threshold: ``log_alpha`` threshold for declaring a weight pruned.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> layer = DenseVariationalDropout(
        ...     in_features=4, out_features=2, pyrox_name="vd"
        ... )
        >>> x = jnp.ones((3, 4))
        >>> with handlers.seed(rng_seed=0):
        ...     y = layer(x)
        >>> y.shape
        (3, 2)
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    bias: bool = eqx.field(static=True, default=True)
    log_alpha_init: float = eqx.field(static=True, default=-5.0)
    threshold: float = eqx.field(static=True, default=3.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        theta = self.pyrox_param(
            "theta",
            jnp.zeros((self.in_features, self.out_features)),
        )
        log_alpha = self.pyrox_param(
            "log_alpha",
            jnp.full(
                (self.in_features, self.out_features),
                float(self.log_alpha_init),
            ),
        )
        log_alpha_clamped = jnp.clip(log_alpha, _VD_LOG_ALPHA_MIN, _VD_LOG_ALPHA_MAX)
        alpha = jnp.exp(log_alpha_clamped)

        gamma = einx.dot("... din, din dout -> ... dout", x, theta)
        delta = einx.dot("... din, din dout -> ... dout", x**2, alpha * theta**2)
        # Floor to a tiny positive value: keeps the sqrt gradient finite at
        # delta = 0 without injecting visible noise (sqrt(1e-30) ≈ 1e-15).
        std = jnp.sqrt(jnp.maximum(delta, 1e-30))
        # numpyro.prng_key returns Array | None at the type level, but is
        # always an Array inside a `seed` handler — which is required for
        # the pyrox_param/factor calls above to succeed in any case.
        key = cast(JaxArray, numpyro.prng_key())
        eps = jax.random.normal(key, gamma.shape, dtype=gamma.dtype)
        out = gamma + std * eps

        if self.bias:
            b = self.pyrox_param("bias", jnp.zeros(self.out_features))
            out = out + b

        numpyro.factor(
            self._pyrox_fullname("kl"),
            jnp.sum(_vd_neg_kl(log_alpha)),
        )
        return out

    def sparsity(self, log_alpha: Float[Array, "D_in D_out"]) -> Float[Array, ""]:
        """Fraction of weights with ``log_alpha > threshold``.

        Pass the trained ``log_alpha`` parameter, typically retrieved
        from the SVI param store under ``f"{pyrox_name}.log_alpha"``.
        """
        return jnp.mean((log_alpha > self.threshold).astype(log_alpha.dtype))

sparsity(log_alpha: Float[Array, 'D_in D_out']) -> Float[Array, '']

Fraction of weights with log_alpha > threshold.

Pass the trained log_alpha parameter, typically retrieved from the SVI param store under f"{pyrox_name}.log_alpha".

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

def sparsity(self, log_alpha: Float[Array, "D_in D_out"]) -> Float[Array, ""]:
    """Fraction of weights with ``log_alpha > threshold``.

    Pass the trained ``log_alpha`` parameter, typically retrieved
    from the SVI param store under ``f"{pyrox_name}.log_alpha"``.
    """
    return jnp.mean((log_alpha > self.threshold).astype(log_alpha.dtype))

DenseNCP

Bases: PyroxModule

Noise Contrastive Prior dense layer (Hafner et al., 2019).

Decomposes a dense layer into a prior-regularized backbone plus a scaled stochastic perturbation:

\[ y = \underbrace{x W_d + b_d}_{\text{backbone}} + \underbrace{\sigma \cdot (x W_s + b_s)}_{\text{perturbation}}, \]

where all weights are pyrox_sample sites with Gaussian priors and \(\sigma\) has a LogNormal prior. The backbone carries the bulk of the signal; the perturbation branch adds calibrated uncertainty that can be trained via a noise contrastive objective.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
out_features	int	Output dimension.
init_scale	float	Initial value for the perturbation scale \(\sigma\).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class DenseNCP(PyroxModule):
    r"""Noise Contrastive Prior dense layer (Hafner et al., 2019).

    Decomposes a dense layer into a prior-regularized backbone plus a
    scaled stochastic perturbation:

    $$
    y = \underbrace{x W_d + b_d}_{\text{backbone}}
      + \underbrace{\sigma \cdot (x W_s + b_s)}_{\text{perturbation}},
    $$

    where all weights are ``pyrox_sample`` sites with Gaussian priors
    and $\sigma$ has a ``LogNormal`` prior. The backbone carries
    the bulk of the signal; the perturbation branch adds calibrated
    uncertainty that can be trained via a noise contrastive objective.

    Attributes:
        in_features: Input dimension.
        out_features: Output dimension.
        init_scale: Initial value for the perturbation scale
            $\sigma$.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    init_scale: float = eqx.field(static=True, default=1.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        W_d = self.pyrox_sample(
            "weight_det",
            dist.Normal(jnp.zeros((self.in_features, self.out_features)), 1.0).to_event(
                2
            ),
        )
        b_d = self.pyrox_sample(
            "bias_det",
            dist.Normal(jnp.zeros(self.out_features), 1.0).to_event(1),
        )
        det = einx.dot("... din, din dout -> ... dout", x, W_d) + b_d

        W_s = self.pyrox_sample(
            "weight_stoch",
            dist.Normal(jnp.zeros((self.in_features, self.out_features)), 1.0).to_event(
                2
            ),
        )
        b_s = self.pyrox_sample(
            "bias_stoch",
            dist.Normal(jnp.zeros(self.out_features), 1.0).to_event(1),
        )
        scale = self.pyrox_sample(
            "scale",
            dist.LogNormal(jnp.log(jnp.maximum(jnp.array(self.init_scale), 1e-6)), 1.0),
        )
        stoch = scale * (einx.dot("... din, din dout -> ... dout", x, W_s) + b_s)

        return det + stoch

NCPContinuousPerturb

Bases: Module

Input perturbation for the Noise Contrastive Prior pattern.

Adds Gaussian noise scaled by a fixed positive scale to the input:

\[ \tilde{x} = x + \sigma \epsilon, \qquad \epsilon \sim \mathcal{N}(0, I). \]

Place before a deterministic network to inject input uncertainty; pair with a Bayesian DenseNCP head for the full NCP architecture (Hafner et al., 2019).

Stochasticity comes from the explicit PRNG key argument.

Attributes:

Name	Type	Description
scale	float | Float[Array, '']	Perturbation scale \(\sigma\).

Examples:

>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> perturb = NCPContinuousPerturb(scale=0.5)
>>> x = jnp.zeros(3)
>>> out = perturb(x, key=jr.PRNGKey(0))  # x̃ = x + σ·ε, ε ~ N(0, I)
>>> out.shape
(3,)

Source code in .venv/lib/python3.12/site-packages/geonnax/ncp.py

class NCPContinuousPerturb(eqx.Module):
    r"""Input perturbation for the Noise Contrastive Prior pattern.

    Adds Gaussian noise scaled by a fixed positive scale to the input:

    $$
    \tilde{x} = x + \sigma \epsilon, \qquad
    \epsilon \sim \mathcal{N}(0, I).
    $$


    Place before a deterministic network to inject input uncertainty;
    pair with a Bayesian ``DenseNCP`` head for the full NCP
    architecture (Hafner et al., 2019).

    Stochasticity comes from the explicit PRNG ``key`` argument.

    Attributes:
        scale: Perturbation scale $\sigma$.

    Examples:
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> perturb = NCPContinuousPerturb(scale=0.5)
        >>> x = jnp.zeros(3)
        >>> out = perturb(x, key=jr.PRNGKey(0))  # x̃ = x + σ·ε, ε ~ N(0, I)
        >>> out.shape
        (3,)
    """

    scale: float | Float[Array, ""] = 1.0

    def __call__(
        self,
        x: Float[Array, " D"],
        *,
        key: Array,
    ) -> Float[Array, " D"]:
        # x̃ = x + σ·ε with ε ~ N(0, I), same shape/dtype as x → (D,).
        eps = jax.random.normal(key, x.shape, dtype=x.dtype)
        return x + self.scale * eps

NCPNormalOutput

Bases: PyroxModule

Output-side Noise Contrastive Prior layer (Hafner et al., 2018).

Completes the NCP pattern in pyrox_nn: pair with NCPContinuousPerturb at the input and a heteroscedastic network (e.g. an MLP terminating in a mean head and a positive-std head — a softplus or exp of a learned log-scale) so the network produces predictions for both the clean batch and the input-perturbed noisy batch. Given the noisy batch's predictive distribution \(\mathcal{N}(\hat{y}_n, \hat{\sigma}_n^2)\), this layer adds the analytic NCP regulariser

\[ \mathcal{L}_\mathrm{NCP} = \sum_{n} \mathrm{KL}\!\bigl[\mathcal{N}(\hat{y}_n, \hat{\sigma}_n^2) \;\big\|\; \mathcal{N}(\mu_\mathrm{prior}, \sigma_\mathrm{prior}^2)\bigr] \]

to the model log density via numpyro.factor. Pulling the noisy-input predictive distribution toward the fixed prior away from the training distribution gives the network calibrated out-of-distribution uncertainty, which is the central claim of NCP.

The closed-form Gaussian KL used here is

\[ \mathrm{KL}\bigl[\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(\mu_p, \sigma_p^2)\bigr] = \log\frac{\sigma_p}{\sigma} + \frac{\sigma^2 + (\mu - \mu_p)^2}{2\sigma_p^2} - \tfrac{1}{2}. \]

Plate semantics

Unlike pyrox's weight-prior KL terms, the NCP KL is data-dependent — every input row contributes its own \(\mathrm{KL}_n\) term. Internally the layer emits the numpyro.factor site as a per-example vector (shape (*batch,)) rather than a pre-summed scalar; that lets NumPyro's plate machinery sum over the batch axis and apply the subsample scaling automatically.

The canonical training pattern is to emit the layer inside numpyro.plate("data", N, subsample_size=B):

def model(x_clean, y_clean, x_noisy):
    clean_mean, _clean_std = network(x_clean)
    noisy_mean, noisy_std = network(x_noisy)
    ncp_out = NCPNormalOutput(prior_std=1.0)
    with numpyro.plate("data", N, subsample_size=B):
        ncp_out(noisy_mean, noisy_std)              # scaled to N
        numpyro.sample("obs",
            dist.Normal(clean_mean, ...), obs=y_clean)

Inside the plate, NumPyro sums the per-example log-densities over the batch dim and multiplies by scale = N / B, producing the standard unbiased estimate of the full-dataset NCP KL Σ_{n=1}^N KL_n. Outside any plate the layer's contribution is just Σ_{n in batch} KL_n (i.e. the raw batch sum), which is the correct full-dataset value when you train on the whole dataset at once.

Attributes:

Name	Type	Description
prior_mean	float	Prior predictive mean \(\mu_\mathrm{prior}\).
prior_std	float	Prior predictive std \(\sigma_\mathrm{prior}\) (must be positive).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> ncp = NCPNormalOutput(
...     prior_mean=0.0, prior_std=1.0, pyrox_name="ncp_out"
... )
>>> noisy_mean = jnp.zeros((4, 1))
>>> noisy_std = 0.5 * jnp.ones((4, 1))
>>> with handlers.seed(rng_seed=0):
...     kl = ncp(noisy_mean, noisy_std)
>>> kl.shape
()

References

Hafner, D., Tran, D., Lillicrap, T., Irpan, A., & Davidson, J. (2018). Noise Contrastive Priors for Functional Uncertainty. UAI.

Source code in packages/pyrox-nn/src/pyrox_nn/_dense.py

class NCPNormalOutput(PyroxModule):
    r"""Output-side Noise Contrastive Prior layer (Hafner et al., 2018).

    Completes the NCP pattern in ``pyrox_nn``: pair with
    `NCPContinuousPerturb` at the input and a heteroscedastic
    network (e.g. an MLP terminating in a mean head and a positive-std
    head — a softplus or ``exp`` of a learned log-scale) so the
    network produces predictions for both the *clean* batch and the
    input-perturbed *noisy* batch. Given the noisy batch's predictive
    distribution $\mathcal{N}(\hat{y}_n, \hat{\sigma}_n^2)$,
    this layer adds the analytic NCP regulariser

    $$
    \mathcal{L}_\mathrm{NCP} =
    \sum_{n} \mathrm{KL}\!\bigl[\mathcal{N}(\hat{y}_n, \hat{\sigma}_n^2)
        \;\big\|\; \mathcal{N}(\mu_\mathrm{prior}, \sigma_\mathrm{prior}^2)\bigr]
    $$

    to the model log density via `numpyro.factor`. Pulling the
    noisy-input predictive distribution toward the fixed prior away
    from the training distribution gives the network calibrated
    out-of-distribution uncertainty, which is the central claim of NCP.

    The closed-form Gaussian KL used here is

    $$
    \mathrm{KL}\bigl[\mathcal{N}(\mu, \sigma^2)
        \,\|\, \mathcal{N}(\mu_p, \sigma_p^2)\bigr]
    = \log\frac{\sigma_p}{\sigma} +
      \frac{\sigma^2 + (\mu - \mu_p)^2}{2\sigma_p^2} - \tfrac{1}{2}.
    $$

    Plate semantics:
        Unlike pyrox's *weight-prior* KL terms, the NCP KL is
        **data-dependent** — every input row contributes its own
        $\mathrm{KL}_n$ term. Internally the layer emits the
        `numpyro.factor` site as a *per-example* vector
        (shape ``(*batch,)``) rather than a pre-summed scalar; that
        lets NumPyro's plate machinery sum over the batch axis and
        apply the subsample scaling automatically.

        The canonical training pattern is to emit the layer **inside**
        ``numpyro.plate("data", N, subsample_size=B)``:

            def model(x_clean, y_clean, x_noisy):
                clean_mean, _clean_std = network(x_clean)
                noisy_mean, noisy_std = network(x_noisy)
                ncp_out = NCPNormalOutput(prior_std=1.0)
                with numpyro.plate("data", N, subsample_size=B):
                    ncp_out(noisy_mean, noisy_std)              # scaled to N
                    numpyro.sample("obs",
                        dist.Normal(clean_mean, ...), obs=y_clean)

        Inside the plate, NumPyro sums the per-example log-densities
        over the batch dim and multiplies by ``scale = N / B``,
        producing the standard unbiased estimate of the full-dataset
        NCP KL ``Σ_{n=1}^N KL_n``. Outside any plate the layer's
        contribution is just ``Σ_{n in batch} KL_n`` (i.e. the raw
        batch sum), which is the correct full-dataset value when
        you train on the whole dataset at once.

    Attributes:
        prior_mean: Prior predictive mean $\mu_\mathrm{prior}$.
        prior_std: Prior predictive std $\sigma_\mathrm{prior}$
            (must be positive).
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> ncp = NCPNormalOutput(
        ...     prior_mean=0.0, prior_std=1.0, pyrox_name="ncp_out"
        ... )
        >>> noisy_mean = jnp.zeros((4, 1))
        >>> noisy_std = 0.5 * jnp.ones((4, 1))
        >>> with handlers.seed(rng_seed=0):
        ...     kl = ncp(noisy_mean, noisy_std)
        >>> kl.shape
        ()

    References:
        Hafner, D., Tran, D., Lillicrap, T., Irpan, A., & Davidson, J.
        (2018). *Noise Contrastive Priors for Functional Uncertainty.*
        UAI.
    """

    prior_mean: float = eqx.field(static=True, default=0.0)
    prior_std: float = eqx.field(static=True, default=1.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    def __post_init__(self) -> None:
        if self.prior_std <= 0:
            raise ValueError(f"prior_std must be > 0; got {self.prior_std}.")

    @pyrox_method
    def __call__(
        self,
        noisy_mean: Float[Array, "*batch D"],
        noisy_std: Float[Array, "*batch D"],
    ) -> Float[Array, ""]:
        # `noisy_std` is a *standard deviation* — the caller is responsible
        # for ensuring it is non-negative (e.g. via softplus/exp on a
        # learned log-scale head). Negative inputs would silently give a
        # finite-but-wrong KL after squaring; an explicit zero produces
        # `+inf` from `log(0)` which surfaces the bug. We do not floor
        # `noisy_var` because doing so asymmetrically (without a matching
        # floor on `prior_var`) would break the `noisy_std == prior_std`
        # → KL = 0 invariant for tiny `prior_std`.
        if noisy_mean.shape != noisy_std.shape:
            raise ValueError(
                f"noisy_mean shape {noisy_mean.shape} != "
                f"noisy_std shape {noisy_std.shape}."
            )
        # Require an explicit feature axis. With a 1-D ``(B,)`` input,
        # summing along axis=-1 below would collapse the *batch* axis
        # itself, producing a scalar factor that gets broadcast across
        # the data plate — exactly the over-counting bug a per-example
        # factor is designed to avoid. For scalar-regression heads,
        # reshape to ``(B, 1)``.
        if noisy_mean.ndim < 2:
            raise ValueError(
                "noisy_mean / noisy_std must have at least 2 dims "
                "(batch + feature). For a scalar regression head, pass "
                "`noisy_mean[:, None]` and `noisy_std[:, None]`. Got "
                f"shape {noisy_mean.shape}."
            )
        prior_var = jnp.asarray(self.prior_std) ** 2
        noisy_var = noisy_std**2
        kl_per_elem = _diag_gaussian_kl(
            noisy_mean,
            noisy_var,
            jnp.log(noisy_var),
            prior_mean=self.prior_mean,
            log_prior_scale=jnp.log(self.prior_std),
            prior_var=prior_var,
        )
        # Sum only over the trailing feature axis. Keeping the leading
        # batch axis intact is what makes NumPyro's plate machinery do
        # the right thing under `plate("data", N, subsample_size=B)`:
        # the plate handler sums log_probs over the batch dim and then
        # multiplies by `N/B`, giving the unbiased full-dataset estimate
        # `(N/B) * sum_{n in batch} kl_n`. Emitting an already-summed
        # scalar instead would let NumPyro broadcast it across the
        # plate dim and over-count by a factor of B.
        kl_per_example = jnp.sum(kl_per_elem, axis=-1)
        # Add -kl_per_example to the model log density site-by-site.
        # Outside any plate this sums to -total_KL; inside a plate the
        # plate handler scales it correctly.
        numpyro.factor(self._pyrox_fullname("kl"), -kl_per_example)
        return jnp.sum(kl_per_example)

RBFFourierFeatures

Bases: PyroxModule

SSGP-style RFF layer with RBF spectral density.

Both the spectral frequencies \(W\) and the lengthscale \(\ell\) are pyrox_sample sites — \(W\) has a standard normal prior (the RBF spectral density) and \(\ell\) has a LogNormal prior. Under SVI, the guide learns a posterior over both; under a seed handler, they are drawn from the prior.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of frequency pairs (output dim 2 * n_features).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class RBFFourierFeatures(PyroxModule):
    r"""SSGP-style RFF layer with RBF spectral density.

    Both the spectral frequencies $W$ and the lengthscale
    $\ell$ are ``pyrox_sample`` sites — $W$ has a
    standard normal prior (the RBF spectral density) and $\ell$
    has a ``LogNormal`` prior. Under SVI, the guide learns a posterior
    over both; under a seed handler, they are drawn from the prior.

    Attributes:
        in_features: Input dimension.
        n_features: Number of frequency pairs (output dim
            ``2 * n_features``).
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        lengthscale: float = 1.0,
    ) -> RBFFourierFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.Normal(0.0, 1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_forward(W, ls, self.n_features, x)

RBFCosineFeatures

Bases: PyroxModule

Cosine-bias variant of random Fourier features for the RBF kernel.

Uses the single-cosine feature map with a bias term:

\[ \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b) \]

where \(W \sim \mathcal{N}(0, I)\) and \(b \sim \mathrm{Uniform}(0, 2\pi)\). This variant produces n_features-dimensional output (half the dimension of the [cos, sin] variant in RBFFourierFeatures) and is commonly used in Random Kitchen Sinks implementations.

All parameters (\(W\), \(b\), \(\ell\)) are pyrox_sample sites.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of random features (= output dimension).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class RBFCosineFeatures(PyroxModule):
    r"""Cosine-bias variant of random Fourier features for the RBF kernel.

    Uses the single-cosine feature map with a bias term:

    $$
    \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b)
    $$

    where $W \sim \mathcal{N}(0, I)$ and
    $b \sim \mathrm{Uniform}(0, 2\pi)$. This variant produces
    ``n_features``-dimensional output (half the dimension of the
    ``[cos, sin]`` variant in `RBFFourierFeatures`) and is
    commonly used in Random Kitchen Sinks implementations.

    All parameters ($W$, $b$, $\ell$) are
    ``pyrox_sample`` sites.

    Attributes:
        in_features: Input dimension.
        n_features: Number of random features (= output dimension).
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        lengthscale: float = 1.0,
    ) -> RBFCosineFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.Normal(0.0, 1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        b = self.pyrox_sample(
            "b",
            dist.Uniform(0.0, 2.0 * jnp.pi).expand([self.n_features]).to_event(1),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_cosine_forward(W, b, ls, self.n_features, x)

MaternFourierFeatures

Bases: PyroxModule

SSGP-style RFF layer with Matern spectral density.

Spectral frequencies \(W\) have a StudentT(df=2\nu) prior (the Matern spectral density). The smoothness \(\nu\) controls the regularity: nu=0.5 (Laplace), nu=1.5 (Matern-3/2), nu=2.5 (Matern-5/2).

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of frequency pairs.
nu	float	Smoothness parameter \(\nu\).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class MaternFourierFeatures(PyroxModule):
    r"""SSGP-style RFF layer with Matern spectral density.

    Spectral frequencies $W$ have a ``StudentT(df=2\nu)`` prior
    (the Matern spectral density). The smoothness $\nu$ controls
    the regularity: ``nu=0.5`` (Laplace), ``nu=1.5`` (Matern-3/2),
    ``nu=2.5`` (Matern-5/2).

    Attributes:
        in_features: Input dimension.
        n_features: Number of frequency pairs.
        nu: Smoothness parameter $\nu$.
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    nu: float = eqx.field(static=True, default=1.5)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        nu: float = 1.5,
        lengthscale: float = 1.0,
    ) -> MaternFourierFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        if nu <= 0:
            raise ValueError(f"nu must be > 0, got {nu}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            nu=nu,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.StudentT(df=2.0 * self.nu, loc=0.0, scale=1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_forward(W, ls, self.n_features, x)

MaternCosineFeatures

Bases: PyroxModule

Cosine-bias variant of random Fourier features for the Matern kernel.

Single-cosine analogue of MaternFourierFeatures:

\[ \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b) \]

where \(W \sim \mathrm{StudentT}(2\nu)\) (the Matern spectral density) and \(b \sim \mathrm{Uniform}(0, 2\pi)\). Output dim is n_features (vs 2 * n_features for the [cos, sin] variant). Approximates the same kernel as MaternFourierFeatures in expectation but with higher variance per draw — see Sutherland & Schneider (2015).

All parameters (\(W\), \(b\), \(\ell\)) are pyrox_sample sites.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of random features (= output dimension).
nu	float	Smoothness parameter \(\nu\).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class MaternCosineFeatures(PyroxModule):
    r"""Cosine-bias variant of random Fourier features for the Matern kernel.

    Single-cosine analogue of `MaternFourierFeatures`:

    $$
    \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b)
    $$

    where $W \sim \mathrm{StudentT}(2\nu)$ (the Matern spectral
    density) and $b \sim \mathrm{Uniform}(0, 2\pi)$. Output dim is
    ``n_features`` (vs ``2 * n_features`` for the ``[cos, sin]``
    variant). Approximates the same kernel as
    `MaternFourierFeatures` in expectation but with higher
    variance per draw — see Sutherland & Schneider (2015).

    All parameters ($W$, $b$, $\ell$) are
    ``pyrox_sample`` sites.

    Attributes:
        in_features: Input dimension.
        n_features: Number of random features (= output dimension).
        nu: Smoothness parameter $\nu$.
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    nu: float = eqx.field(static=True, default=1.5)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        nu: float = 1.5,
        lengthscale: float = 1.0,
    ) -> MaternCosineFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        if nu <= 0:
            raise ValueError(f"nu must be > 0, got {nu}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            nu=nu,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.StudentT(df=2.0 * self.nu, loc=0.0, scale=1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        b = self.pyrox_sample(
            "b",
            dist.Uniform(0.0, 2.0 * jnp.pi).expand([self.n_features]).to_event(1),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_cosine_forward(W, b, ls, self.n_features, x)

LaplaceFourierFeatures

Bases: PyroxModule

SSGP-style RFF layer with Laplace (Matern-1/2) spectral density.

Spectral frequencies \(W\) have a Cauchy prior (Student-t with df = 1).

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of frequency pairs.
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class LaplaceFourierFeatures(PyroxModule):
    r"""SSGP-style RFF layer with Laplace (Matern-1/2) spectral density.

    Spectral frequencies $W$ have a ``Cauchy`` prior (Student-t
    with ``df = 1``).

    Attributes:
        in_features: Input dimension.
        n_features: Number of frequency pairs.
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        lengthscale: float = 1.0,
    ) -> LaplaceFourierFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.StudentT(df=1.0, loc=0.0, scale=1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_forward(W, ls, self.n_features, x)

LaplaceCosineFeatures

Bases: PyroxModule

Cosine-bias variant of random Fourier features for the Laplace kernel.

Single-cosine analogue of LaplaceFourierFeatures (the Matern-1/2 kernel):

\[ \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b) \]

where \(W \sim \mathrm{Cauchy}(0, 1)\) (Student-t with df = 1) and \(b \sim \mathrm{Uniform}(0, 2\pi)\). Output dim is n_features.

All parameters (\(W\), \(b\), \(\ell\)) are pyrox_sample sites.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of random features (= output dimension).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class LaplaceCosineFeatures(PyroxModule):
    r"""Cosine-bias variant of random Fourier features for the Laplace kernel.

    Single-cosine analogue of `LaplaceFourierFeatures` (the
    Matern-1/2 kernel):

    $$
    \phi(x) = \sqrt{2 / D}\,\cos(x W / \ell + b)
    $$

    where $W \sim \mathrm{Cauchy}(0, 1)$ (Student-t with
    ``df = 1``) and $b \sim \mathrm{Uniform}(0, 2\pi)$. Output
    dim is ``n_features``.

    All parameters ($W$, $b$, $\ell$) are
    ``pyrox_sample`` sites.

    Attributes:
        in_features: Input dimension.
        n_features: Number of random features (= output dimension).
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        lengthscale: float = 1.0,
    ) -> LaplaceCosineFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.StudentT(df=1.0, loc=0.0, scale=1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        b = self.pyrox_sample(
            "b",
            dist.Uniform(0.0, 2.0 * jnp.pi).expand([self.n_features]).to_event(1),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_cosine_forward(W, b, ls, self.n_features, x)

ArcCosineFourierFeatures

Bases: PyroxModule

Random features for the arc-cosine kernel (Cho & Saul, 2009).

The arc-cosine kernel of order \(p\) corresponds to an infinite-width single-layer ReLU network. The random feature map is:

\[ \phi(x) = \sqrt{2 / D}\,\max(0,\, x W / \ell)^p \]

where \(W \sim \mathcal{N}(0, I)\).

order=0 gives the Heaviside (step) feature; order=1 gives the ReLU feature (the most common); order=2 gives the squared ReLU feature.

Attributes:

Name	Type	Description
in_features	int	Input dimension.
n_features	int	Number of random features (= output dimension).
order	int	Kernel order (0, 1, or 2).
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class ArcCosineFourierFeatures(PyroxModule):
    r"""Random features for the arc-cosine kernel (Cho & Saul, 2009).

    The arc-cosine kernel of order $p$ corresponds to an
    infinite-width single-layer ReLU network. The random feature map
    is:

    $$
    \phi(x) = \sqrt{2 / D}\,\max(0,\, x W / \ell)^p
    $$

    where $W \sim \mathcal{N}(0, I)$.

    ``order=0`` gives the Heaviside (step) feature; ``order=1`` gives
    the ReLU feature (the most common); ``order=2`` gives the squared
    ReLU feature.

    Attributes:
        in_features: Input dimension.
        n_features: Number of random features (= output dimension).
        order: Kernel order (0, 1, or 2).
        init_lengthscale: Prior location for the lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    order: int = eqx.field(static=True, default=1)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        order: int = 1,
        lengthscale: float = 1.0,
    ) -> ArcCosineFourierFeatures:
        return cls(
            in_features=in_features,
            n_features=n_features,
            order=order,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        W = self.pyrox_sample(
            "W",
            dist.Normal(0.0, 1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        z = einx.dot("... din, din f -> ... f", x, W) / ls
        if self.order == 0:
            h = (z > 0.0).astype(x.dtype)
        else:
            h = jnp.maximum(z, 0.0) ** self.order
        return jnp.sqrt(2.0 / self.n_features) * h

RandomKitchenSinks

Bases: PyroxModule

Random Kitchen Sinks: RFF + a learned linear head.

Composes any RFF layer (RBFFourierFeatures, MaternFourierFeatures, LaplaceFourierFeatures) with a trainable linear projection:

\[ y = \phi(x)\, \beta + b \]

The linear head (beta, bias) is registered via pyrox_sample with Normal priors.

Attributes:

Name	Type	Description
rff	RBFFourierFeatures | MaternFourierFeatures | LaplaceFourierFeatures	The underlying RFF feature layer.
init_beta	Float[Array, 'D_rff D_out']	Initial linear weights.
init_bias	Float[Array, ' D_out']	Initial bias vector.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class RandomKitchenSinks(PyroxModule):
    r"""Random Kitchen Sinks: RFF + a learned linear head.

    Composes any RFF layer (`RBFFourierFeatures`,
    `MaternFourierFeatures`, `LaplaceFourierFeatures`)
    with a trainable linear projection:

    $$
    y = \phi(x)\, \beta + b
    $$

    The linear head (``beta``, ``bias``) is registered via
    ``pyrox_sample`` with ``Normal`` priors.

    Attributes:
        rff: The underlying RFF feature layer.
        init_beta: Initial linear weights.
        init_bias: Initial bias vector.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    rff: RBFFourierFeatures | MaternFourierFeatures | LaplaceFourierFeatures
    init_beta: Float[Array, "D_rff D_out"]
    init_bias: Float[Array, " D_out"]
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        rff: RBFFourierFeatures | MaternFourierFeatures | LaplaceFourierFeatures,
        out_features: int,
    ) -> RandomKitchenSinks:
        """Construct from a pre-built RFF layer with zero-initialized head."""
        beta = jnp.zeros((2 * rff.n_features, out_features))
        bias = jnp.zeros(out_features)
        return cls(rff=rff, init_beta=beta, init_bias=bias)

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        phi = self.rff(x)
        beta = self.pyrox_sample(
            "beta",
            dist.Normal(self.init_beta, 1.0).to_event(2),
        )
        bias = self.pyrox_sample(
            "bias",
            dist.Normal(self.init_bias, 1.0).to_event(1),
        )
        return einx.dot("... r, r dout -> ... dout", phi, beta) + bias

init(rff: RBFFourierFeatures | MaternFourierFeatures | LaplaceFourierFeatures, out_features: int) -> RandomKitchenSinks classmethod

Construct from a pre-built RFF layer with zero-initialized head.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

@classmethod
def init(
    cls,
    rff: RBFFourierFeatures | MaternFourierFeatures | LaplaceFourierFeatures,
    out_features: int,
) -> RandomKitchenSinks:
    """Construct from a pre-built RFF layer with zero-initialized head."""
    beta = jnp.zeros((2 * rff.n_features, out_features))
    bias = jnp.zeros(out_features)
    return cls(rff=rff, init_beta=beta, init_bias=bias)

Wave-4 spectral layers (#41)

VariationalFourierFeatures

Bases: PyroxModule

VSSGP — RFF with a learnable variational posterior over frequencies.

Standard RFF (e.g. RBFFourierFeatures) treats the spectral frequencies \(W\) as a frozen prior draw; VSSGP (Gal & Turner, 2015) treats \(W\) as a latent with a learnable mean-field posterior, recovering spectral uncertainty on top of the feature-space uncertainty.

Prior: \(p(W) = \mathcal{N}(0, I)\) (RBF spectral density in lengthscale-1 units). The lengthscale is itself a sampled site (LogNormal(log init_lengthscale, 1)) so that frequencies are rescaled to the physical kernel.

Under SVI, attach an AutoNormal to learn the posterior on W; under prior-only seeds, behaves identically to RBFFourierFeatures.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D\).
n_features	int	Number of frequency pairs (output dim 2 * n_features).
init_lengthscale	float	Prior location for the kernel lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class VariationalFourierFeatures(PyroxModule):
    r"""VSSGP — RFF with a learnable variational posterior over frequencies.

    Standard RFF (e.g. `RBFFourierFeatures`) treats the spectral
    frequencies $W$ as a frozen prior draw; VSSGP (Gal & Turner,
    2015) treats $W$ as a latent with a learnable mean-field
    posterior, recovering spectral *uncertainty* on top of the
    feature-space uncertainty.

    Prior: $p(W) = \mathcal{N}(0, I)$ (RBF spectral density in
    lengthscale-1 units). The lengthscale is itself a sampled site
    (``LogNormal(log init_lengthscale, 1)``) so that frequencies are
    rescaled to the physical kernel.

    Under SVI, attach an `AutoNormal` to
    learn the posterior on ``W``; under prior-only seeds, behaves
    identically to `RBFFourierFeatures`.

    Attributes:
        in_features: Input dimension $D$.
        n_features: Number of frequency pairs (output dim ``2 * n_features``).
        init_lengthscale: Prior location for the kernel lengthscale.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    init_lengthscale: float = 1.0
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        lengthscale: float = 1.0,
    ) -> VariationalFourierFeatures:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        return cls(
            in_features=in_features,
            n_features=n_features,
            init_lengthscale=lengthscale,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        # Same prior as RBFFourierFeatures — the *posterior* is what differs
        # under SVI: an attached AutoGuide learns q(W) instead of forcing W
        # to its prior draw.
        W = self.pyrox_sample(
            "W",
            dist.Normal(0.0, 1.0)
            .expand([self.in_features, self.n_features])
            .to_event(2),
        )
        ls = self.pyrox_sample(
            "lengthscale",
            dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
        )
        return _vmap_rff_forward(W, ls, self.n_features, x)

OrthogonalRandomFeatures

Bases: Module

Orthogonal Random Features (Yu et al., 2016) — variance-reduced RFF.

Frequencies are drawn from blocks of Haar-orthogonal matrices scaled by independent chi-distributed magnitudes, giving the same RBF kernel approximation as plain RBFFourierFeatures in expectation but with provably lower variance for finite n_features.

Frozen at construction time — no priors, no SVI on W. The frequency matrix is built once from a key and stored as a static array.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D\).
n_features	int	Number of feature pairs. Must satisfy n_features % in_features == 0 so that ORF blocks tile cleanly.
lengthscale	Float[Array, '']	Fixed kernel lengthscale (no prior; pass a value).
W	Float[Array, 'D_in D_orf']	Pre-built frequency matrix of shape (in_features, n_features).

The feature map is the shared RFF map \(\phi(x) = \sqrt{1/D}\,[\cos(W^\top x/\ell),\,\sin(W^\top x/\ell)]\), so calling the module on a (in_features,) vector yields a (2 * n_features,) feature vector.

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from geonnax.randfeat import OrthogonalRandomFeatures
>>> orf = OrthogonalRandomFeatures.init(4, 8, key=jr.PRNGKey(0))
>>> orf(jnp.ones(4)).shape
(16,)

Source code in .venv/lib/python3.12/site-packages/geonnax/randfeat.py

class OrthogonalRandomFeatures(eqx.Module):
    r"""Orthogonal Random Features (Yu et al., 2016) — variance-reduced RFF.

    Frequencies are drawn from blocks of Haar-orthogonal matrices scaled by
    independent chi-distributed magnitudes, giving the same RBF kernel
    approximation as plain ``RBFFourierFeatures`` *in expectation* but
    with provably lower variance for finite ``n_features``.

    Frozen at construction time — no priors, no SVI on ``W``. The frequency
    matrix is built once from a ``key`` and stored as a static array.

    Attributes:
        in_features: Input dimension $D$.
        n_features: Number of feature pairs. Must satisfy
            ``n_features % in_features == 0`` so that ORF blocks tile cleanly.
        lengthscale: Fixed kernel lengthscale (no prior; pass a value).
        W: Pre-built frequency matrix of shape ``(in_features, n_features)``.

    The feature map is the shared RFF map
    $\phi(x) = \sqrt{1/D}\,[\cos(W^\top x/\ell),\,\sin(W^\top x/\ell)]$,
    so calling the module on a ``(in_features,)`` vector yields a
    ``(2 * n_features,)`` feature vector.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from geonnax.randfeat import OrthogonalRandomFeatures
        >>> orf = OrthogonalRandomFeatures.init(4, 8, key=jr.PRNGKey(0))
        >>> orf(jnp.ones(4)).shape
        (16,)
    """

    in_features: int = eqx.field(static=True)
    n_features: int = eqx.field(static=True)
    lengthscale: Float[Array, ""]
    W: Float[Array, "D_in D_orf"]

    @classmethod
    def init(
        cls,
        in_features: int,
        n_features: int,
        *,
        key: jax.Array,
        lengthscale: float = 1.0,
    ) -> OrthogonalRandomFeatures:
        """Build the frozen ORF frequency matrix and wrap it in the module.

        ``n_features`` must be divisible by ``in_features`` so the
        Haar-orthogonal blocks tile cleanly.

        Examples:
            >>> import jax.random as jr
            >>> from geonnax.randfeat import OrthogonalRandomFeatures
            >>> orf = OrthogonalRandomFeatures.init(4, 8, key=jr.PRNGKey(0))
            >>> orf.W.shape
            (4, 8)
        """
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        if in_features <= 0 or n_features <= 0:
            raise ValueError(
                "in_features and n_features must be > 0; got "
                f"in_features={in_features}, n_features={n_features}."
            )
        if n_features % in_features != 0:
            raise ValueError(
                f"n_features ({n_features}) must be divisible by in_features "
                f"({in_features}) so ORF blocks tile cleanly."
            )
        n_blocks = n_features // in_features
        W = orthogonal_blocks(in_features, n_blocks, key=key)
        return cls(
            in_features=in_features,
            n_features=n_features,
            lengthscale=jnp.asarray(lengthscale),
            W=W,
        )

    def __call__(self, x: Float[Array, " D_in"]) -> Float[Array, " D_rff"]:
        r"""Map ``x`` to its ORF feature vector ``(2 * n_features,)``.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.randfeat import OrthogonalRandomFeatures
            >>> orf = OrthogonalRandomFeatures.init(2, 4, key=jr.PRNGKey(0))
            >>> orf(jnp.ones(2)).shape
            (8,)
        """
        # (D_in,) -> (2 * n_features,)
        return rff_forward(self.W, self.lengthscale, self.n_features, x)

init(in_features: int, n_features: int, *, key: jax.Array, lengthscale: float = 1.0) -> OrthogonalRandomFeatures classmethod

Build the frozen ORF frequency matrix and wrap it in the module.

n_features must be divisible by in_features so the Haar-orthogonal blocks tile cleanly.

Examples:

>>> import jax.random as jr
>>> from geonnax.randfeat import OrthogonalRandomFeatures
>>> orf = OrthogonalRandomFeatures.init(4, 8, key=jr.PRNGKey(0))
>>> orf.W.shape
(4, 8)

Source code in .venv/lib/python3.12/site-packages/geonnax/randfeat.py

@classmethod
def init(
    cls,
    in_features: int,
    n_features: int,
    *,
    key: jax.Array,
    lengthscale: float = 1.0,
) -> OrthogonalRandomFeatures:
    """Build the frozen ORF frequency matrix and wrap it in the module.

    ``n_features`` must be divisible by ``in_features`` so the
    Haar-orthogonal blocks tile cleanly.

    Examples:
        >>> import jax.random as jr
        >>> from geonnax.randfeat import OrthogonalRandomFeatures
        >>> orf = OrthogonalRandomFeatures.init(4, 8, key=jr.PRNGKey(0))
        >>> orf.W.shape
        (4, 8)
    """
    if lengthscale <= 0:
        raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
    if in_features <= 0 or n_features <= 0:
        raise ValueError(
            "in_features and n_features must be > 0; got "
            f"in_features={in_features}, n_features={n_features}."
        )
    if n_features % in_features != 0:
        raise ValueError(
            f"n_features ({n_features}) must be divisible by in_features "
            f"({in_features}) so ORF blocks tile cleanly."
        )
    n_blocks = n_features // in_features
    W = orthogonal_blocks(in_features, n_blocks, key=key)
    return cls(
        in_features=in_features,
        n_features=n_features,
        lengthscale=jnp.asarray(lengthscale),
        W=W,
    )

HSGPFeatures

Bases: PyroxModule

Hilbert-Space Gaussian Process feature layer (Riutort-Mayol et al., 2023).

A deterministic Laplacian-eigenfunction basis on the bounded box \([-L, L]^D\) plus learnable per-basis amplitudes with a kernel-spectral-density prior:

\[ \hat{f}(x) = \sum_{j=1}^{M} \alpha_j\,\sqrt{S(\sqrt{\lambda_j})}\,\phi_j(x), \quad \alpha_j \sim \mathcal{N}(0, 1). \]

This is the NN-side dual of pyrox_gp.FourierInducingFeatures — same basis, different prior wiring. As M and L grow, the induced GP converges to the kernel passed in.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D\).
num_basis_per_dim	tuple[int, ...]	Per-axis number of 1D eigenfunctions; total basis count is prod(num_basis_per_dim).
L	tuple[float, ...]	Per-axis box half-width.
kernel	Kernel	A stationary kernel from pyrox_gp whose spectral density supplies the per-basis prior variance. Currently pyrox_gp.RBF and pyrox_gp.Matern are supported by pyrox_gp._basis.spectral_density.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class HSGPFeatures(PyroxModule):
    r"""Hilbert-Space Gaussian Process feature layer (Riutort-Mayol et al., 2023).

    A *deterministic* Laplacian-eigenfunction basis on the bounded box
    $[-L, L]^D$ plus learnable per-basis amplitudes with a
    kernel-spectral-density prior:

    $$
    \hat{f}(x) = \sum_{j=1}^{M} \alpha_j\,\sqrt{S(\sqrt{\lambda_j})}\,\phi_j(x),
    \quad \alpha_j \sim \mathcal{N}(0, 1).
    $$

    This is the NN-side dual of `pyrox_gp.FourierInducingFeatures`
    — same basis, different prior wiring. As ``M`` and ``L`` grow, the
    induced GP converges to the kernel passed in.

    Attributes:
        in_features: Input dimension $D$.
        num_basis_per_dim: Per-axis number of 1D eigenfunctions; total
            basis count is ``prod(num_basis_per_dim)``.
        L: Per-axis box half-width.
        kernel: A stationary kernel from `pyrox_gp` whose spectral
            density supplies the per-basis prior variance. Currently
            `pyrox_gp.RBF` and `pyrox_gp.Matern` are
            supported by `pyrox_gp._basis.spectral_density`.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int = eqx.field(static=True)
    num_basis_per_dim: tuple[int, ...] = eqx.field(static=True)
    L: tuple[float, ...] = eqx.field(static=True)
    kernel: Kernel
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        in_features: int,
        num_basis_per_dim: int | tuple[int, ...],
        L: float | tuple[float, ...],
        *,
        kernel: Kernel,
    ) -> HSGPFeatures:
        if isinstance(num_basis_per_dim, int):
            num_basis_per_dim = (num_basis_per_dim,) * in_features
        if isinstance(L, int | float):
            L = (float(L),) * in_features
        if len(num_basis_per_dim) != in_features:
            raise ValueError(
                f"num_basis_per_dim length ({len(num_basis_per_dim)}) "
                f"must match in_features ({in_features})."
            )
        if len(L) != in_features:
            raise ValueError(
                f"L length ({len(L)}) must match in_features ({in_features})."
            )
        if any(L_d <= 0 for L_d in L):
            raise ValueError(f"L must be all positive; got {L}.")
        if any(M_d < 1 for M_d in num_basis_per_dim):
            raise ValueError(
                f"num_basis_per_dim must be all >= 1; got {num_basis_per_dim}."
            )
        return cls(
            in_features=in_features,
            num_basis_per_dim=tuple(num_basis_per_dim),
            L=tuple(float(L_d) for L_d in L),
            kernel=kernel,
        )

    @property
    def num_basis(self) -> int:
        n = 1
        for m in self.num_basis_per_dim:
            n *= m
        return n

    @pyrox_method
    def __call__(self, x: Float[Array, "N D_in"]) -> Float[Array, " N"]:
        Phi, lam = fourier_basis(x, self.num_basis_per_dim, self.L)  # (N, M), (M,)
        # Spectral density evaluated under the kernel's own context so any
        # priors on (variance, lengthscale) register exactly once.
        with _kernel_context(self.kernel):
            S = spectral_density(self.kernel, lam, D=self.in_features)
        sqrt_S = jnp.sqrt(S)
        alpha = self.pyrox_sample(
            "alpha",
            dist.Normal(0.0, 1.0).expand([self.num_basis]).to_event(1),
        )
        return einx.dot("n m, m -> n", Phi, sqrt_S * alpha)

Spherical-harmonic GP features

SphericalHarmonicFeatures is the \(S^2\) analogue of HSGPFeatures: real spherical harmonics up to degree \(L\) with \(\mathcal{N}(0, a_l)\) weights, where \(a_l\) are the kernel's Funk-Hecke coefficients. With fixed kernel hyperparameters it is a GP with the truncated kernel, linear in the number of observations, so it is an optimal-interpolation baseline that runs through the same inference code as a neural field. The sketch below fits it and a BayesianNeuralField on the same synthetic global field with the same pyrox.inference.ensemble_map loop, so the comparison isolates the network.

import jax
import jax.numpy as jnp
import jax.random as jr
import numpyro.distributions as dist
from numpyro import handlers
from pyrox.inference import ensemble_map
from pyrox_gp import Matern
from pyrox_nn import BayesianNeuralField, SphericalHarmonicFeatures

# Synthetic global field: 2,000 noisy observations at random lon / lat (degrees).
k_lon, k_lat, k_eps = jr.split(jr.key(0), 3)
lon = jr.uniform(k_lon, (2000,), minval=-180.0, maxval=180.0)
lat = jnp.rad2deg(jnp.arcsin(jr.uniform(k_lat, (2000,), minval=-1.0, maxval=1.0)))
x = jnp.stack([lon, lat], axis=-1)  # (N, 2)
lam, phi = jnp.deg2rad(lon), jnp.deg2rad(lat)
y = jnp.sin(2 * lam) * jnp.cos(phi) + 0.5 * jnp.sin(phi) ** 2
y = y + 0.1 * jr.normal(k_eps, y.shape)


def fit(layer, sigma=0.1):
    """MAP-fit every sample site of ``layer`` under a Normal likelihood."""

    def log_joint(params, xb, yb):
        with handlers.trace() as tr, handlers.substitute(data=params):
            f = layer(xb)
        logprior = sum(
            s["fn"].log_prob(s["value"]).sum()
            for s in tr.values()
            if s["type"] == "sample"
        )
        return dist.Normal(f, sigma).log_prob(yb).sum(), logprior

    def init_fn(key):
        with handlers.trace() as tr, handlers.seed(rng_seed=key):
            layer(x[:1])
        return {n: s["value"] for n, s in tr.items() if s["type"] == "sample"}

    params, _ = ensemble_map(
        log_joint, init_fn, ensemble_size=4, num_epochs=2000,
        data=(x, y), seed=jr.key(1),
    )
    # Ensemble-mean prediction at the training points.
    return jax.vmap(lambda p: handlers.substitute(layer, data=p)(x))(params).mean(0)


gp = SphericalHarmonicFeatures.init(
    l_max=30,
    kernel=Matern(nu=1.5, init_lengthscale=0.5),  # fixed hyperparameters
    input_mode="lonlat",
    input_unit="degrees",
)
bnf = BayesianNeuralField(
    input_scales=(180.0, 90.0),
    fourier_degrees=(4, 4),
    interactions=((0, 1),),
    seasonality_periods=(),
    num_seasonal_harmonics=(),
    width=64,
    depth=2,
)
for name, layer in [("SH-GP", gp), ("BNF", bnf)]:
    rmse = jnp.sqrt(jnp.mean((fit(layer) - y) ** 2))
    print(f"{name}: train RMSE {rmse:.3f}")

The kernel's hyperparameters are param sites here, so they stay fixed; give the kernel priors (kernel.set_prior(...)) to fit them as sample sites too.

SphericalHarmonicFeatures

Bases: PyroxModule

Weight-space GP on the sphere \(S^2\) from real spherical harmonics.

The \(S^2\) analogue of HSGPFeatures. For a zonal kernel \(k(n, n') = \kappa(n \cdot n')\) with Funk-Hecke coefficients \(a_l = 2\pi \int_{-1}^{1} \kappa(t) P_l(t)\,dt\), the addition theorem gives the Mercer expansion

\[ k(n, n') = \sum_{l=0}^{\infty} a_l \sum_{m=-l}^{l} Y_{lm}(n)\,Y_{lm}(n'), \]

with \(Y_{lm}\) the orthonormal real spherical harmonics. Truncating at degree \(L\) = l_max gives the layer

\[ f(n) = \sum_{l=0}^{L} \sum_{m=-l}^{l} \sqrt{a_l}\,\alpha_{lm}\,Y_{lm}(n), \qquad \alpha_{lm} \sim \mathcal{N}(0, 1), \]

a GP whose covariance is the truncated kernel \(k_L(n, n') = \sum_{l \le L} a_l \sum_m Y_{lm}(n) Y_{lm}(n')\), linear in the number of points. This is the NN-side dual of pyrox_gp.SphericalHarmonicInducingFeatures (VISH, Dutordoir et al., 2020): same harmonics, same pyrox_gp.funk_hecke_coefficients.

The kernel is any isotropic pyrox_gp kernel, evaluated on the embedded unit vectors (chordal distance), as in VISH. A kernel with a per-dimension (ARD) lengthscale is not zonal and raises NotImplementedError. Coefficients that quadrature round-off pushes below zero are clipped to zero (a kernel restricted from \(\mathbb{R}^3\) to the sphere has \(a_l \ge 0\)).

Sites. <pyrox_name>.alpha with prior \(\mathcal{N}(0, I_{(L+1)^2})\); any priors on the kernel's hyperparameters register once per call under the kernel's own scope (evaluated inside pyrox_gp._context._kernel_context).

Attributes:

Name	Type	Description
l_max	int	Maximum harmonic degree \(L\) (inclusive); the layer has (l_max + 1) ** 2 features.
kernel	Kernel	An isotropic pyrox_gp kernel (e.g. pyrox_gp.Matern, pyrox_gp.RBF) built without input_dim.
input_mode	Literal['xyz', 'lonlat']	"xyz" for unit vectors (N, 3), or "lonlat" for (N, 2) longitude / latitude columns.
input_unit	Literal['degrees', 'radians'] | None	Unit of the "lonlat" columns, "degrees" or "radians"; None for "xyz".
num_quadrature	int	Gauss-Legendre nodes for the Funk-Hecke integral.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> from pyrox_gp import Matern
>>> layer = SphericalHarmonicFeatures.init(
...     l_max=10,
...     kernel=Matern(nu=1.5, init_lengthscale=0.5),
...     input_mode="lonlat",
...     input_unit="degrees",
... )
>>> lonlat = jnp.array([[0.0, 0.0], [90.0, 45.0], [-120.0, -60.0]])
>>> with handlers.seed(rng_seed=0):
...     f = layer(lonlat)
>>> f.shape, layer.num_basis
((3,), 121)

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

class SphericalHarmonicFeatures(PyroxModule):
    r"""Weight-space GP on the sphere $S^2$ from real spherical harmonics.

    The $S^2$ analogue of `HSGPFeatures`. For a zonal kernel
    $k(n, n') = \kappa(n \cdot n')$ with Funk-Hecke coefficients
    $a_l = 2\pi \int_{-1}^{1} \kappa(t) P_l(t)\,dt$, the addition theorem
    gives the Mercer expansion

    $$
    k(n, n') = \sum_{l=0}^{\infty} a_l \sum_{m=-l}^{l} Y_{lm}(n)\,Y_{lm}(n'),
    $$

    with $Y_{lm}$ the orthonormal real spherical harmonics. Truncating at
    degree $L$ = ``l_max`` gives the layer

    $$
    f(n) = \sum_{l=0}^{L} \sum_{m=-l}^{l} \sqrt{a_l}\,\alpha_{lm}\,Y_{lm}(n),
    \qquad \alpha_{lm} \sim \mathcal{N}(0, 1),
    $$

    a GP whose covariance is the truncated kernel
    $k_L(n, n') = \sum_{l \le L} a_l \sum_m Y_{lm}(n) Y_{lm}(n')$, linear in
    the number of points. This is the NN-side dual of
    `pyrox_gp.SphericalHarmonicInducingFeatures` (VISH, Dutordoir et al.,
    2020): same harmonics, same `pyrox_gp.funk_hecke_coefficients`.

    The kernel is any isotropic `pyrox_gp` kernel, evaluated on the
    embedded unit vectors (chordal distance), as in VISH. A kernel with a
    per-dimension (ARD) lengthscale is not zonal and raises
    ``NotImplementedError``. Coefficients that quadrature round-off pushes
    below zero are clipped to zero (a kernel restricted from
    $\mathbb{R}^3$ to the sphere has $a_l \ge 0$).

    **Sites.** ``<pyrox_name>.alpha`` with prior
    $\mathcal{N}(0, I_{(L+1)^2})$; any priors on the kernel's
    hyperparameters register once per call under the kernel's own scope
    (evaluated inside `pyrox_gp._context._kernel_context`).

    Attributes:
        l_max: Maximum harmonic degree $L$ (inclusive); the layer has
            ``(l_max + 1) ** 2`` features.
        kernel: An isotropic `pyrox_gp` kernel (e.g. `pyrox_gp.Matern`,
            `pyrox_gp.RBF`) built without ``input_dim``.
        input_mode: ``"xyz"`` for unit vectors ``(N, 3)``, or ``"lonlat"``
            for ``(N, 2)`` longitude / latitude columns.
        input_unit: Unit of the ``"lonlat"`` columns, ``"degrees"`` or
            ``"radians"``; ``None`` for ``"xyz"``.
        num_quadrature: Gauss-Legendre nodes for the Funk-Hecke integral.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> from pyrox_gp import Matern
        >>> layer = SphericalHarmonicFeatures.init(
        ...     l_max=10,
        ...     kernel=Matern(nu=1.5, init_lengthscale=0.5),
        ...     input_mode="lonlat",
        ...     input_unit="degrees",
        ... )
        >>> lonlat = jnp.array([[0.0, 0.0], [90.0, 45.0], [-120.0, -60.0]])
        >>> with handlers.seed(rng_seed=0):
        ...     f = layer(lonlat)
        >>> f.shape, layer.num_basis
        ((3,), 121)
    """

    kernel: Kernel
    l_max: int = eqx.field(static=True)
    input_mode: Literal["xyz", "lonlat"] = eqx.field(static=True, default="xyz")
    input_unit: Literal["degrees", "radians"] | None = eqx.field(
        static=True, default=None
    )
    num_quadrature: int = eqx.field(static=True, default=256)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @classmethod
    def init(
        cls,
        *,
        l_max: int,
        kernel: Kernel,
        input_mode: Literal["xyz", "lonlat"] = "xyz",
        input_unit: Literal["degrees", "radians"] | None = None,
        num_quadrature: int = 256,
        pyrox_name: str | None = None,
    ) -> SphericalHarmonicFeatures:
        """Validate the configuration and build the layer.

        Raises:
            ValueError: ``l_max < 0``, ``num_quadrature < 1``, an unknown
                ``input_mode``, or an ``input_unit`` that is missing for
                ``"lonlat"``, unknown, or given for ``"xyz"``.
        """
        if l_max < 0:
            raise ValueError(f"l_max must be >= 0; got {l_max}.")
        if num_quadrature < 1:
            raise ValueError(f"num_quadrature must be >= 1; got {num_quadrature}.")
        if input_mode not in ("xyz", "lonlat"):
            raise ValueError(
                f"input_mode must be 'xyz' or 'lonlat'; got {input_mode!r}."
            )
        if input_mode == "lonlat" and input_unit not in ("degrees", "radians"):
            raise ValueError(
                "input_mode='lonlat' needs an explicit input_unit, 'degrees' or "
                f"'radians'; got {input_unit!r}."
            )
        if input_mode == "xyz" and input_unit is not None:
            raise ValueError(
                f"input_unit applies to input_mode='lonlat' only; got {input_unit!r}."
            )
        return cls(
            kernel=kernel,
            l_max=l_max,
            input_mode=input_mode,
            input_unit=input_unit,
            num_quadrature=num_quadrature,
            pyrox_name=pyrox_name,
        )

    @property
    def num_basis(self) -> int:
        """Number of harmonics, ``(l_max + 1) ** 2``."""
        return (self.l_max + 1) ** 2

    @pyrox_method
    def __call__(self, x: Float[Array, "N D_in"]) -> Float[Array, " N"]:
        """Evaluate one GP draw at ``x``.

        Args:
            x: Unit vectors ``(N, 3)`` (``input_mode="xyz"``) or lon / lat
                ``(N, 2)`` in ``input_unit`` (``input_mode="lonlat"``).

        Returns:
            The latent function values, shape ``(N,)``.
        """
        d_in = 3 if self.input_mode == "xyz" else 2
        if x.ndim != 2 or x.shape[-1] != d_in:
            raise ValueError(
                f"input_mode={self.input_mode!r} expects x of shape (N, {d_in}); "
                f"got {x.shape}."
            )
        if self.input_mode == "lonlat":
            if self.input_unit is None:  # built without init(): no default unit
                raise ValueError(
                    "input_mode='lonlat' needs an explicit input_unit, "
                    "'degrees' or 'radians'."
                )
            x = lonlat_to_cartesian3d(x, input_unit=self.input_unit)  # (N, 3)
        Y = real_spherical_harmonics(x, self.l_max)  # (N, M)
        # Funk-Hecke coefficients; the helper rejects ARD lengthscales. It
        # opens its own kernel context, so this outer one is a no-op today
        # (contexts are reentrant); it is kept, as in HSGPFeatures, so that
        # any kernel call added here shares the same hyperprior draw.
        with _kernel_context(self.kernel):
            a = funk_hecke_coefficients(
                self.kernel, self.l_max, num_quadrature=self.num_quadrature
            )  # (l_max + 1,)
        # a_l >= 0 for a kernel restricted from R^3; quadrature round-off can
        # push tail coefficients below 0 (often in float32). Clip them with a
        # double where: sqrt(maximum(a, 0)) has a 0 * inf = NaN gradient there.
        pos = a > 0
        sqrt_a_l = jnp.where(pos, jnp.sqrt(jnp.where(pos, a, 1.0)), 0.0)
        degrees = jnp.asarray(harmonic_degrees(self.l_max))  # (M,) degree l per column
        sqrt_a = sqrt_a_l[degrees]  # (M,)
        alpha = self.pyrox_sample(
            "alpha",
            dist.Normal(0.0, 1.0).expand([self.num_basis]).to_event(1),
        )
        # The quadrature runs in the default float dtype; keep the input's.
        return einx.dot("n m, m -> n", Y, (sqrt_a * alpha).astype(Y.dtype))

num_basis: int property

Number of harmonics, (l_max + 1) ** 2.

init(*, l_max: int, kernel: Kernel, input_mode: Literal['xyz', 'lonlat'] = 'xyz', input_unit: Literal['degrees', 'radians'] | None = None, num_quadrature: int = 256, pyrox_name: str | None = None) -> SphericalHarmonicFeatures classmethod

Validate the configuration and build the layer.

Raises:

Type	Description
ValueError	l_max < 0, num_quadrature < 1, an unknown input_mode, or an input_unit that is missing for "lonlat", unknown, or given for "xyz".

Source code in packages/pyrox-nn/src/pyrox_nn/_features.py

@classmethod
def init(
    cls,
    *,
    l_max: int,
    kernel: Kernel,
    input_mode: Literal["xyz", "lonlat"] = "xyz",
    input_unit: Literal["degrees", "radians"] | None = None,
    num_quadrature: int = 256,
    pyrox_name: str | None = None,
) -> SphericalHarmonicFeatures:
    """Validate the configuration and build the layer.

    Raises:
        ValueError: ``l_max < 0``, ``num_quadrature < 1``, an unknown
            ``input_mode``, or an ``input_unit`` that is missing for
            ``"lonlat"``, unknown, or given for ``"xyz"``.
    """
    if l_max < 0:
        raise ValueError(f"l_max must be >= 0; got {l_max}.")
    if num_quadrature < 1:
        raise ValueError(f"num_quadrature must be >= 1; got {num_quadrature}.")
    if input_mode not in ("xyz", "lonlat"):
        raise ValueError(
            f"input_mode must be 'xyz' or 'lonlat'; got {input_mode!r}."
        )
    if input_mode == "lonlat" and input_unit not in ("degrees", "radians"):
        raise ValueError(
            "input_mode='lonlat' needs an explicit input_unit, 'degrees' or "
            f"'radians'; got {input_unit!r}."
        )
    if input_mode == "xyz" and input_unit is not None:
        raise ValueError(
            f"input_unit applies to input_mode='lonlat' only; got {input_unit!r}."
        )
    return cls(
        kernel=kernel,
        l_max=l_max,
        input_mode=input_mode,
        input_unit=input_unit,
        num_quadrature=num_quadrature,
        pyrox_name=pyrox_name,
    )

SIREN — Sinusoidal Representation Networks

SIREN (Sitzmann, Martel, Bergman, Lindell, Wetzstein — NeurIPS 2020) replaces ReLU/GELU with sin and prescribes a three-regime initialisation scheme that keeps pre-activation variance stable across depth.

Three-regime weight initialisation (Theorem 1)

Layer	W init	Activation
"first"	U(-1/d_in, 1/d_in)	sin(ω₀ · (W x + b))
"hidden"	U(-√(c/d_in)/ω, √(c/d_in)/ω)	sin(ω · (W x + b))
"last"	U(-√(c/d_in), √(c/d_in))	none (linear) — W x + b

Bias b is initialised U(-1/√d_in, 1/√d_in) for every regime. Typical choice: ω₀ = ω = 30 for image / high-frequency INR tasks.

Usage

import jax.random as jr, jax.numpy as jnp
from pyrox_nn import SirenDense, SIREN, BayesianSIREN

# Single layer
layer = SirenDense.init(3, 64, key=jr.PRNGKey(0), layer_type="first")
y = layer(jnp.ones((5, 3)))  # (5, 64)

# Multi-layer network (depth=5 → first + 3 hidden + last)
net = SIREN.init(2, 64, 1, depth=5, key=jr.PRNGKey(0))
y = net(jnp.zeros((100, 2)))  # (100, 1)

# Bayesian variant (no key needed — weights come from the prior)
from numpyro import handlers
bnet = BayesianSIREN.init(2, 32, 1, depth=3)
with handlers.seed(rng_seed=0):
    y = bnet(jnp.zeros((10, 2)))  # (10, 1)

Alternative INR backbone

SIREN and GaborNet / FourierNet (MFN, #87) are complementary INR backbones: SIREN composes nonlinearities deeply, while MFN uses a product of Gabor filters. Choose based on the signal's smoothness profile.

SirenDense

Bases: Module

Sine-activated dense layer: y = sin(ω · (W x + b)) or y = W x + b.

Single primitive of a SIREN network with three init regimes (Sitzmann et al. 2020, Theorem 1):

+----------+-------------------------------------------+-------------+ | Regime | W init | Activation | +==========+===========================================+=============+ | first | U(-1/d_in, 1/d_in) | sin(ω··)| +----------+-------------------------------------------+-------------+ | hidden | U(-√(c/d_in)/ω, √(c/d_in)/ω) | sin(ω··)| +----------+-------------------------------------------+-------------+ | last | U(-√(c/d_in), √(c/d_in)) | none | +----------+-------------------------------------------+-------------+

Bias b is initialised U(-1/√d_in, 1/√d_in) for every regime.

Attributes:

Name	Type	Description
W	Float[Array, 'in_features out_features']	Weight matrix of shape (in_features, out_features).
b	Float[Array, ' out_features']	Bias vector of shape (out_features,).
omega	float	Frequency multiplier applied inside the sine.
in_features	int	Input dimension.
out_features	int	Output dimension.
layer_type	SirenLayerType	One of "first", "hidden", "last".
c	float	Constant from Theorem 1 (default 6.0).

Source code in .venv/lib/python3.12/site-packages/geonnax/siren.py

class SirenDense(eqx.Module):
    r"""Sine-activated dense layer: ``y = sin(ω · (W x + b))`` or ``y = W x + b``.

    Single primitive of a SIREN network with three init regimes
    (Sitzmann et al. 2020, Theorem 1):

    +----------+-------------------------------------------+-------------+
    | Regime   | ``W`` init                                | Activation  |
    +==========+===========================================+=============+
    | first    | ``U(-1/d_in, 1/d_in)``                    | ``sin(ω··)``|
    +----------+-------------------------------------------+-------------+
    | hidden   | ``U(-√(c/d_in)/ω, √(c/d_in)/ω)``         | ``sin(ω··)``|
    +----------+-------------------------------------------+-------------+
    | last     | ``U(-√(c/d_in), √(c/d_in))``              | none        |
    +----------+-------------------------------------------+-------------+

    Bias ``b`` is initialised ``U(-1/√d_in, 1/√d_in)`` for every regime.

    Attributes:
        W: Weight matrix of shape ``(in_features, out_features)``.
        b: Bias vector of shape ``(out_features,)``.
        omega: Frequency multiplier applied inside the sine.
        in_features: Input dimension.
        out_features: Output dimension.
        layer_type: One of ``"first"``, ``"hidden"``, ``"last"``.
        c: Constant from Theorem 1 (default 6.0).
    """

    W: Float[Array, "in_features out_features"]
    b: Float[Array, " out_features"]
    omega: float = eqx.field(static=True)
    in_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    layer_type: SirenLayerType = eqx.field(static=True)
    c: float = eqx.field(static=True, default=6.0)

    @classmethod
    def init(
        cls,
        in_features: int,
        out_features: int,
        *,
        key: Array,
        omega: float = 30.0,
        layer_type: SirenLayerType = "hidden",
        c: float = 6.0,
    ) -> SirenDense:
        """Construct a ``SirenDense`` with Sitzmann-regime weight initialisation.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.siren import SirenDense
            >>> layer = SirenDense.init(
            ...     3, 8, key=jr.PRNGKey(0), layer_type="first"
            ... )
            >>> layer(jnp.ones(3)).shape  # (3,) -> (8,)
            (8,)
        """
        _require_positive(
            in_features=in_features,
            out_features=out_features,
            omega=omega,
            c=c,
        )
        k_w, k_b = jax.random.split(key)
        # siren_W_limit validates layer_type.
        w_limit = siren_W_limit(layer_type, in_features, omega, c)
        W = jax.random.uniform(
            k_w, (in_features, out_features), minval=-w_limit, maxval=w_limit
        )
        b_limit = 1.0 / math.sqrt(in_features)
        b = jax.random.uniform(k_b, (out_features,), minval=-b_limit, maxval=b_limit)
        return cls(
            W=W,
            b=b,
            omega=omega,
            in_features=in_features,
            out_features=out_features,
            layer_type=layer_type,
            c=c,
        )

    def __call__(self, x: Float[Array, " D_in"]) -> Float[Array, " D_out"]:
        r"""Apply the layer: ``sin(ω · (W x + b))``, or ``W x + b`` if ``last``.

        Args:
            x: Input vector of shape ``(in_features,)``.

        Returns:
            Output vector of shape ``(out_features,)``.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.siren import SirenDense
            >>> layer = SirenDense.init(4, 5, key=jr.PRNGKey(0))
            >>> layer(jnp.ones(4)).shape  # (4,) -> (5,)
            (5,)
        """
        # Affine map: (D_in,) · (D_in, D_out) -> (D_out,), then add bias.
        pre = einx.dot("i, i o -> o", x, self.W) + self.b
        if self.layer_type == "last":
            return pre  # readout layer has no activation
        # Sine activation with frequency multiplier ω: y = sin(ω · pre).
        return jnp.sin(self.omega * pre)

init(in_features: int, out_features: int, *, key: Array, omega: float = 30.0, layer_type: SirenLayerType = 'hidden', c: float = 6.0) -> SirenDense classmethod

Construct a SirenDense with Sitzmann-regime weight initialisation.

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from geonnax.siren import SirenDense
>>> layer = SirenDense.init(
...     3, 8, key=jr.PRNGKey(0), layer_type="first"
... )
>>> layer(jnp.ones(3)).shape  # (3,) -> (8,)
(8,)

Source code in .venv/lib/python3.12/site-packages/geonnax/siren.py

@classmethod
def init(
    cls,
    in_features: int,
    out_features: int,
    *,
    key: Array,
    omega: float = 30.0,
    layer_type: SirenLayerType = "hidden",
    c: float = 6.0,
) -> SirenDense:
    """Construct a ``SirenDense`` with Sitzmann-regime weight initialisation.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from geonnax.siren import SirenDense
        >>> layer = SirenDense.init(
        ...     3, 8, key=jr.PRNGKey(0), layer_type="first"
        ... )
        >>> layer(jnp.ones(3)).shape  # (3,) -> (8,)
        (8,)
    """
    _require_positive(
        in_features=in_features,
        out_features=out_features,
        omega=omega,
        c=c,
    )
    k_w, k_b = jax.random.split(key)
    # siren_W_limit validates layer_type.
    w_limit = siren_W_limit(layer_type, in_features, omega, c)
    W = jax.random.uniform(
        k_w, (in_features, out_features), minval=-w_limit, maxval=w_limit
    )
    b_limit = 1.0 / math.sqrt(in_features)
    b = jax.random.uniform(k_b, (out_features,), minval=-b_limit, maxval=b_limit)
    return cls(
        W=W,
        b=b,
        omega=omega,
        in_features=in_features,
        out_features=out_features,
        layer_type=layer_type,
        c=c,
    )

SIREN

Bases: Module

Multi-layer sinusoidal representation network (Sitzmann et al., NeurIPS 2020).

Topology:

\[ \begin{aligned} z_1 &= \sin(\omega_0 (W_0 x + b_0)), \\ z_{i+1} &= \sin(\omega (W_i z_i + b_i)), \quad i = 1 \ldots L-1, \\ y &= W_L z_L + b_L. \end{aligned} \]

Each layer uses the corresponding Sitzmann Theorem 1 init regime (SirenDense): "first" for layer 0, "hidden" for intermediate layers, and "last" for the readout.

depth counts all layers including the readout; depth=2 gives one first-layer + one last-layer (no hidden layers); depth=5 gives first + 3 hidden + last. Must be ≥ 2.

Source code in .venv/lib/python3.12/site-packages/geonnax/siren.py

class SIREN(eqx.Module):
    r"""Multi-layer sinusoidal representation network (Sitzmann et al., NeurIPS 2020).

    Topology:

    $$
    \begin{aligned}
    z_1 &= \sin(\omega_0 (W_0 x + b_0)), \\
    z_{i+1} &= \sin(\omega (W_i z_i + b_i)), \quad i = 1 \ldots L-1, \\
    y &= W_L z_L + b_L.
    \end{aligned}
    $$


    Each layer uses the corresponding Sitzmann Theorem 1 init regime
    (`SirenDense`):  ``"first"`` for layer 0, ``"hidden"`` for
    intermediate layers, and ``"last"`` for the readout.

    ``depth`` counts *all* layers including the readout; ``depth=2`` gives
    one first-layer + one last-layer (no hidden layers); ``depth=5`` gives
    first + 3 hidden + last.  Must be ≥ 2.
    """

    layers: list[SirenDense]
    in_features: int = eqx.field(static=True)
    hidden_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    depth: int = eqx.field(static=True)
    first_omega: float = eqx.field(static=True)
    hidden_omega: float = eqx.field(static=True)

    @classmethod
    def init(
        cls,
        in_features: int,
        hidden_features: int,
        out_features: int,
        *,
        depth: int,
        key: Array,
        first_omega: float = 30.0,
        hidden_omega: float = 30.0,
        c: float = 6.0,
    ) -> SIREN:
        """Construct a SIREN with the correct per-layer init regimes.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.siren import SIREN
            >>> net = SIREN.init(2, 16, 1, depth=4, key=jr.PRNGKey(0))
            >>> net(jnp.zeros(2)).shape  # (2,) -> (1,)
            (1,)
            >>> len(net.layers)  # first + 2 hidden + last
            4
        """
        if depth < 2:
            raise ValueError(f"depth must be >= 2 (first + last); got depth={depth}")
        _require_positive(
            in_features=in_features,
            hidden_features=hidden_features,
            out_features=out_features,
            first_omega=first_omega,
            hidden_omega=hidden_omega,
            c=c,
        )
        specs = build_siren_specs(
            in_features,
            hidden_features,
            out_features,
            depth,
            first_omega,
            hidden_omega,
            c,
        )
        keys = jax.random.split(key, depth)
        layers = [
            SirenDense.init(
                spec.in_features,
                spec.out_features,
                key=k,
                omega=spec.omega,
                layer_type=spec.layer_type,
                c=spec.c,
            )
            for spec, k in zip(specs, keys, strict=True)
        ]
        return cls(
            layers=layers,
            in_features=in_features,
            hidden_features=hidden_features,
            out_features=out_features,
            depth=depth,
            first_omega=first_omega,
            hidden_omega=hidden_omega,
        )

    def __call__(self, x: Float[Array, " D_in"]) -> Float[Array, " D_out"]:
        r"""Run the full forward pass, composing the sine-activated layers.

        Args:
            x: Input vector of shape ``(in_features,)``.

        Returns:
            Output vector of shape ``(out_features,)``.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.siren import SIREN
            >>> net = SIREN.init(3, 32, 2, depth=3, key=jr.PRNGKey(0))
            >>> net(jnp.ones(3)).shape  # (3,) -> (2,)
            (2,)
        """
        # z_1 = sin(ω₀(W₀x+b₀)); z_{i+1}=sin(ω(W_i z_i+b_i)); y=W_L z_L+b_L.
        # Shapes: (D_in,) -> (H,) -> … -> (H,) -> (D_out,).
        z = x
        for layer in self.layers:
            z = layer(z)
        return z

init(in_features: int, hidden_features: int, out_features: int, *, depth: int, key: Array, first_omega: float = 30.0, hidden_omega: float = 30.0, c: float = 6.0) -> SIREN classmethod

Construct a SIREN with the correct per-layer init regimes.

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from geonnax.siren import SIREN
>>> net = SIREN.init(2, 16, 1, depth=4, key=jr.PRNGKey(0))
>>> net(jnp.zeros(2)).shape  # (2,) -> (1,)
(1,)
>>> len(net.layers)  # first + 2 hidden + last
4

Source code in .venv/lib/python3.12/site-packages/geonnax/siren.py

@classmethod
def init(
    cls,
    in_features: int,
    hidden_features: int,
    out_features: int,
    *,
    depth: int,
    key: Array,
    first_omega: float = 30.0,
    hidden_omega: float = 30.0,
    c: float = 6.0,
) -> SIREN:
    """Construct a SIREN with the correct per-layer init regimes.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from geonnax.siren import SIREN
        >>> net = SIREN.init(2, 16, 1, depth=4, key=jr.PRNGKey(0))
        >>> net(jnp.zeros(2)).shape  # (2,) -> (1,)
        (1,)
        >>> len(net.layers)  # first + 2 hidden + last
        4
    """
    if depth < 2:
        raise ValueError(f"depth must be >= 2 (first + last); got depth={depth}")
    _require_positive(
        in_features=in_features,
        hidden_features=hidden_features,
        out_features=out_features,
        first_omega=first_omega,
        hidden_omega=hidden_omega,
        c=c,
    )
    specs = build_siren_specs(
        in_features,
        hidden_features,
        out_features,
        depth,
        first_omega,
        hidden_omega,
        c,
    )
    keys = jax.random.split(key, depth)
    layers = [
        SirenDense.init(
            spec.in_features,
            spec.out_features,
            key=k,
            omega=spec.omega,
            layer_type=spec.layer_type,
            c=spec.c,
        )
        for spec, k in zip(specs, keys, strict=True)
    ]
    return cls(
        layers=layers,
        in_features=in_features,
        hidden_features=hidden_features,
        out_features=out_features,
        depth=depth,
        first_omega=first_omega,
        hidden_omega=hidden_omega,
    )

BayesianSIREN

Bases: PyroxModule

SIREN with regime-scaled Normal priors on all layer weights.

Replaces the deterministic weight matrices of SIREN with NumPyro sample sites. For layer \(i\) with Sitzmann Theorem 1 half-width \(a_i\) (the uniform bound used by SirenDense):

\[ W_i \sim \mathcal{N}\!\left(0,\, \sigma_0 \cdot \frac{a_i}{\sqrt{3}}\right), \qquad b_i \sim \mathcal{N}\!\left(0,\, \sigma_0 \cdot \frac{1}{\sqrt{3 \, d_i}}\right), \]

where \(\sigma_0\) is prior_std and \(d_i\) is the input dimension of layer \(i\). The \(a_i / \sqrt{3}\) factor makes \(\operatorname{Var}(W_i)\) equal to the variance of Sitzmann's \(\mathcal{U}(-a_i, a_i)\) init exactly, so the Bayesian prior preserves the activation variance prescribed by Theorem 1 — avoiding the saturated-sine pathology that a flat \(\mathcal{N}(0, 1)\) prior would cause.

Registered sites: {scope}.layer_0.W, {scope}.layer_0.b, …, {scope}.layer_{depth-1}.W, {scope}.layer_{depth-1}.b — exactly 2 · depth sites per forward call.

Attributes:

Name	Type	Description
specs	tuple[SirenLayerSpec, ...]	Tuple of per-layer specs (static). Holds each layer's layer_type, in_features, out_features, omega, and c — i.e. everything needed to scale the priors.
in_features	int	Input dimension.
hidden_features	int	Hidden dimension.
out_features	int	Output dimension.
depth	int	Total layers including readout. Must be ≥ 2.
first_omega	float	Frequency multiplier for the first layer.
hidden_omega	float	Frequency multiplier for hidden layers.
prior_std	float	Scale factor for the regime-scaled Normal prior (default 1.0).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.random as jr, jax.numpy as jnp
>>> from numpyro import handlers
>>> net = BayesianSIREN.init(2, 32, 1, depth=3)
>>> with handlers.seed(rng_seed=0):
...     y = net(jnp.zeros((4, 2)))
>>> y.shape
(4, 1)

Source code in packages/pyrox-nn/src/pyrox_nn/_siren.py

class BayesianSIREN(PyroxModule):
    r"""SIREN with regime-scaled Normal priors on all layer weights.

    Replaces the deterministic weight matrices of `SIREN` with NumPyro
    sample sites.  For layer $i$ with Sitzmann Theorem 1 half-width
    $a_i$ (the uniform bound used by `SirenDense`):

    $$
    W_i \sim \mathcal{N}\!\left(0,\, \sigma_0 \cdot \frac{a_i}{\sqrt{3}}\right),
    \qquad
    b_i \sim \mathcal{N}\!\left(0,\,
        \sigma_0 \cdot \frac{1}{\sqrt{3 \, d_i}}\right),
    $$

    where $\sigma_0$ is ``prior_std`` and $d_i$ is the input
    dimension of layer $i$.  The $a_i / \sqrt{3}$ factor makes
    $\operatorname{Var}(W_i)$ equal to the variance of Sitzmann's
    $\mathcal{U}(-a_i, a_i)$ init exactly, so the Bayesian prior
    preserves the activation variance prescribed by Theorem 1 — avoiding
    the saturated-sine pathology that a flat $\mathcal{N}(0, 1)$
    prior would cause.

    Registered sites: ``{scope}.layer_0.W``, ``{scope}.layer_0.b``, …,
    ``{scope}.layer_{depth-1}.W``, ``{scope}.layer_{depth-1}.b``
    — exactly ``2 · depth`` sites per forward call.

    Attributes:
        specs: Tuple of per-layer specs (static).  Holds each layer's
            ``layer_type``, ``in_features``, ``out_features``, ``omega``,
            and ``c`` — i.e. everything needed to scale the priors.
        in_features: Input dimension.
        hidden_features: Hidden dimension.
        out_features: Output dimension.
        depth: Total layers including readout.  Must be ≥ 2.
        first_omega: Frequency multiplier for the first layer.
        hidden_omega: Frequency multiplier for hidden layers.
        prior_std: Scale factor for the regime-scaled Normal prior (default 1.0).
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.random as jr, jax.numpy as jnp
        >>> from numpyro import handlers
        >>> net = BayesianSIREN.init(2, 32, 1, depth=3)
        >>> with handlers.seed(rng_seed=0):
        ...     y = net(jnp.zeros((4, 2)))
        >>> y.shape
        (4, 1)
    """

    specs: tuple[SirenLayerSpec, ...] = eqx.field(static=True)
    in_features: int = eqx.field(static=True)
    hidden_features: int = eqx.field(static=True)
    out_features: int = eqx.field(static=True)
    depth: int = eqx.field(static=True)
    first_omega: float = eqx.field(static=True)
    hidden_omega: float = eqx.field(static=True)
    prior_std: float = eqx.field(static=True, default=1.0)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @classmethod
    def init(
        cls,
        in_features: int,
        hidden_features: int,
        out_features: int,
        *,
        depth: int,
        first_omega: float = 30.0,
        hidden_omega: float = 30.0,
        c: float = 6.0,
        prior_std: float = 1.0,
        pyrox_name: str | None = None,
    ) -> BayesianSIREN:
        """Construct a `BayesianSIREN`.

        All weights come from the prior, so no PRNG key is needed at
        construction time — the key enters when sampling inside a
        ``numpyro`` handler (``handlers.seed``, SVI, etc.).

        Args:
            in_features: Input dimension.
            hidden_features: Hidden dimension.
            out_features: Output dimension.
            depth: Total layers including readout.  Must be ≥ 2.
            first_omega: Frequency for the first layer.
            hidden_omega: Frequency for hidden layers.
            c: Theorem-1 constant.
            prior_std: Scale factor for the Normal priors (default 1.0, must be > 0).
            pyrox_name: Optional explicit scope name for NumPyro.

        Returns:
            Initialised `BayesianSIREN`.

        Raises:
            ValueError: If ``depth < 2``, or any of the feature dimensions,
                omegas, ``c``, or ``prior_std`` is non-positive.
        """
        if depth < 2:
            raise ValueError(f"depth must be >= 2 (first + last); got depth={depth}")
        _require_positive(
            in_features=in_features,
            hidden_features=hidden_features,
            out_features=out_features,
            first_omega=first_omega,
            hidden_omega=hidden_omega,
            c=c,
            prior_std=prior_std,
        )
        specs = build_siren_specs(
            in_features,
            hidden_features,
            out_features,
            depth,
            first_omega,
            hidden_omega,
            c,
        )
        return cls(
            specs=specs,
            in_features=in_features,
            hidden_features=hidden_features,
            out_features=out_features,
            depth=depth,
            first_omega=first_omega,
            hidden_omega=hidden_omega,
            prior_std=prior_std,
            pyrox_name=pyrox_name,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        """Sample weights from regime-scaled priors and run the forward pass.

        Registers ``layer_{i}.W`` and ``layer_{i}.b`` NumPyro sample sites
        for each layer ``i`` in ``[0, depth)``.

        Args:
            x: Input tensor of shape ``(*batch, in_features)``.

        Returns:
            Output tensor of shape ``(*batch, out_features)``.
        """
        # Normal stddev = (Uniform half-width) / √3 so Var(W) matches
        # Sitzmann's U(-a, a) init exactly.
        inv_sqrt3 = 1.0 / math.sqrt(3.0)
        z = x
        for i, spec in enumerate(self.specs):
            a = siren_W_limit(spec.layer_type, spec.in_features, spec.omega, spec.c)
            w_scale = self.prior_std * a * inv_sqrt3
            b_scale = self.prior_std * inv_sqrt3 / math.sqrt(spec.in_features)
            W = self.pyrox_sample(
                f"layer_{i}.W",
                dist.Normal(0.0, w_scale)
                .expand([spec.in_features, spec.out_features])
                .to_event(2),
            )
            b = self.pyrox_sample(
                f"layer_{i}.b",
                dist.Normal(0.0, b_scale).expand([spec.out_features]).to_event(1),
            )
            pre = einx.dot("... i, i o -> ... o", z, W) + b
            z = pre if spec.layer_type == "last" else jnp.sin(spec.omega * pre)
        return z

init(in_features: int, hidden_features: int, out_features: int, *, depth: int, first_omega: float = 30.0, hidden_omega: float = 30.0, c: float = 6.0, prior_std: float = 1.0, pyrox_name: str | None = None) -> BayesianSIREN classmethod

Construct a BayesianSIREN.

All weights come from the prior, so no PRNG key is needed at construction time — the key enters when sampling inside a numpyro handler (handlers.seed, SVI, etc.).

Parameters:

Name	Type	Description	Default
in_features	int	Input dimension.	required
hidden_features	int	Hidden dimension.	required
out_features	int	Output dimension.	required
depth	int	Total layers including readout. Must be ≥ 2.	required
first_omega	float	Frequency for the first layer.	30.0
hidden_omega	float	Frequency for hidden layers.	30.0
c	float	Theorem-1 constant.	6.0
prior_std	float	Scale factor for the Normal priors (default 1.0, must be > 0).	1.0
pyrox_name	str | None	Optional explicit scope name for NumPyro.	None

Returns:

Type	Description
BayesianSIREN	Initialised BayesianSIREN.

Raises:

Type	Description
ValueError	If depth < 2, or any of the feature dimensions, omegas, c, or prior_std is non-positive.

Source code in packages/pyrox-nn/src/pyrox_nn/_siren.py

@classmethod
def init(
    cls,
    in_features: int,
    hidden_features: int,
    out_features: int,
    *,
    depth: int,
    first_omega: float = 30.0,
    hidden_omega: float = 30.0,
    c: float = 6.0,
    prior_std: float = 1.0,
    pyrox_name: str | None = None,
) -> BayesianSIREN:
    """Construct a `BayesianSIREN`.

    All weights come from the prior, so no PRNG key is needed at
    construction time — the key enters when sampling inside a
    ``numpyro`` handler (``handlers.seed``, SVI, etc.).

    Args:
        in_features: Input dimension.
        hidden_features: Hidden dimension.
        out_features: Output dimension.
        depth: Total layers including readout.  Must be ≥ 2.
        first_omega: Frequency for the first layer.
        hidden_omega: Frequency for hidden layers.
        c: Theorem-1 constant.
        prior_std: Scale factor for the Normal priors (default 1.0, must be > 0).
        pyrox_name: Optional explicit scope name for NumPyro.

    Returns:
        Initialised `BayesianSIREN`.

    Raises:
        ValueError: If ``depth < 2``, or any of the feature dimensions,
            omegas, ``c``, or ``prior_std`` is non-positive.
    """
    if depth < 2:
        raise ValueError(f"depth must be >= 2 (first + last); got depth={depth}")
    _require_positive(
        in_features=in_features,
        hidden_features=hidden_features,
        out_features=out_features,
        first_omega=first_omega,
        hidden_omega=hidden_omega,
        c=c,
        prior_std=prior_std,
    )
    specs = build_siren_specs(
        in_features,
        hidden_features,
        out_features,
        depth,
        first_omega,
        hidden_omega,
        c,
    )
    return cls(
        specs=specs,
        in_features=in_features,
        hidden_features=hidden_features,
        out_features=out_features,
        depth=depth,
        first_omega=first_omega,
        hidden_omega=hidden_omega,
        prior_std=prior_std,
        pyrox_name=pyrox_name,
    )

SNGP — spectral-normalised GP head

The SNGP output layer (Liu et al., 2020): a random-feature GP last layer whose posterior covariance comes from a Laplace approximation (LaplaceRandomFeatureCovariance, re-exported from geonnax), giving distance-aware uncertainty from a single deterministic forward pass.

RandomFeatureGaussianProcess

Bases: PyroxModule

SNGP output layer (Liu et al., 2020).

A random Fourier feature map followed by a learnable linear head, plus a Laplace-approximation covariance over the linear weights. The forward pass returns the mean prediction and (optionally) a per-input predictive variance summarising distance from the training distribution.

Forward (mean):

\[ \phi(x) = \sqrt{\tfrac{2}{D}}\,\cos\!\bigl(W\, x / \ell + b\bigr), \qquad \mu(x) = \phi(x)\, H + b_H. \]

The frequencies \(W\) and bias \(b\) of the RFF map are frozen (they implicitly define the kernel approximation): they are registered as pyrox_param sites for substitution and checkpointing, then guarded with jax.lax.stop_gradient inside feature_map so SGD-style optimisers leave them untouched. The lengthscale \(\ell\), the linear head \(H, b_H\), and the Laplace precision are the trainable / updated quantities.

Predictive variance — when \(\hat{\Lambda}\) is the current precision matrix:

\[ \sigma^2(x_*) = \phi(x_*)^\top \hat{\Lambda}^{-1}\, \phi(x_*). \]

Training pattern (one minibatch):

1. mean = layer(x) registers / reuses the trainable params and returns the mean prediction. Compute the loss, take a gradient step on the SVI parameter store as usual.
2. After the gradient step, call new_layer = layer.update_precision(features) where features is the result of feature_map evaluated on the same minibatch using the updated parameters. This returns a new layer with the LRFC's precision EMA-updated.

At inference, mean, var = layer(x, return_cov=True) produces the mean and the Laplace per-input predictive variance.

Plate semantics

Same as the rest of pyrox_nn's Bayesian / heteroscedastic dense layers — call this layer outside numpyro.plate("data", ..., subsample_size=...) and only plate the observation likelihood.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_\mathrm{in}\).
num_features	int	Number of random Fourier features \(D\).
out_features	int	Output dimension \(D_\mathrm{out}\).
init_lengthscale	float	Initial lengthscale \(\ell\). Optimised during training as a positive pyrox_param.
W_init	Float[Array, 'D_in D']	Frozen RFF frequencies, shape (D_in, D), drawn from a standard Normal (the RBF spectral density).
bias_init	Float[Array, ' D']	Frozen RFF biases, shape (D,), drawn from Uniform(0, 2 pi).
output_linear_init	Float[Array, 'D D_out']	Init for the linear head, shape (D, D_out).
covariance	LaplaceRandomFeatureCovariance	The LaplaceRandomFeatureCovariance instance.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

References

Liu, J. Z., et al. (2020). Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness. NeurIPS.

Source code in packages/pyrox-nn/src/pyrox_nn/_sngp.py

class RandomFeatureGaussianProcess(PyroxModule):
    r"""SNGP output layer (Liu et al., 2020).

    A random Fourier feature map followed by a learnable linear head,
    plus a Laplace-approximation covariance over the linear weights.
    The forward pass returns the mean prediction and (optionally) a
    per-input predictive variance summarising distance from the
    training distribution.

    Forward (mean):

    $$
    \phi(x) = \sqrt{\tfrac{2}{D}}\,\cos\!\bigl(W\, x / \ell + b\bigr),
    \qquad \mu(x) = \phi(x)\, H + b_H.
    $$

    The frequencies $W$ and bias $b$ of the RFF map are
    *frozen* (they implicitly define the kernel approximation): they
    are registered as ``pyrox_param`` sites for substitution and
    checkpointing, then guarded with `jax.lax.stop_gradient`
    inside `feature_map` so SGD-style optimisers leave them
    untouched. The lengthscale $\ell$, the linear head
    $H, b_H$, and the Laplace precision are the trainable /
    updated quantities.

    Predictive variance — when $\hat{\Lambda}$ is the current
    precision matrix:

    $$
    \sigma^2(x_*) = \phi(x_*)^\top \hat{\Lambda}^{-1}\, \phi(x_*).
    $$

    Training pattern (one minibatch):

    1. ``mean = layer(x)`` registers / reuses the trainable params and
       returns the mean prediction. Compute the loss, take a gradient
       step on the SVI parameter store as usual.
    2. After the gradient step, call
       ``new_layer = layer.update_precision(features)`` where
       ``features`` is the result of `feature_map` evaluated on
       the same minibatch using the *updated* parameters. This returns
       a new layer with the LRFC's precision EMA-updated.

    At inference, ``mean, var = layer(x, return_cov=True)`` produces
    the mean and the Laplace per-input predictive variance.

    Plate semantics:
        Same as the rest of ``pyrox_nn``'s Bayesian / heteroscedastic
        dense layers — call this layer outside
        ``numpyro.plate("data", ..., subsample_size=...)`` and only
        plate the observation likelihood.

    Attributes:
        in_features: Input dimension $D_\mathrm{in}$.
        num_features: Number of random Fourier features $D$.
        out_features: Output dimension $D_\mathrm{out}$.
        init_lengthscale: Initial lengthscale $\ell$. Optimised
            during training as a positive ``pyrox_param``.
        W_init: Frozen RFF frequencies, shape ``(D_in, D)``, drawn from
            a standard Normal (the RBF spectral density).
        bias_init: Frozen RFF biases, shape ``(D,)``, drawn from
            ``Uniform(0, 2 pi)``.
        output_linear_init: Init for the linear head, shape ``(D, D_out)``.
        covariance: The `LaplaceRandomFeatureCovariance` instance.
        pyrox_name: Explicit scope name for NumPyro site registration.

    References:
        Liu, J. Z., et al. (2020). *Simple and Principled Uncertainty
        Estimation with Deterministic Deep Learning via Distance
        Awareness.* NeurIPS.
    """

    core: geonnax.RandomFeatureGaussianProcess
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        key: PRNGKeyArray,
        in_features: int,
        num_features: int,
        out_features: int,
        *,
        init_lengthscale: float = 1.0,
        momentum: float = 0.999,
        ridge: float = 1.0,
        head_scale: float = 0.01,
        pyrox_name: str | None = None,
    ) -> RandomFeatureGaussianProcess:
        """Construct an SNGP head with frozen RFF freqs and an empty precision."""
        # geonnax validates positive dims and init_lengthscale > 0;
        # momentum / ridge constraints are validated inside the LRFC init.
        core = geonnax.RandomFeatureGaussianProcess.init(
            in_features=in_features,
            num_features=num_features,
            out_features=out_features,
            key=key,
            init_lengthscale=init_lengthscale,
            momentum=momentum,
            ridge=ridge,
            head_scale=head_scale,
        )
        return cls(core=core, pyrox_name=pyrox_name)

    # Read-only property accessors retain the pre-refactor attribute names so
    # external callers (tests, user code reading static dims) keep working.
    @property
    def in_features(self) -> int:
        return self.core.in_features

    @property
    def num_features(self) -> int:
        return self.core.num_features

    @property
    def out_features(self) -> int:
        return self.core.out_features

    @property
    def init_lengthscale(self) -> float:
        return float(self.core.lengthscale)

    @property
    def W_init(self) -> Float[Array, "D_in D"]:
        return self.core.W

    @property
    def bias_init(self) -> Float[Array, " D"]:
        return self.core.bias

    @property
    def output_linear_init(self) -> Float[Array, "D D_out"]:
        return self.core.output_linear

    @property
    def covariance(self) -> geonnax.LaplaceRandomFeatureCovariance:
        return self.core.covariance

    def _swap_feature_core(self) -> geonnax.RandomFeatureGaussianProcess:
        """Register the RFF-map params and swap them into the core.

        Only the frequency/bias/lengthscale arrays are needed for the
        feature map; the linear-head params are registered inside
        ``__call__`` so disabled branches (e.g. pure feature-map usage
        via `feature_map`) don't materialise unused sites.
        """
        W = jax.lax.stop_gradient(self.pyrox_param("W", self.core.W))
        b = jax.lax.stop_gradient(self.pyrox_param("bias", self.core.bias))
        ls = self.pyrox_param(
            "lengthscale",
            self.core.lengthscale,
            constraint=dist.constraints.positive,
        )
        return eqx.tree_at(
            lambda c: (c.W, c.bias, c.lengthscale),
            self.core,
            (W, b, ls),
        )

    @pyrox_method
    def feature_map(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D"]:
        r"""Random Fourier feature map: $\phi(x) = \sqrt{2/D}\,\cos(Wx/\ell + b)$.

        Frequencies and bias are registered as ``pyrox_param`` sites for
        substitution / checkpointing, but `jax.lax.stop_gradient`
        is applied so SVI's gradient-based optimisers leave them
        frozen at their init values. The lengthscale is the active
        bandwidth control and is constrained positive.
        """
        new_core = self._swap_feature_core()
        # geonnax feature_map is single-example `(D_in,) -> (D,)`.
        return vmap_over_flat_batch(new_core.feature_map, x)

    @pyrox_method
    def __call__(
        self,
        x: Float[Array, "*batch D_in"],
        *,
        return_cov: bool = False,
    ) -> (
        Float[Array, "*batch D_out"]
        | tuple[Float[Array, "*batch D_out"], Float[Array, " *batch"]]
    ):
        # Register the RFF + linear-head params and swap them into the core.
        new_core = self._swap_feature_core()
        H = self.pyrox_param("output_linear", self.core.output_linear)
        b_out = self.pyrox_param(
            "output_bias", jnp.zeros(self.out_features, dtype=x.dtype)
        )
        new_core = eqx.tree_at(
            lambda c: (c.output_linear, c.output_bias),
            new_core,
            (H, b_out),
        )

        # geonnax `__call__` is single-example `(D_in,)`. `return_cov=False`
        # keeps the output a plain array; `return_cov=True` returns
        # (mean, var) per example and the helper restores both leaves.
        if return_cov:
            mean, var = vmap_over_flat_batch(
                lambda xi: new_core(xi, return_cov=True), x
            )
            return mean, var

        return vmap_over_flat_batch(new_core, x)

    def update_precision(
        self, features: Float[Array, "*batch D"]
    ) -> RandomFeatureGaussianProcess:
        """Return a new layer with an EMA-updated Laplace precision.

        Pure-functional: ``self`` is unchanged. Pass features computed
        on the current minibatch (e.g. via `feature_map`) — the
        update folds the empirical second moment into the EMA. Call
        this once per training batch *after* the gradient step.
        """
        # Flatten any leading batch dims down to the single batch axis the
        # geonnax `update_precision` expects.
        flat = einx.id("b... d -> (b...) d", features)
        new_core = self.core.update_precision(flat)
        return eqx.tree_at(lambda layer: layer.core, self, new_core)

init(key: PRNGKeyArray, in_features: int, num_features: int, out_features: int, *, init_lengthscale: float = 1.0, momentum: float = 0.999, ridge: float = 1.0, head_scale: float = 0.01, pyrox_name: str | None = None) -> RandomFeatureGaussianProcess classmethod

Construct an SNGP head with frozen RFF freqs and an empty precision.

Source code in packages/pyrox-nn/src/pyrox_nn/_sngp.py

@classmethod
def init(
    cls,
    key: PRNGKeyArray,
    in_features: int,
    num_features: int,
    out_features: int,
    *,
    init_lengthscale: float = 1.0,
    momentum: float = 0.999,
    ridge: float = 1.0,
    head_scale: float = 0.01,
    pyrox_name: str | None = None,
) -> RandomFeatureGaussianProcess:
    """Construct an SNGP head with frozen RFF freqs and an empty precision."""
    # geonnax validates positive dims and init_lengthscale > 0;
    # momentum / ridge constraints are validated inside the LRFC init.
    core = geonnax.RandomFeatureGaussianProcess.init(
        in_features=in_features,
        num_features=num_features,
        out_features=out_features,
        key=key,
        init_lengthscale=init_lengthscale,
        momentum=momentum,
        ridge=ridge,
        head_scale=head_scale,
    )
    return cls(core=core, pyrox_name=pyrox_name)

feature_map(x: Float[Array, '*batch D_in']) -> Float[Array, '*batch D']

Random Fourier feature map: \(\phi(x) = \sqrt{2/D}\,\cos(Wx/\ell + b)\).

Frequencies and bias are registered as pyrox_param sites for substitution / checkpointing, but jax.lax.stop_gradient is applied so SVI's gradient-based optimisers leave them frozen at their init values. The lengthscale is the active bandwidth control and is constrained positive.

Source code in packages/pyrox-nn/src/pyrox_nn/_sngp.py

@pyrox_method
def feature_map(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D"]:
    r"""Random Fourier feature map: $\phi(x) = \sqrt{2/D}\,\cos(Wx/\ell + b)$.

    Frequencies and bias are registered as ``pyrox_param`` sites for
    substitution / checkpointing, but `jax.lax.stop_gradient`
    is applied so SVI's gradient-based optimisers leave them
    frozen at their init values. The lengthscale is the active
    bandwidth control and is constrained positive.
    """
    new_core = self._swap_feature_core()
    # geonnax feature_map is single-example `(D_in,) -> (D,)`.
    return vmap_over_flat_batch(new_core.feature_map, x)

update_precision(features: Float[Array, '*batch D']) -> RandomFeatureGaussianProcess

Return a new layer with an EMA-updated Laplace precision.

Pure-functional: self is unchanged. Pass features computed on the current minibatch (e.g. via feature_map) — the update folds the empirical second moment into the EMA. Call this once per training batch after the gradient step.

Source code in packages/pyrox-nn/src/pyrox_nn/_sngp.py

def update_precision(
    self, features: Float[Array, "*batch D"]
) -> RandomFeatureGaussianProcess:
    """Return a new layer with an EMA-updated Laplace precision.

    Pure-functional: ``self`` is unchanged. Pass features computed
    on the current minibatch (e.g. via `feature_map`) — the
    update folds the empirical second moment into the EMA. Call
    this once per training batch *after* the gradient step.
    """
    # Flatten any leading batch dims down to the single batch axis the
    # geonnax `update_precision` expects.
    flat = einx.id("b... d -> (b...) d", features)
    new_core = self.core.update_precision(flat)
    return eqx.tree_at(lambda layer: layer.core, self, new_core)

LaplaceRandomFeatureCovariance

Bases: Module

Laplace-approximation precision for an SNGP output head.

Stores the precision matrix \(\hat{\Lambda} \in \mathbb{R}^{D \times D}\) over the linear weights of the output layer. Updated as an exponential moving average of feature outer products during training:

\[ \hat{\Lambda}_{t+1} \leftarrow m\,\hat{\Lambda}_t + (1 - m)\,\frac{1}{B} \sum_{b=1}^{B} \phi(x_b)\,\phi(x_b)^\top. \]

At test time the predictive variance for a feature vector \(\phi(x_*)\) is

\[ \sigma^2(x_*) = \phi(x_*)^\top \hat{\Sigma}\, \phi(x_*), \qquad \hat{\Sigma} = \hat{\Lambda}^{-1}, \]

computed stably via a Cholesky solve.

The container is pure-functional: update returns a new instance with an updated precision rather than mutating self, matching how Equinox composes immutable PyTrees with optimisers. A small ridge \(\lambda\) initialises the precision at \(\lambda I\) and is also added at solve-time inside covariance and variance_at so the Cholesky stays numerically well-conditioned even after many EMA steps with low momentum (which would otherwise let the ridge contribution decay geometrically and the precision approach singularity). Equivalently \(\hat\Sigma = (\hat\Lambda + \lambda I)^{-1}\) — the Bayesian-linear-regression interpretation of SNGP, where \(\lambda I\) is a Gaussian prior precision on the head weights.

Attributes:

Name	Type	Description
precision	Float[Array, 'D D']	Current precision matrix \(\hat{\Lambda}\).
momentum	float	EMA momentum \(m \in [0, 1]\). Higher values give slower updates; 0.999 works well for most settings.
ridge	float	Diagonal ridge \(\lambda\). Used both as the init value of precision and as a solve-time jitter to keep the Cholesky well-defined.

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.sngp import LaplaceRandomFeatureCovariance
>>> cov = LaplaceRandomFeatureCovariance.init(4, ridge=1.0)
>>> cov.precision.shape
(4, 4)
>>> cov.variance_at(jnp.eye(4)).shape
(4,)

Source code in .venv/lib/python3.12/site-packages/geonnax/sngp.py

class LaplaceRandomFeatureCovariance(eqx.Module):
    r"""Laplace-approximation precision for an SNGP output head.

    Stores the precision matrix $\hat{\Lambda} \in \mathbb{R}^{D \times D}$
    over the linear weights of the output layer. Updated as an
    exponential moving average of feature outer products during
    training:

    $$
    \hat{\Lambda}_{t+1} \leftarrow m\,\hat{\Lambda}_t
    + (1 - m)\,\frac{1}{B} \sum_{b=1}^{B} \phi(x_b)\,\phi(x_b)^\top.
    $$


    At test time the predictive variance for a feature vector
    $\phi(x_*)$ is

    $$
    \sigma^2(x_*) = \phi(x_*)^\top \hat{\Sigma}\, \phi(x_*),
    \qquad \hat{\Sigma} = \hat{\Lambda}^{-1},
    $$


    computed stably via a Cholesky solve.

    The container is *pure-functional*: `update` returns a new
    instance with an updated precision rather than mutating ``self``,
    matching how Equinox composes immutable PyTrees with optimisers.
    A small ridge $\lambda$ initialises the precision at
    $\lambda I$ and is *also* added at solve-time inside
    `covariance` and `variance_at` so the Cholesky stays
    numerically well-conditioned even after many EMA steps with low
    momentum (which would otherwise let the ridge contribution decay
    geometrically and the precision approach singularity).
    Equivalently $\hat\Sigma = (\hat\Lambda + \lambda I)^{-1}$ —
    the Bayesian-linear-regression interpretation of SNGP, where
    $\lambda I$ is a Gaussian prior precision on the head weights.

    Attributes:
        precision: Current precision matrix $\hat{\Lambda}$.
        momentum: EMA momentum $m \in [0, 1]$. Higher values give
            slower updates; ``0.999`` works well for most settings.
        ridge: Diagonal ridge $\lambda$. Used both as the init
            value of ``precision`` and as a solve-time jitter to keep
            the Cholesky well-defined.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
        >>> cov = LaplaceRandomFeatureCovariance.init(4, ridge=1.0)
        >>> cov.precision.shape
        (4, 4)
        >>> cov.variance_at(jnp.eye(4)).shape
        (4,)
    """

    precision: Float[Array, "D D"]
    momentum: float = eqx.field(static=True, default=0.999)
    ridge: float = eqx.field(static=True, default=1.0)

    @classmethod
    def init(
        cls,
        num_features: int,
        *,
        momentum: float = 0.999,
        ridge: float = 1.0,
    ) -> LaplaceRandomFeatureCovariance:
        r"""Construct a fresh covariance container with ``ridge * I`` precision.

        Examples:
            >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
            >>> cov = LaplaceRandomFeatureCovariance.init(4, momentum=0.9)
            >>> cov.precision.shape
            (4, 4)
        """
        if num_features <= 0:
            raise ValueError(f"num_features must be > 0; got {num_features}.")
        if not 0.0 <= momentum <= 1.0:
            raise ValueError(f"momentum must lie in [0, 1]; got {momentum}.")
        if ridge <= 0:
            raise ValueError(f"ridge must be > 0; got {ridge}.")
        return cls(
            precision=ridge * jnp.eye(num_features),
            momentum=momentum,
            ridge=ridge,
        )

    def update(self, features: Float[Array, "B D"]) -> LaplaceRandomFeatureCovariance:
        r"""Return a new container with EMA-updated precision.

        $\hat\Lambda \leftarrow m\,\hat\Lambda + (1-m)\,\Phi^\top\Phi/B$.

        Examples:
            >>> import jax.numpy as jnp
            >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
            >>> cov = LaplaceRandomFeatureCovariance.init(4)
            >>> new = cov.update(jnp.ones((8, 4)))
            >>> new is cov
            False
        """
        B = features.shape[0]
        # Feature Gram ΦᵀΦ / B: contract the batch axis b → (D, D).
        outer = einx.dot("b d, b e -> d e", features, features) / B
        new_precision = self.momentum * self.precision + (1.0 - self.momentum) * outer
        return eqx.tree_at(lambda c: c.precision, self, new_precision)

    def _chol(self) -> Float[Array, "D D"]:
        # Symmetrise to absorb floating-point asymmetry in the EMA, then
        # add ridge jitter so the matrix is guaranteed positive-definite
        # regardless of how the EMA has evolved.
        sym = 0.5 * (self.precision + einx.id("i j -> j i", self.precision))
        D = sym.shape[0]
        return jnp.linalg.cholesky(sym + self.ridge * jnp.eye(D, dtype=sym.dtype))

    def covariance(self) -> Float[Array, "D D"]:
        r"""Inverse of the precision matrix (one-shot Cholesky inversion).

        $\hat\Sigma = (\hat\Lambda + \lambda I)^{-1}$, shape ``(D, D)``.

        Examples:
            >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
            >>> cov = LaplaceRandomFeatureCovariance.init(4)
            >>> cov.covariance().shape
            (4, 4)
        """
        L = self._chol()
        D = self.precision.shape[0]
        return jax.scipy.linalg.cho_solve((L, True), jnp.eye(D))

    def variance_at(self, features: Float[Array, "N D"]) -> Float[Array, " N"]:
        r"""Per-row predictive variance $\phi(x_n)^\top \hat{\Sigma}\,\phi(x_n)$.

        Computed via a triangular solve to avoid materialising the full
        $D \times D$ covariance:

        $$
        y = L^{-1} \phi(x_n)^\top, \qquad
        \sigma^2(x_n) = \lVert y \rVert_2^2
        = \phi(x_n)^\top (L L^\top)^{-1} \phi(x_n).
        $$


        Examples:
            >>> import jax.numpy as jnp
            >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
            >>> cov = LaplaceRandomFeatureCovariance.init(4)
            >>> cov.variance_at(jnp.eye(4)).shape
            (4,)
        """
        L = self._chol()
        # Solve L y = Φᵀ ; (D, D)\(D, N) -> (D, N), one column per row of Φ.
        y = jax.scipy.linalg.solve_triangular(
            L, einx.id("n d -> d n", features), lower=True
        )
        # σ²(xₙ) = ‖yₙ‖² ; sum over D -> (N,)
        return einx.sum("[d] n", y * y)

init(num_features: int, *, momentum: float = 0.999, ridge: float = 1.0) -> LaplaceRandomFeatureCovariance classmethod

Construct a fresh covariance container with ridge * I precision.

Examples:

>>> from geonnax.sngp import LaplaceRandomFeatureCovariance
>>> cov = LaplaceRandomFeatureCovariance.init(4, momentum=0.9)
>>> cov.precision.shape
(4, 4)

Source code in .venv/lib/python3.12/site-packages/geonnax/sngp.py

@classmethod
def init(
    cls,
    num_features: int,
    *,
    momentum: float = 0.999,
    ridge: float = 1.0,
) -> LaplaceRandomFeatureCovariance:
    r"""Construct a fresh covariance container with ``ridge * I`` precision.

    Examples:
        >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
        >>> cov = LaplaceRandomFeatureCovariance.init(4, momentum=0.9)
        >>> cov.precision.shape
        (4, 4)
    """
    if num_features <= 0:
        raise ValueError(f"num_features must be > 0; got {num_features}.")
    if not 0.0 <= momentum <= 1.0:
        raise ValueError(f"momentum must lie in [0, 1]; got {momentum}.")
    if ridge <= 0:
        raise ValueError(f"ridge must be > 0; got {ridge}.")
    return cls(
        precision=ridge * jnp.eye(num_features),
        momentum=momentum,
        ridge=ridge,
    )

update(features: Float[Array, 'B D']) -> LaplaceRandomFeatureCovariance

Return a new container with EMA-updated precision.

\(\hat\Lambda \leftarrow m\,\hat\Lambda + (1-m)\,\Phi^\top\Phi/B\).

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.sngp import LaplaceRandomFeatureCovariance
>>> cov = LaplaceRandomFeatureCovariance.init(4)
>>> new = cov.update(jnp.ones((8, 4)))
>>> new is cov
False

Source code in .venv/lib/python3.12/site-packages/geonnax/sngp.py

def update(self, features: Float[Array, "B D"]) -> LaplaceRandomFeatureCovariance:
    r"""Return a new container with EMA-updated precision.

    $\hat\Lambda \leftarrow m\,\hat\Lambda + (1-m)\,\Phi^\top\Phi/B$.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
        >>> cov = LaplaceRandomFeatureCovariance.init(4)
        >>> new = cov.update(jnp.ones((8, 4)))
        >>> new is cov
        False
    """
    B = features.shape[0]
    # Feature Gram ΦᵀΦ / B: contract the batch axis b → (D, D).
    outer = einx.dot("b d, b e -> d e", features, features) / B
    new_precision = self.momentum * self.precision + (1.0 - self.momentum) * outer
    return eqx.tree_at(lambda c: c.precision, self, new_precision)

covariance() -> Float[Array, 'D D']

Inverse of the precision matrix (one-shot Cholesky inversion).

\(\hat\Sigma = (\hat\Lambda + \lambda I)^{-1}\), shape (D, D).

Examples:

>>> from geonnax.sngp import LaplaceRandomFeatureCovariance
>>> cov = LaplaceRandomFeatureCovariance.init(4)
>>> cov.covariance().shape
(4, 4)

Source code in .venv/lib/python3.12/site-packages/geonnax/sngp.py

def covariance(self) -> Float[Array, "D D"]:
    r"""Inverse of the precision matrix (one-shot Cholesky inversion).

    $\hat\Sigma = (\hat\Lambda + \lambda I)^{-1}$, shape ``(D, D)``.

    Examples:
        >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
        >>> cov = LaplaceRandomFeatureCovariance.init(4)
        >>> cov.covariance().shape
        (4, 4)
    """
    L = self._chol()
    D = self.precision.shape[0]
    return jax.scipy.linalg.cho_solve((L, True), jnp.eye(D))

variance_at(features: Float[Array, 'N D']) -> Float[Array, ' N']

Per-row predictive variance \(\phi(x_n)^\top \hat{\Sigma}\,\phi(x_n)\).

Computed via a triangular solve to avoid materialising the full \(D \times D\) covariance:

\[ y = L^{-1} \phi(x_n)^\top, \qquad \sigma^2(x_n) = \lVert y \rVert_2^2 = \phi(x_n)^\top (L L^\top)^{-1} \phi(x_n). \]

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.sngp import LaplaceRandomFeatureCovariance
>>> cov = LaplaceRandomFeatureCovariance.init(4)
>>> cov.variance_at(jnp.eye(4)).shape
(4,)

Source code in .venv/lib/python3.12/site-packages/geonnax/sngp.py

def variance_at(self, features: Float[Array, "N D"]) -> Float[Array, " N"]:
    r"""Per-row predictive variance $\phi(x_n)^\top \hat{\Sigma}\,\phi(x_n)$.

    Computed via a triangular solve to avoid materialising the full
    $D \times D$ covariance:

    $$
    y = L^{-1} \phi(x_n)^\top, \qquad
    \sigma^2(x_n) = \lVert y \rVert_2^2
    = \phi(x_n)^\top (L L^\top)^{-1} \phi(x_n).
    $$


    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.sngp import LaplaceRandomFeatureCovariance
        >>> cov = LaplaceRandomFeatureCovariance.init(4)
        >>> cov.variance_at(jnp.eye(4)).shape
        (4,)
    """
    L = self._chol()
    # Solve L y = Φᵀ ; (D, D)\(D, N) -> (D, N), one column per row of Φ.
    y = jax.scipy.linalg.solve_triangular(
        L, einx.id("n d -> d n", features), lower=True
    )
    # σ²(xₙ) = ‖yₙ‖² ; sum over D -> (N,)
    return einx.sum("[d] n", y * y)

Deep spectral GPs

DeepVSSGP

Bases: PyroxModule

Deep Random Feature Expansion for Variational SSGP (Cutajar et al. 2017).

A stack of \(L\) variational SSGP layers, each with random spectral frequencies \(\Omega_l\) and random projection weights \(W_l\):

\[ \begin{aligned} F_0 &= X, \\ F_{l+1} &= \Phi_l(F_l;\, \Omega_l, \ell_l)\, W_l, \quad l = 0, \ldots, L-1, \\ \Phi_l(F;\, \Omega_l, \ell_l) &= \sqrt{1/M}\, [\cos(F\,\Omega_l/\ell_l), \sin(F\,\Omega_l/\ell_l)]. \end{aligned} \]

Each layer registers three sample sites:

• layer_{l}.W_freq — RFF frequencies, prior \(\mathcal{N}(0, 1)\) (RBF spectral density in lengthscale-1 units).
• layer_{l}.lengthscale — kernel lengthscale, prior \(\mathrm{LogNormal}(\log \ell_{\mathrm{init}}, 1)\).
• layer_{l}.W_proj — projection weights, prior \(\mathcal{N}(0, \sigma_W^2)\).

Under SVI an AutoNormal learns mean-field Gaussian posteriors over all \(3L\) sites — one MC sample per forward pass gives the doubly-stochastic reparameterised ELBO of Cutajar et al. (2017).

At depth=1 this reduces to a single VSSGP layer mapping in_features -> out_features via the RFF basis (same model class as VariationalFourierFeatures followed by a DenseReparameterization head). Stacking adds non-stationarity at the cost of a non-Gaussian aggregate likelihood — the layer-wise marginalisation that makes single-layer SSGP closed-form is no longer available, hence the variational treatment.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_{\mathrm{in}}\).
hidden_features	int	Inter-layer dimension \(D_h\) (constant across hidden layers).
out_features	int	Output dimension \(D_{\mathrm{out}}\).
n_features	int	Per-layer Fourier-feature pair count \(M\) (so each layer's hidden state is \(2M\)-dim before projection).
depth	int	Total number of stacked SSGP layers \(L\). Must be \(\ge 1\).
init_lengthscale	float	Prior location for each layer's lengthscale.
prior_std	float	Standard deviation of the per-layer projection prior \(\mathcal{N}(0, \sigma_W^2)\).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.random as jr, jax.numpy as jnp
>>> from numpyro import handlers
>>> net = DeepVSSGP.init(in_features=2, hidden_features=4,
...                       out_features=1, depth=3, n_features=16)
>>> with handlers.seed(rng_seed=0):
...     y = net(jnp.zeros((8, 2)))
>>> y.shape
(8, 1)

Source code in packages/pyrox-nn/src/pyrox_nn/_vssgp.py

class DeepVSSGP(PyroxModule):
    r"""Deep Random Feature Expansion for Variational SSGP (Cutajar et al. 2017).

    A stack of $L$ variational SSGP layers, each with random
    spectral frequencies $\Omega_l$ and random projection weights
    $W_l$:

    $$
    \begin{aligned}
    F_0 &= X, \\
    F_{l+1} &= \Phi_l(F_l;\, \Omega_l, \ell_l)\, W_l,
        \quad l = 0, \ldots, L-1, \\
    \Phi_l(F;\, \Omega_l, \ell_l) &=
        \sqrt{1/M}\,
        [\cos(F\,\Omega_l/\ell_l), \sin(F\,\Omega_l/\ell_l)].
    \end{aligned}
    $$

    Each layer registers three sample sites:

    - ``layer_{l}.W_freq`` — RFF frequencies, prior $\mathcal{N}(0, 1)$
      (RBF spectral density in lengthscale-1 units).
    - ``layer_{l}.lengthscale`` — kernel lengthscale, prior
      $\mathrm{LogNormal}(\log \ell_{\mathrm{init}}, 1)$.
    - ``layer_{l}.W_proj`` — projection weights, prior
      $\mathcal{N}(0, \sigma_W^2)$.

    Under SVI an
    `AutoNormal` learns mean-field
    Gaussian posteriors over all $3L$ sites — one MC sample per
    forward pass gives the doubly-stochastic reparameterised ELBO of
    Cutajar et al. (2017).

    At ``depth=1`` this reduces to a single VSSGP layer mapping
    ``in_features -> out_features`` via the RFF basis (same model class
    as `VariationalFourierFeatures` followed by a
    `DenseReparameterization` head). Stacking adds
    non-stationarity at the cost of a non-Gaussian aggregate likelihood
    — the layer-wise marginalisation that makes single-layer SSGP
    closed-form is no longer available, hence the variational
    treatment.

    Attributes:
        in_features: Input dimension $D_{\mathrm{in}}$.
        hidden_features: Inter-layer dimension $D_h$ (constant
            across hidden layers).
        out_features: Output dimension $D_{\mathrm{out}}$.
        n_features: Per-layer Fourier-feature pair count $M$ (so
            each layer's hidden state is $2M$-dim before
            projection).
        depth: Total number of stacked SSGP layers $L$. Must
            be $\ge 1$.
        init_lengthscale: Prior location for each layer's lengthscale.
        prior_std: Standard deviation of the per-layer projection
            prior $\mathcal{N}(0, \sigma_W^2)$.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.random as jr, jax.numpy as jnp
        >>> from numpyro import handlers
        >>> net = DeepVSSGP.init(in_features=2, hidden_features=4,
        ...                       out_features=1, depth=3, n_features=16)
        >>> with handlers.seed(rng_seed=0):
        ...     y = net(jnp.zeros((8, 2)))
        >>> y.shape
        (8, 1)
    """

    core: DeepVSSGPCore
    init_lengthscale: float = 1.0
    prior_std: float = 1.0
    pyrox_name: str | None = None

    # Structural accessors so wrapper exposes the same field surface as before.
    @property
    def in_features(self) -> int:
        return self.core.in_features

    @property
    def hidden_features(self) -> int:
        return self.core.hidden_features

    @property
    def out_features(self) -> int:
        return self.core.out_features

    @property
    def n_features(self) -> int:
        return self.core.n_features

    @property
    def depth(self) -> int:
        return self.core.depth

    @classmethod
    def init(
        cls,
        in_features: int,
        hidden_features: int,
        out_features: int,
        *,
        depth: int,
        n_features: int = 64,
        lengthscale: float = 1.0,
        prior_std: float = 1.0,
        pyrox_name: str | None = None,
    ) -> DeepVSSGP:
        """Construct a `DeepVSSGP`.

        Args:
            in_features: Input dimension. Must be $\\ge 1$.
            hidden_features: Hidden dimension. Must be $\\ge 1$.
            out_features: Output dimension. Must be $\\ge 1$.
            depth: Total stacked SSGP layers (including readout).
                Must be $\\ge 1$.
            n_features: Per-layer Fourier-feature pair count. Must be
                $\\ge 1$.
            lengthscale: Prior location for each layer's lengthscale.
                Must be $> 0$.
            prior_std: Per-layer projection prior standard deviation.
                Must be $> 0$.
            pyrox_name: Optional explicit scope name for NumPyro site
                registration.

        Returns:
            Initialised `DeepVSSGP`.

        Raises:
            ValueError: If ``depth``, any feature dimension, or
                ``n_features`` is $< 1$, or if ``lengthscale`` /
                ``prior_std`` is $\\le 0$.
        """
        if depth < 1:
            raise ValueError(f"depth must be >= 1, got {depth}.")
        if n_features < 1:
            raise ValueError(f"n_features must be >= 1, got {n_features}.")
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        if prior_std <= 0:
            raise ValueError(f"prior_std must be > 0, got {prior_std}.")
        for name, dim in (
            ("in_features", in_features),
            ("hidden_features", hidden_features),
            ("out_features", out_features),
        ):
            if dim < 1:
                raise ValueError(f"{name} must be >= 1, got {dim}.")
        # The Bayesian wrapper samples every per-layer array on each
        # forward call, so the core's stored arrays are only used as a
        # structural skeleton.  We build it with a dummy PRNG key — its
        # weight values are irrelevant because eqx.tree_at swaps them
        # for sampled values below.
        core = DeepVSSGPCore.init(
            in_features,
            hidden_features,
            out_features,
            depth=depth,
            key=jax.random.PRNGKey(0),
            n_features=n_features,
            lengthscale=lengthscale,
            prior_std=prior_std,
        )
        return cls(
            core=core,
            init_lengthscale=lengthscale,
            prior_std=prior_std,
            pyrox_name=pyrox_name,
        )

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        sampled_W_freqs: list[Float[Array, "d_in n_features"]] = []
        sampled_W_projs: list[Float[Array, "d_rff d_out"]] = []
        sampled_lengthscales: list[Float[Array, ""]] = []
        for layer_idx in range(self.core.depth):
            in_dim = (
                self.core.in_features if layer_idx == 0 else self.core.hidden_features
            )
            out_dim = (
                self.core.out_features
                if layer_idx == self.core.depth - 1
                else self.core.hidden_features
            )
            W_freq = self.pyrox_sample(
                f"layer_{layer_idx}.W_freq",
                dist.Normal(0.0, 1.0)
                .expand([in_dim, self.core.n_features])
                .to_event(2),
            )
            ls = self.pyrox_sample(
                f"layer_{layer_idx}.lengthscale",
                dist.LogNormal(jnp.log(jnp.asarray(self.init_lengthscale)), 1.0),
            )
            W_proj = self.pyrox_sample(
                f"layer_{layer_idx}.W_proj",
                dist.Normal(0.0, self.prior_std)
                .expand([2 * self.core.n_features, out_dim])
                .to_event(2),
            )
            sampled_W_freqs.append(W_freq)
            sampled_W_projs.append(W_proj)
            sampled_lengthscales.append(ls)

        sampled_core = eqx.tree_at(
            lambda c: (c.W_freqs, c.W_projs, c.lengthscales),
            self.core,
            (
                sampled_W_freqs,
                sampled_W_projs,
                jnp.stack(sampled_lengthscales),
            ),
        )
        # Single-example geonnax core over arbitrary leading batch dims.
        return vmap_over_flat_batch(sampled_core, x)

init(in_features: int, hidden_features: int, out_features: int, *, depth: int, n_features: int = 64, lengthscale: float = 1.0, prior_std: float = 1.0, pyrox_name: str | None = None) -> DeepVSSGP classmethod

Construct a DeepVSSGP.

Parameters:

Name	Type	Description	Default
in_features	int	Input dimension. Must be \(\ge 1\).	required
hidden_features	int	Hidden dimension. Must be \(\ge 1\).	required
out_features	int	Output dimension. Must be \(\ge 1\).	required
depth	int	Total stacked SSGP layers (including readout). Must be \(\ge 1\).	required
n_features	int	Per-layer Fourier-feature pair count. Must be \(\ge 1\).	64
lengthscale	float	Prior location for each layer's lengthscale. Must be \(> 0\).	1.0
prior_std	float	Per-layer projection prior standard deviation. Must be \(> 0\).	1.0
pyrox_name	str | None	Optional explicit scope name for NumPyro site registration.	None

Returns:

Type	Description
DeepVSSGP	Initialised DeepVSSGP.

Raises:

Type	Description
ValueError	If depth, any feature dimension, or n_features is \(< 1\), or if lengthscale / prior_std is \(\le 0\).

Source code in packages/pyrox-nn/src/pyrox_nn/_vssgp.py

@classmethod
def init(
    cls,
    in_features: int,
    hidden_features: int,
    out_features: int,
    *,
    depth: int,
    n_features: int = 64,
    lengthscale: float = 1.0,
    prior_std: float = 1.0,
    pyrox_name: str | None = None,
) -> DeepVSSGP:
    """Construct a `DeepVSSGP`.

    Args:
        in_features: Input dimension. Must be $\\ge 1$.
        hidden_features: Hidden dimension. Must be $\\ge 1$.
        out_features: Output dimension. Must be $\\ge 1$.
        depth: Total stacked SSGP layers (including readout).
            Must be $\\ge 1$.
        n_features: Per-layer Fourier-feature pair count. Must be
            $\\ge 1$.
        lengthscale: Prior location for each layer's lengthscale.
            Must be $> 0$.
        prior_std: Per-layer projection prior standard deviation.
            Must be $> 0$.
        pyrox_name: Optional explicit scope name for NumPyro site
            registration.

    Returns:
        Initialised `DeepVSSGP`.

    Raises:
        ValueError: If ``depth``, any feature dimension, or
            ``n_features`` is $< 1$, or if ``lengthscale`` /
            ``prior_std`` is $\\le 0$.
    """
    if depth < 1:
        raise ValueError(f"depth must be >= 1, got {depth}.")
    if n_features < 1:
        raise ValueError(f"n_features must be >= 1, got {n_features}.")
    if lengthscale <= 0:
        raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
    if prior_std <= 0:
        raise ValueError(f"prior_std must be > 0, got {prior_std}.")
    for name, dim in (
        ("in_features", in_features),
        ("hidden_features", hidden_features),
        ("out_features", out_features),
    ):
        if dim < 1:
            raise ValueError(f"{name} must be >= 1, got {dim}.")
    # The Bayesian wrapper samples every per-layer array on each
    # forward call, so the core's stored arrays are only used as a
    # structural skeleton.  We build it with a dummy PRNG key — its
    # weight values are irrelevant because eqx.tree_at swaps them
    # for sampled values below.
    core = DeepVSSGPCore.init(
        in_features,
        hidden_features,
        out_features,
        depth=depth,
        key=jax.random.PRNGKey(0),
        n_features=n_features,
        lengthscale=lengthscale,
        prior_std=prior_std,
    )
    return cls(
        core=core,
        init_lengthscale=lengthscale,
        prior_std=prior_std,
        pyrox_name=pyrox_name,
    )

Ensembles — BatchEnsemble / rank-1

Efficient deep ensembles that share one weight matrix and learn per-member rank-1 perturbations (Wen et al., 2020; Dusenberry et al., 2020).

DenseRank1

Bases: PyroxModule

Rank-1 ensemble dense layer.

Implements the BatchEnsemble (Wen et al., 2020) / rank-1 BNN (Dusenberry et al., 2020) parameterization: a single shared kernel \(W \in \mathbb{R}^{D_\mathrm{in} \times D_\mathrm{out}}\) and per-member rank-1 multiplicative perturbations \(s_i \in \mathbb{R}^{D_\mathrm{in}}\), \(r_i \in \mathbb{R}^{D_\mathrm{out}}\) for \(i = 1, \ldots, M\). The per-member effective weight is

\[ W_i = (s_i \otimes r_i) \circ W, \]

and the efficient forward pass avoids materialising \(W_i\):

\[ y_i = \bigl((x \circ s_i)\, W\bigr) \circ r_i + b_i. \]

Two modes via the bayesian flag:

• bayesian=False (default) — BatchEnsemble. \(r, s, W, b\) are all deterministic pyrox_param sites and per-member diversity comes purely from the random initialisation of \(r_i, s_i\). Use this for ensemble training under a single shared SGD trajectory.
• bayesian=True — rank-1 BNN. \(r, s\) are pyrox_sample sites with Normal priors centered at the per-member init values; \(W, b\) remain deterministic. Plug into NumPyro's SVI machinery (an AutoNormal guide on r, s recovers Dusenberry et al., 2020).

Plate semantics

Identical to other pyrox Bayesian dense layers — call this layer outside numpyro.plate("data", ..., subsample_size=...) and only plate the observation likelihood. The model log density picks up \(\log p(r_i)\) and \(\log p(s_i)\) once per layer (not once per example) under the canonical pattern.

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_\mathrm{in}\).
out_features	int	Output dimension \(D_\mathrm{out}\).
ensemble_size	int	Number of ensemble members \(M\).
bias	bool	Whether to include a per-member bias.
bayesian	bool	If True, place Normal priors on \(r, s\).
prior_scale	float	Std of the Bayesian priors on \(r, s\). Only used when bayesian=True.
W_init	float	Shared kernel init, shape (D_in, D_out).
r_init	float	Per-member output-side init, shape (M, D_out).
s_init	float	Per-member input-side init, shape (M, D_in).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.random as jr
>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> layer = DenseRank1.init(
...     jr.PRNGKey(0),
...     in_features=4,
...     out_features=2,
...     ensemble_size=3,
... )
>>> x = jnp.ones((5, 4))
>>> with handlers.seed(rng_seed=0):
...     y = layer(x)
>>> y.shape
(3, 5, 2)

References

Wen, Y., Tran, D., & Ba, J. (2020). BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning. ICLR.

Dusenberry, M. W., et al. (2020). Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors. ICML.

Source code in packages/pyrox-nn/src/pyrox_nn/_ensemble.py

class DenseRank1(PyroxModule):
    r"""Rank-1 ensemble dense layer.

    Implements the BatchEnsemble (Wen et al., 2020) / rank-1 BNN
    (Dusenberry et al., 2020) parameterization: a single shared kernel
    $W \in \mathbb{R}^{D_\mathrm{in} \times D_\mathrm{out}}$ and
    per-member rank-1 multiplicative perturbations
    $s_i \in \mathbb{R}^{D_\mathrm{in}}$,
    $r_i \in \mathbb{R}^{D_\mathrm{out}}$ for
    $i = 1, \ldots, M$. The per-member effective weight is

    $$
    W_i = (s_i \otimes r_i) \circ W,
    $$

    and the efficient forward pass avoids materialising $W_i$:

    $$
    y_i = \bigl((x \circ s_i)\, W\bigr) \circ r_i + b_i.
    $$

    Two modes via the ``bayesian`` flag:

    * ``bayesian=False`` (default) — BatchEnsemble. $r, s, W, b$
      are all deterministic ``pyrox_param`` sites and per-member
      diversity comes purely from the random initialisation of
      $r_i, s_i$. Use this for ensemble training under a single
      shared SGD trajectory.
    * ``bayesian=True`` — rank-1 BNN. $r, s$ are
      ``pyrox_sample`` sites with Normal priors centered at the
      per-member init values; $W, b$ remain deterministic. Plug
      into NumPyro's SVI machinery (an ``AutoNormal`` guide on
      ``r, s`` recovers Dusenberry et al., 2020).

    Plate semantics:
        Identical to other pyrox Bayesian dense layers — call this
        layer **outside** ``numpyro.plate("data", ..., subsample_size=...)``
        and only plate the observation likelihood. The model log
        density picks up $\log p(r_i)$ and $\log p(s_i)$
        once per layer (not once per example) under the canonical
        pattern.

    Attributes:
        in_features: Input dimension $D_\mathrm{in}$.
        out_features: Output dimension $D_\mathrm{out}$.
        ensemble_size: Number of ensemble members $M$.
        bias: Whether to include a per-member bias.
        bayesian: If ``True``, place Normal priors on $r, s$.
        prior_scale: Std of the Bayesian priors on $r, s$. Only
            used when ``bayesian=True``.
        W_init: Shared kernel init, shape ``(D_in, D_out)``.
        r_init: Per-member output-side init, shape ``(M, D_out)``.
        s_init: Per-member input-side init, shape ``(M, D_in)``.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.random as jr
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> layer = DenseRank1.init(
        ...     jr.PRNGKey(0),
        ...     in_features=4,
        ...     out_features=2,
        ...     ensemble_size=3,
        ... )
        >>> x = jnp.ones((5, 4))
        >>> with handlers.seed(rng_seed=0):
        ...     y = layer(x)
        >>> y.shape
        (3, 5, 2)

    References:
        Wen, Y., Tran, D., & Ba, J. (2020). *BatchEnsemble: An
        Alternative Approach to Efficient Ensemble and Lifelong
        Learning.* ICLR.

        Dusenberry, M. W., et al. (2020). *Efficient and Scalable
        Bayesian Neural Nets with Rank-1 Factors.* ICML.
    """

    core: geonnax.DenseRank1
    bayesian: bool = eqx.field(static=True, default=False)
    prior_scale: float = 0.5
    pyrox_name: str | None = None

    @classmethod
    def init(
        cls,
        key: PRNGKeyArray,
        in_features: int,
        out_features: int,
        ensemble_size: int,
        *,
        bias: bool = True,
        bayesian: bool = False,
        init_scale: float = 0.5,
        prior_scale: float = 0.5,
        pyrox_name: str | None = None,
    ) -> DenseRank1:
        """Construct a layer with random per-member init vectors."""
        # `prior_scale` is only consulted in Bayesian mode — don't reject
        # configs that pass through a sentinel default in deterministic mode.
        if bayesian and prior_scale <= 0:
            raise ValueError(
                f"prior_scale must be > 0 when bayesian=True; got {prior_scale}."
            )
        # geonnax handles positive-dim / init_scale validation.
        core = geonnax.DenseRank1.init(
            key,
            in_features=in_features,
            out_features=out_features,
            ensemble_size=ensemble_size,
            bias=bias,
            init_scale=init_scale,
        )
        return cls(
            core=core,
            bayesian=bayesian,
            prior_scale=prior_scale,
            pyrox_name=pyrox_name,
        )

    # Convenience read-only attribute access.
    @property
    def in_features(self) -> int:
        return self.core.in_features

    @property
    def out_features(self) -> int:
        return self.core.out_features

    @property
    def ensemble_size(self) -> int:
        return self.core.ensemble_size

    @property
    def bias(self) -> bool:
        return self.core.bias

    @pyrox_method
    def __call__(
        self, x: Float[Array, "*batch D_in"]
    ) -> Float[Array, "M *batch D_out"]:
        W = self.pyrox_param("W", self.core.W)

        if self.bayesian:
            r = self.pyrox_sample(
                "r",
                dist.Normal(self.core.r, self.prior_scale).to_event(2),
            )
            s = self.pyrox_sample(
                "s",
                dist.Normal(self.core.s, self.prior_scale).to_event(2),
            )
        else:
            r = self.pyrox_param("r", self.core.r)
            s = self.pyrox_param("s", self.core.s)

        b = self.pyrox_param("b", self.core.b) if self.bias else self.core.b

        new_core = eqx.tree_at(
            lambda c: (c.W, c.r, c.s, c.b),
            self.core,
            (W, r, s, b),
        )
        return _vmap_collapsed(new_core, x)

init(key: PRNGKeyArray, in_features: int, out_features: int, ensemble_size: int, *, bias: bool = True, bayesian: bool = False, init_scale: float = 0.5, prior_scale: float = 0.5, pyrox_name: str | None = None) -> DenseRank1 classmethod

Construct a layer with random per-member init vectors.

Source code in packages/pyrox-nn/src/pyrox_nn/_ensemble.py

@classmethod
def init(
    cls,
    key: PRNGKeyArray,
    in_features: int,
    out_features: int,
    ensemble_size: int,
    *,
    bias: bool = True,
    bayesian: bool = False,
    init_scale: float = 0.5,
    prior_scale: float = 0.5,
    pyrox_name: str | None = None,
) -> DenseRank1:
    """Construct a layer with random per-member init vectors."""
    # `prior_scale` is only consulted in Bayesian mode — don't reject
    # configs that pass through a sentinel default in deterministic mode.
    if bayesian and prior_scale <= 0:
        raise ValueError(
            f"prior_scale must be > 0 when bayesian=True; got {prior_scale}."
        )
    # geonnax handles positive-dim / init_scale validation.
    core = geonnax.DenseRank1.init(
        key,
        in_features=in_features,
        out_features=out_features,
        ensemble_size=ensemble_size,
        bias=bias,
        init_scale=init_scale,
    )
    return cls(
        core=core,
        bayesian=bayesian,
        prior_scale=prior_scale,
        pyrox_name=pyrox_name,
    )

LayerNormEnsemble

Bases: PyroxModule

Per-ensemble-member LayerNorm.

Drop-in replacement for LayerNorm inside BatchEnsemble / Rank1 architectures. Computes the standard LayerNorm normalisation over the trailing feature dimension and applies a per-member affine transform — each ensemble member \(i \in \{1, \ldots, M\}\) gets its own learnable scale \(\gamma_i \in \mathbb{R}^D\) and bias \(\beta_i \in \mathbb{R}^D\):

\[ \hat{x}_i = \frac{x_i - \mu(x_i)}{\sqrt{\sigma^2(x_i) + \epsilon}}, \qquad y_i = \gamma_i \odot \hat{x}_i + \beta_i, \]

where \(\mu\) and \(\sigma^2\) are the empirical mean and variance over the trailing feature axis (computed independently for each member-batch slice). Without per-member scale/bias, sharing a single LayerNorm across the ensemble would couple all members and erase the diversity introduced by DenseRank1 or any other BatchEnsemble layer upstream.

Input is expected to carry a leading ensemble axis of size ensemble_size and a trailing feature axis of size feature_dim. Any number of intermediate batch / time axes are supported and pass through unchanged.

Attributes:

Name	Type	Description
ensemble_size	int	Number of ensemble members \(M\).
feature_dim	int	Trailing feature dimension \(D\) over which the normalisation is computed.
eps	float	Small positive constant added to the variance for numerical stability.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> ln = LayerNormEnsemble(
...     ensemble_size=3, feature_dim=4, pyrox_name="ln"
... )
>>> x = jnp.ones((3, 5, 4))  # (M, batch, D)
>>> with handlers.seed(rng_seed=0):
...     y = ln(x)
>>> y.shape
(3, 5, 4)

Source code in packages/pyrox-nn/src/pyrox_nn/_ensemble.py

class LayerNormEnsemble(PyroxModule):
    r"""Per-ensemble-member LayerNorm.

    Drop-in replacement for ``LayerNorm`` inside BatchEnsemble / Rank1
    architectures. Computes the standard LayerNorm normalisation over
    the trailing feature dimension and applies a *per-member* affine
    transform — each ensemble member $i \in \{1, \ldots, M\}$
    gets its own learnable scale $\gamma_i \in \mathbb{R}^D$
    and bias $\beta_i \in \mathbb{R}^D$:

    $$
    \hat{x}_i = \frac{x_i - \mu(x_i)}{\sqrt{\sigma^2(x_i) + \epsilon}},
    \qquad
    y_i = \gamma_i \odot \hat{x}_i + \beta_i,
    $$

    where $\mu$ and $\sigma^2$ are the empirical mean and
    variance over the trailing feature axis (computed independently
    for each member-batch slice). Without per-member scale/bias,
    sharing a single LayerNorm across the ensemble would couple all
    members and erase the diversity introduced by `DenseRank1`
    or any other BatchEnsemble layer upstream.

    Input is expected to carry a leading ensemble axis of size
    ``ensemble_size`` and a trailing feature axis of size
    ``feature_dim``. Any number of intermediate batch / time axes
    are supported and pass through unchanged.

    Attributes:
        ensemble_size: Number of ensemble members $M$.
        feature_dim: Trailing feature dimension $D$ over which
            the normalisation is computed.
        eps: Small positive constant added to the variance for
            numerical stability.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> ln = LayerNormEnsemble(
        ...     ensemble_size=3, feature_dim=4, pyrox_name="ln"
        ... )
        >>> x = jnp.ones((3, 5, 4))  # (M, batch, D)
        >>> with handlers.seed(rng_seed=0):
        ...     y = ln(x)
        >>> y.shape
        (3, 5, 4)
    """

    core: geonnax.LayerNormEnsemble
    pyrox_name: str | None = eqx.field(static=True, default=None)

    def __init__(
        self,
        ensemble_size: int,
        feature_dim: int,
        eps: float = 1e-5,
        pyrox_name: str | None = None,
        *,
        core: geonnax.LayerNormEnsemble | None = None,
    ) -> None:
        # Keep the prior call signature (positional dims, kw eps / pyrox_name)
        # so existing callers don't need to know about the geonnax core
        # construction.
        if core is None:
            # geonnax validates positive ensemble_size / feature_dim / eps.
            core = geonnax.LayerNormEnsemble.init(
                ensemble_size=ensemble_size,
                feature_dim=feature_dim,
                eps=eps,
            )
        self.core = core
        self.pyrox_name = pyrox_name

    @property
    def ensemble_size(self) -> int:
        return self.core.ensemble_size

    @property
    def feature_dim(self) -> int:
        return self.core.feature_dim

    @property
    def eps(self) -> float:
        return self.core.eps

    @pyrox_method
    def __call__(self, x: Float[Array, "M *batch D"]) -> Float[Array, "M *batch D"]:
        if x.ndim < 2:
            raise ValueError(
                f"x must have at least 2 dims (M and D); got shape {x.shape}."
            )
        if x.shape[0] != self.ensemble_size:
            raise ValueError(
                f"x.shape[0] = {x.shape[0]} does not match "
                f"ensemble_size = {self.ensemble_size}."
            )
        if x.shape[-1] != self.feature_dim:
            raise ValueError(
                f"x.shape[-1] = {x.shape[-1]} does not match "
                f"feature_dim = {self.feature_dim}."
            )

        scales = self.pyrox_param("scales", self.core.scales)
        biases = self.pyrox_param("biases", self.core.biases)

        new_core = eqx.tree_at(
            lambda c: (c.scales, c.biases),
            self.core,
            (scales, biases),
        )

        # Core takes (M, D) → (M, D); collapse arbitrary intermediate batch
        # dims to a single leading axis, vmap over it, then restore. The
        # ensemble axis stays at position 0 because the core is per-member.
        batch_shape = x.shape[1:-1]
        flat = einx.id("m b... d -> m (b...) d", x)  # (M, B, D)
        # vmap over the B axis (position 1 on both input and output).
        out = jax.vmap(new_core, in_axes=1, out_axes=1)(flat)  # (M, B, D)
        return einx.id("m (b...) d -> m b... d", out, b=batch_shape)

MultiHeadAttentionBE

Bases: PyroxModule

Multi-head attention with BatchEnsemble rank-1 projections.

Standard scaled-dot-product multi-head attention where each of the four linear projections — query, key, value, and output — uses a BatchEnsemble parameterisation: a shared full-rank kernel plus per-ensemble-member rank-1 multiplicative perturbations. So for member \(i \in \{1, \ldots, M\}\) and projection \(P \in \{Q, K, V, O\}\),

\[ W_i^{(P)} = (s_i^{(P)} \otimes r_i^{(P)}) \circ W^{(P)}, \]

and the attention itself is the usual

\[ \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\! \Bigl(\frac{Q K^\top}{\sqrt{d_k}}\Bigr) V. \]

The forward consumes un-ensembled inputs (query, key, value of shape (T, D) / (S, D)), adds the ensemble axis when projecting to Q, K, V, runs per-member attention in parallel, and returns the per-member output of shape (M, T, D). Equivalent to running M independent attention heads with rank-1 weight perturbations and stacking their outputs.

Plate semantics

Same convention as DenseRank1 and the rest of the pyrox_nn ensemble / Bayesian dense family — call this layer outside numpyro.plate("data", ..., subsample_size=...) and only plate the observation likelihood. All four projections register their parameters as pyrox_param sites; nothing about the layer is data-dependent so plate-scaling does not come into play unless the user puts the call inside a subsampled plate.

Attributes:

Name	Type	Description
embed_dim	int	Total feature dimension \(D\) of query / key / value (must be divisible by num_heads).
num_heads	int	Number of attention heads \(H\). Each head sees embed_dim // num_heads features.
ensemble_size	int	Number of ensemble members \(M\).
bias	bool	Whether each of the four projections includes a per-member bias. When False, no bias param sites are registered for any of Q / K / V / O.
q_init	/ k_init / v_init / o_init	Per-projection BatchEnsemble init arrays. Build via init.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.random as jr
>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> mha = MultiHeadAttentionBE.init(
...     jr.PRNGKey(0),
...     embed_dim=8, num_heads=2, ensemble_size=3,
... )
>>> x = jnp.ones((5, 8))
>>> with handlers.seed(rng_seed=0):
...     y = mha(x, x, x)         # self-attention
>>> y.shape
(3, 5, 8)

Source code in packages/pyrox-nn/src/pyrox_nn/_ensemble.py

class MultiHeadAttentionBE(PyroxModule):
    r"""Multi-head attention with BatchEnsemble rank-1 projections.

    Standard scaled-dot-product multi-head attention where each of the
    four linear projections — query, key, value, and output — uses a
    BatchEnsemble parameterisation: a shared full-rank kernel plus
    per-ensemble-member rank-1 multiplicative perturbations. So for
    member $i \in \{1, \ldots, M\}$ and projection
    $P \in \{Q, K, V, O\}$,

    $$
    W_i^{(P)} = (s_i^{(P)} \otimes r_i^{(P)}) \circ W^{(P)},
    $$

    and the attention itself is the usual

    $$
    \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\!
        \Bigl(\frac{Q K^\top}{\sqrt{d_k}}\Bigr) V.
    $$

    The forward consumes un-ensembled inputs (``query``, ``key``,
    ``value`` of shape ``(T, D)`` / ``(S, D)``), adds the ensemble
    axis when projecting to ``Q``, ``K``, ``V``, runs per-member
    attention in parallel, and returns the per-member output of
    shape ``(M, T, D)``. Equivalent to running ``M`` independent
    attention heads with rank-1 weight perturbations and stacking
    their outputs.

    Plate semantics:
        Same convention as `DenseRank1` and the rest of the
        ``pyrox_nn`` ensemble / Bayesian dense family — call this
        layer **outside** ``numpyro.plate("data", ..., subsample_size=...)``
        and only plate the observation likelihood. All four projections
        register their parameters as ``pyrox_param`` sites; nothing
        about the layer is data-dependent so plate-scaling does not
        come into play unless the user puts the call inside a
        subsampled plate.

    Attributes:
        embed_dim: Total feature dimension $D$ of query / key /
            value (must be divisible by ``num_heads``).
        num_heads: Number of attention heads $H$. Each head sees
            ``embed_dim // num_heads`` features.
        ensemble_size: Number of ensemble members $M$.
        bias: Whether each of the four projections includes a
            per-member bias. When ``False``, no bias param sites are
            registered for any of Q / K / V / O.
        q_init / k_init / v_init / o_init: Per-projection
            BatchEnsemble init arrays. Build via `init`.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.random as jr
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> mha = MultiHeadAttentionBE.init(
        ...     jr.PRNGKey(0),
        ...     embed_dim=8, num_heads=2, ensemble_size=3,
        ... )
        >>> x = jnp.ones((5, 8))
        >>> with handlers.seed(rng_seed=0):
        ...     y = mha(x, x, x)         # self-attention
        >>> y.shape
        (3, 5, 8)
    """

    core: geonnax.MultiHeadAttentionBE
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @classmethod
    def init(
        cls,
        key: PRNGKeyArray,
        embed_dim: int,
        num_heads: int,
        ensemble_size: int,
        *,
        bias: bool = True,
        init_scale: float = 0.5,
        pyrox_name: str | None = None,
    ) -> MultiHeadAttentionBE:
        """Construct an MHA-BE layer with random Q/K/V/O projection inits."""
        # geonnax validates positive dims, divisibility, init_scale.
        core = geonnax.MultiHeadAttentionBE.init(
            key,
            embed_dim=embed_dim,
            num_heads=num_heads,
            ensemble_size=ensemble_size,
            bias=bias,
            init_scale=init_scale,
        )
        return cls(core=core, pyrox_name=pyrox_name)

    @property
    def embed_dim(self) -> int:
        return self.core.embed_dim

    @property
    def num_heads(self) -> int:
        return self.core.num_heads

    @property
    def ensemble_size(self) -> int:
        return self.core.ensemble_size

    @property
    def bias(self) -> bool:
        return self.core.bias

    @property
    def head_dim(self) -> int:
        return self.core.head_dim

    # Maintain the pre-refactor attribute names for any external accessors.
    @property
    def q_init(self) -> Rank1ProjInit:
        return self.core.q_proj

    @property
    def k_init(self) -> Rank1ProjInit:
        return self.core.k_proj

    @property
    def v_init(self) -> Rank1ProjInit:
        return self.core.v_proj

    @property
    def o_init(self) -> Rank1ProjInit:
        return self.core.o_proj

    def _register_proj(self, name: str, init: Rank1ProjInit) -> Rank1ProjInit:
        """Register a projection's arrays as ``pyrox_param`` sites.

        Skips the ``b`` site when ``self.bias`` is ``False`` so disabled
        biases don't leak unused params into the SVI parameter store.
        """
        b = (
            self.pyrox_param(f"{name}_b", init.b)
            if self.bias
            else jnp.zeros_like(init.b)
        )
        return Rank1ProjInit(
            W=self.pyrox_param(f"{name}_W", init.W),
            r=self.pyrox_param(f"{name}_r", init.r),
            s=self.pyrox_param(f"{name}_s", init.s),
            b=b,
        )

    @pyrox_method
    def __call__(
        self,
        query: Float[Array, "T D"],
        key: Float[Array, "S D"],
        value: Float[Array, "S D"],
    ) -> Float[Array, "M T D"]:
        # Register the four projections.
        q_proj = self._register_proj("q", self.core.q_proj)
        k_proj = self._register_proj("k", self.core.k_proj)
        v_proj = self._register_proj("v", self.core.v_proj)
        o_proj = self._register_proj("o", self.core.o_proj)

        # Swap all four projections into the core; geonnax's __call__
        # validates query/key/value shapes and runs the per-member
        # attention. The core forward matches the pyrox call signature
        # (un-ensembled `(T, D)` / `(S, D)` in → `(M, T, D)` out), so no
        # vmap is needed here — the ensemble axis is intrinsic to the
        # core, not a data batch.
        new_core = eqx.tree_at(
            lambda c: (c.q_proj, c.k_proj, c.v_proj, c.o_proj),
            self.core,
            (q_proj, k_proj, v_proj, o_proj),
        )
        return new_core(query, key, value)

init(key: PRNGKeyArray, embed_dim: int, num_heads: int, ensemble_size: int, *, bias: bool = True, init_scale: float = 0.5, pyrox_name: str | None = None) -> MultiHeadAttentionBE classmethod

Construct an MHA-BE layer with random Q/K/V/O projection inits.

Source code in packages/pyrox-nn/src/pyrox_nn/_ensemble.py

@classmethod
def init(
    cls,
    key: PRNGKeyArray,
    embed_dim: int,
    num_heads: int,
    ensemble_size: int,
    *,
    bias: bool = True,
    init_scale: float = 0.5,
    pyrox_name: str | None = None,
) -> MultiHeadAttentionBE:
    """Construct an MHA-BE layer with random Q/K/V/O projection inits."""
    # geonnax validates positive dims, divisibility, init_scale.
    core = geonnax.MultiHeadAttentionBE.init(
        key,
        embed_dim=embed_dim,
        num_heads=num_heads,
        ensemble_size=ensemble_size,
        bias=bias,
        init_scale=init_scale,
    )
    return cls(core=core, pyrox_name=pyrox_name)

Heteroscedastic output heads

Monte-Carlo sigmoid / softmax output layers with factor-analysis noise (Collier et al., 2021) for input-dependent label noise, and a Bayesian regression head that returns a mean and an input-dependent log standard deviation.

GaussianRegressionHead

Bases: PyroxModule

Bayesian heteroscedastic regression head (mean plus log-scale).

Wraps geonnax.GaussianRegressionHead and samples its four weight arrays. For a feature vector \(h \in \mathbb{R}^{D_\mathrm{in}}\),

\[ \mu = W_\mu^\top h + b_\mu, \qquad \log\sigma = \operatorname{clip}\bigl(W_s^\top h + b_s,\ -7,\ 3\bigr), \qquad y \sim \mathcal{N}\bigl(\mu, \sigma^2\bigr), \]

with independent Gaussian priors

\[ W_\mu \sim \mathcal{N}(0, s^2), \qquad b_\mu \sim \mathcal{N}(0, s^2), \qquad W_s \sim \mathcal{N}(0, s_\sigma^2), \qquad b_s \sim \mathcal{N}(\log\sigma_\mathrm{init}, s^2), \]

where \(s\) is prior_std and \(s_\sigma\) is log_scale_prior_std. Under the prior, \(W_s^\top h\) has standard deviation \(s_\sigma \lVert h \rVert\) per output, so \(s_\sigma\) sets how far \(\sigma\) may vary across inputs a priori. With \(s_\sigma = s = 1\) and 15 RBF features on \([-1, 1]\), \(\log\sigma\) at \(x = 0\) and \(x = 0.95\) differ by about 1.4 nats at one prior sd (a 4-fold change in \(\sigma\)), and some draws reach the clip: a strongly heteroscedastic prior. The default \(s_\sigma = 0.1\) gives a log-scale sd of \(0.1\,\lVert h \rVert\), i.e. about 10 % variation in \(\sigma\) per unit of feature norm, so the prior says "roughly constant noise" and leaves input dependence to the data; the overall noise level, \(b_s\), keeps the broad \(s\). Rescale \(s_\sigma\) with the norm of your features.

Only the prior median is homoscedastic: the wrapped core holds the prior locations (zero weights, \(b_s = \log\sigma_\mathrm{init}\)), so init_to_median (or substituting the core's arrays) starts inference from a head that predicts \(\mu = 0\) and \(\sigma = \sigma_\mathrm{init}\) everywhere; individual prior draws are mildly heteroscedastic.

The layer returns \((\mu, \log\sigma)\); the caller writes the likelihood. With a known per-observation error \(\sigma_i\), use \(\sigma_i^2 + \sigma(h)^2\) as the variance.

Sites (scope pyrox_name, default "GaussianRegressionHead"): <scope>.W_loc \((D_\mathrm{in}, K)\), <scope>.b_loc \((K,)\), <scope>.W_log_scale \((D_\mathrm{in}, K)\) and <scope>.b_log_scale \((K,)\), each sampled once per model call.

Plate semantics

Call outside numpyro.plate("data", ..., subsample_size=...) so the weight sites are not scaled by the subsample factor; put the observation site inside the plate.

Attributes:

Name	Type	Description
core	GaussianRegressionHead	The deterministic geonnax head holding the prior locations.
prior_std	float	Prior standard deviation \(s\) of \(W_\mu\), \(b_\mu\), \(b_s\).
log_scale_prior_std	float	Prior standard deviation \(s_\sigma\) of \(W_s\).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> head = GaussianRegressionHead.init(in_features=8, out_features=2)
>>> with handlers.seed(rng_seed=0):
...     mu, log_sigma = head(jnp.ones((5, 8)))
>>> mu.shape, log_sigma.shape
((5, 2), (5, 2))

References

Nix, D. A., & Weigend, A. S. (1994). Estimating the mean and variance of the target probability distribution. ICNN.

Source code in packages/pyrox-nn/src/pyrox_nn/_heteroscedastic.py

class GaussianRegressionHead(PyroxModule):
    r"""Bayesian heteroscedastic regression head (mean plus log-scale).

    Wraps `geonnax.GaussianRegressionHead` and samples its four weight
    arrays. For a feature vector $h \in \mathbb{R}^{D_\mathrm{in}}$,

    $$
    \mu = W_\mu^\top h + b_\mu, \qquad
    \log\sigma = \operatorname{clip}\bigl(W_s^\top h + b_s,\ -7,\ 3\bigr),
    \qquad y \sim \mathcal{N}\bigl(\mu, \sigma^2\bigr),
    $$

    with independent Gaussian priors

    $$
    W_\mu \sim \mathcal{N}(0, s^2), \qquad
    b_\mu \sim \mathcal{N}(0, s^2), \qquad
    W_s \sim \mathcal{N}(0, s_\sigma^2), \qquad
    b_s \sim \mathcal{N}(\log\sigma_\mathrm{init}, s^2),
    $$

    where $s$ is ``prior_std`` and $s_\sigma$ is ``log_scale_prior_std``.
    Under the prior, $W_s^\top h$ has standard deviation
    $s_\sigma \lVert h \rVert$ per output, so $s_\sigma$ sets how far
    $\sigma$ may vary across inputs *a priori*. With $s_\sigma = s = 1$ and
    15 RBF features on $[-1, 1]$, $\log\sigma$ at $x = 0$ and $x = 0.95$
    differ by about 1.4 nats at one prior sd (a 4-fold change in $\sigma$),
    and some draws reach the clip: a strongly heteroscedastic prior. The
    default $s_\sigma = 0.1$ gives a
    log-scale sd of $0.1\,\lVert h \rVert$, i.e. about 10 % variation in
    $\sigma$ per unit of feature norm, so the prior says "roughly constant
    noise" and leaves input dependence to the data; the overall noise
    level, $b_s$, keeps the broad $s$. Rescale $s_\sigma$ with the norm of
    your features.

    Only the prior *median* is homoscedastic: the wrapped core holds the
    prior locations (zero weights, $b_s = \log\sigma_\mathrm{init}$), so
    ``init_to_median`` (or substituting the core's arrays) starts inference
    from a head that predicts $\mu = 0$ and $\sigma = \sigma_\mathrm{init}$
    everywhere; individual prior draws are mildly heteroscedastic.

    The layer returns $(\mu, \log\sigma)$; the caller writes the
    likelihood. With a known per-observation error $\sigma_i$, use
    $\sigma_i^2 + \sigma(h)^2$ as the variance.

    Sites (scope ``pyrox_name``, default ``"GaussianRegressionHead"``):
    ``<scope>.W_loc`` $(D_\mathrm{in}, K)$, ``<scope>.b_loc`` $(K,)$,
    ``<scope>.W_log_scale`` $(D_\mathrm{in}, K)$ and
    ``<scope>.b_log_scale`` $(K,)$, each sampled once per model call.

    Plate semantics:
        Call outside ``numpyro.plate("data", ..., subsample_size=...)`` so
        the weight sites are not scaled by the subsample factor; put the
        observation site inside the plate.

    Attributes:
        core: The deterministic geonnax head holding the prior locations.
        prior_std: Prior standard deviation $s$ of $W_\mu$, $b_\mu$, $b_s$.
        log_scale_prior_std: Prior standard deviation $s_\sigma$ of $W_s$.
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> head = GaussianRegressionHead.init(in_features=8, out_features=2)
        >>> with handlers.seed(rng_seed=0):
        ...     mu, log_sigma = head(jnp.ones((5, 8)))
        >>> mu.shape, log_sigma.shape
        ((5, 2), (5, 2))

    References:
        Nix, D. A., & Weigend, A. S. (1994). *Estimating the mean and
        variance of the target probability distribution.* ICNN.
    """

    core: geonnax.GaussianRegressionHead
    prior_std: float = eqx.field(static=True, default=1.0)
    log_scale_prior_std: float = eqx.field(static=True, default=0.1)
    pyrox_name: str | None = eqx.field(static=True, default=None)

    @classmethod
    def init(
        cls,
        in_features: int,
        out_features: int = 1,
        *,
        prior_std: float = 1.0,
        log_scale_prior_std: float = 0.1,
        sigma_init: float = 1.0,
        pyrox_name: str | None = None,
    ) -> Self:
        r"""Build the head with priors centred on a homoscedastic core.

        Args:
            in_features: Feature dimension $D_\mathrm{in}$; must be > 0.
            out_features: Number of targets $K$; must be > 0.
            prior_std: Prior standard deviation $s$ of the mean weights and
                both biases; must be > 0.
            log_scale_prior_std: Prior standard deviation $s_\sigma$ of the
                log-scale weights $W_s$; must be > 0. See the class
                docstring for the default.
            sigma_init: Prior location of $\sigma$, i.e.
                $b_s \sim \mathcal{N}(\log\sigma_\mathrm{init}, s^2)$;
                $\log\sigma_\mathrm{init}$ must lie strictly inside the
                core's clip range $(-7, 3)$.
            pyrox_name: Explicit scope name for the sites.

        Returns:
            The head.

        Raises:
            ValueError: If a dimension, ``prior_std`` or
                ``log_scale_prior_std`` is not positive, or ``sigma_init`` is
                not positive or puts $\log\sigma_\mathrm{init}$ on or outside
                the clip range.
        """
        if in_features <= 0 or out_features <= 0:
            raise ValueError(
                "in_features and out_features must be > 0; "
                f"got {in_features=}, {out_features=}."
            )
        for name, value in (
            ("prior_std", prior_std),
            ("log_scale_prior_std", log_scale_prior_std),
            ("sigma_init", sigma_init),
        ):
            if not value > 0:
                raise ValueError(f"{name} must be > 0; got {value}.")
        # The weights are sampled, so the core only carries the prior
        # locations; it is built directly (no key, no random draw).
        log_sigma_init = math.log(sigma_init)
        core = geonnax.GaussianRegressionHead(
            W_loc=jnp.zeros((in_features, out_features)),
            b_loc=jnp.zeros(out_features),
            W_log_scale=jnp.zeros((in_features, out_features)),
            b_log_scale=jnp.full((out_features,), log_sigma_init),
            in_features=in_features,
            out_features=out_features,
        )
        # Strict, as in the core's init: on the bound the clip would start
        # half-saturated.
        if not core.log_sigma_min < log_sigma_init < core.log_sigma_max:
            raise ValueError(
                "log(sigma_init) must lie strictly inside "
                f"({core.log_sigma_min}, {core.log_sigma_max}); "
                f"got sigma_init={sigma_init}."
            )
        return cls(
            core=core,
            prior_std=prior_std,
            log_scale_prior_std=log_scale_prior_std,
            pyrox_name=pyrox_name,
        )

    @property
    def in_features(self) -> int:
        return self.core.in_features

    @property
    def out_features(self) -> int:
        return self.core.out_features

    def _prior(self, name: str) -> dist.Distribution:
        loc = getattr(self.core, name)
        scale = self.log_scale_prior_std if name == "W_log_scale" else self.prior_std
        return dist.Normal(loc, scale).to_event(loc.ndim)

    @pyrox_method
    def __call__(
        self, h: Float[Array, "*batch D_in"]
    ) -> tuple[Float[Array, "*batch K"], Float[Array, "*batch K"]]:
        r"""Sample the weights once and map features to $(\mu, \log\sigma)$.

        Args:
            h: Features of shape ``(*batch, D_in)``; ``(D_in,)`` for one.

        Returns:
            ``(mu, log_sigma)``, each of shape ``(*batch, K)``, with
            ``log_sigma`` clamped to $[-7, 3]$.
        """
        c = self.core
        names = ("W_loc", "b_loc", "W_log_scale", "b_log_scale")
        sampled = tuple(self.pyrox_sample(name, self._prior(name)) for name in names)
        core = eqx.tree_at(
            lambda m: (m.W_loc, m.b_loc, m.W_log_scale, m.b_log_scale), c, sampled
        )
        return vmap_over_flat_batch(core, h)

init(in_features: int, out_features: int = 1, *, prior_std: float = 1.0, log_scale_prior_std: float = 0.1, sigma_init: float = 1.0, pyrox_name: str | None = None) -> Self classmethod

Build the head with priors centred on a homoscedastic core.

Parameters:

Name	Type	Description	Default
in_features	int	Feature dimension \(D_\mathrm{in}\); must be > 0.	required
out_features	int	Number of targets \(K\); must be > 0.	1
prior_std	float	Prior standard deviation \(s\) of the mean weights and both biases; must be > 0.	1.0
log_scale_prior_std	float	Prior standard deviation \(s_\sigma\) of the log-scale weights \(W_s\); must be > 0. See the class docstring for the default.	0.1
sigma_init	float	Prior location of \(\sigma\), i.e. \(b_s \sim \mathcal{N}(\log\sigma_\mathrm{init}, s^2)\); \(\log\sigma_\mathrm{init}\) must lie strictly inside the core's clip range \((-7, 3)\).	1.0
pyrox_name	str | None	Explicit scope name for the sites.	None

Returns:

Type	Description
Self	The head.

Raises:

Type	Description
ValueError	If a dimension, prior_std or log_scale_prior_std is not positive, or sigma_init is not positive or puts \(\log\sigma_\mathrm{init}\) on or outside the clip range.

Source code in packages/pyrox-nn/src/pyrox_nn/_heteroscedastic.py

@classmethod
def init(
    cls,
    in_features: int,
    out_features: int = 1,
    *,
    prior_std: float = 1.0,
    log_scale_prior_std: float = 0.1,
    sigma_init: float = 1.0,
    pyrox_name: str | None = None,
) -> Self:
    r"""Build the head with priors centred on a homoscedastic core.

    Args:
        in_features: Feature dimension $D_\mathrm{in}$; must be > 0.
        out_features: Number of targets $K$; must be > 0.
        prior_std: Prior standard deviation $s$ of the mean weights and
            both biases; must be > 0.
        log_scale_prior_std: Prior standard deviation $s_\sigma$ of the
            log-scale weights $W_s$; must be > 0. See the class
            docstring for the default.
        sigma_init: Prior location of $\sigma$, i.e.
            $b_s \sim \mathcal{N}(\log\sigma_\mathrm{init}, s^2)$;
            $\log\sigma_\mathrm{init}$ must lie strictly inside the
            core's clip range $(-7, 3)$.
        pyrox_name: Explicit scope name for the sites.

    Returns:
        The head.

    Raises:
        ValueError: If a dimension, ``prior_std`` or
            ``log_scale_prior_std`` is not positive, or ``sigma_init`` is
            not positive or puts $\log\sigma_\mathrm{init}$ on or outside
            the clip range.
    """
    if in_features <= 0 or out_features <= 0:
        raise ValueError(
            "in_features and out_features must be > 0; "
            f"got {in_features=}, {out_features=}."
        )
    for name, value in (
        ("prior_std", prior_std),
        ("log_scale_prior_std", log_scale_prior_std),
        ("sigma_init", sigma_init),
    ):
        if not value > 0:
            raise ValueError(f"{name} must be > 0; got {value}.")
    # The weights are sampled, so the core only carries the prior
    # locations; it is built directly (no key, no random draw).
    log_sigma_init = math.log(sigma_init)
    core = geonnax.GaussianRegressionHead(
        W_loc=jnp.zeros((in_features, out_features)),
        b_loc=jnp.zeros(out_features),
        W_log_scale=jnp.zeros((in_features, out_features)),
        b_log_scale=jnp.full((out_features,), log_sigma_init),
        in_features=in_features,
        out_features=out_features,
    )
    # Strict, as in the core's init: on the bound the clip would start
    # half-saturated.
    if not core.log_sigma_min < log_sigma_init < core.log_sigma_max:
        raise ValueError(
            "log(sigma_init) must lie strictly inside "
            f"({core.log_sigma_min}, {core.log_sigma_max}); "
            f"got sigma_init={sigma_init}."
        )
    return cls(
        core=core,
        prior_std=prior_std,
        log_scale_prior_std=log_scale_prior_std,
        pyrox_name=pyrox_name,
    )

MCSigmoidDenseFA

Bases: _HeteroscedasticBase

Heteroscedastic multi-label output layer (FA noise + sigmoid).

Identical low-rank-plus-diagonal logit-noise model as MCSoftmaxDenseFA, but the per-class outputs are independent Bernoullis — final probabilities are the MC average of element-wise sigmoids, not a softmax. Use this for multi-label classification or independent binary heads.

\[ \hat{p}(y_k = 1 \mid x) \approx \frac{1}{S}\sum_{s=1}^{S} \sigma\!\bigl(\eta(x) + \epsilon_s\bigr)_k. \]

See MCSoftmaxDenseFA for the noise model, plate semantics, init API, and references.

Source code in packages/pyrox-nn/src/pyrox_nn/_heteroscedastic.py

class MCSigmoidDenseFA(_HeteroscedasticBase):
    r"""Heteroscedastic multi-label output layer (FA noise + sigmoid).

    Identical low-rank-plus-diagonal logit-noise model as
    `MCSoftmaxDenseFA`, but the per-class outputs are
    independent Bernoullis — final probabilities are the MC average of
    *element-wise* sigmoids, not a softmax. Use this for multi-label
    classification or independent binary heads.

    $$
    \hat{p}(y_k = 1 \mid x) \approx
    \frac{1}{S}\sum_{s=1}^{S}
    \sigma\!\bigl(\eta(x) + \epsilon_s\bigr)_k.
    $$

    See `MCSoftmaxDenseFA` for the noise model, plate semantics,
    init API, and references.
    """

    @classmethod
    def _core_cls(cls) -> type[geonnax.HeteroscedasticHead]:
        return geonnax.MCSigmoidDenseFA

    @pyrox_method
    def __call__(self, x: Float[Array, "N D_in"]) -> Float[Array, "N C"]:
        new_core = self._swap_core()
        logits = _batched_logits(new_core, x)
        return jnp.mean(jax.nn.sigmoid(logits), axis=0)

MCSoftmaxDenseFA

Bases: _HeteroscedasticBase

Heteroscedastic multi-class output layer (FA noise + softmax).

Implements Collier et al. (2021): the logit covariance is input-dependent low-rank-plus-diagonal,

\[ \eta(x) = W_\mu x + b_\mu + \epsilon, \qquad \Sigma(x) = V(x) V(x)^\top + \operatorname{diag}\!\bigl(\sigma^2(x)\bigr), \;\; \epsilon \sim \mathcal{N}(0, \Sigma(x)), \]

where \(V(x) = \mathrm{reshape}(W_V x + b_V, [C, r])\) and \(\sigma(x) = \exp(W_\sigma x + b_\sigma)\). Output is the Monte Carlo average of softmaxed perturbed logits

\[ \hat{p}(y = k \mid x) \approx \frac{1}{S}\sum_{s=1}^{S} \mathrm{softmax}_k\!\bigl(\eta(x) + \epsilon_s\bigr). \]

All linear factors are deterministic pyrox_param sites — the layer is heteroscedastic but not Bayesian over its weights. Use it as a drop-in head for classification when label noise is known to be input-dependent (label disagreement, fine-grained categories).

Plate semantics

Same as other pyrox_nn Bayesian dense layers — call outside numpyro.plate("data", ..., subsample_size=...) so the parameter sites are unscaled. The MC noise is drawn from numpyro.prng_key().

Attributes:

Name	Type	Description
in_features	int	Input dimension \(D_\mathrm{in}\).
num_classes	int	Number of classes \(C\).
rank	int	Rank \(r\) of the low-rank factor \(V(x)\).
num_mc_samples	int	Number of MC softmax samples \(S\) per forward call.
diag_init_bias	float	Initial value for the diagonal-scale bias b_diag (a small negative number keeps initial noise small).
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Examples:

>>> import jax.random as jr
>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> layer = MCSoftmaxDenseFA.init(
...     jr.PRNGKey(0), in_features=4, num_classes=3, rank=2,
... )
>>> x = jnp.ones((5, 4))
>>> with handlers.seed(rng_seed=0):
...     probs = layer(x)
>>> probs.shape
(5, 3)
>>> bool(jnp.allclose(probs.sum(axis=-1), 1.0))
True

References

Collier, M., Mustafa, B., Kokiopoulou, E., Jenatton, R., & Berent, J. (2021). Correlated Input-Dependent Label Noise in Large-Scale Image Classification. CVPR.

Source code in packages/pyrox-nn/src/pyrox_nn/_heteroscedastic.py

class MCSoftmaxDenseFA(_HeteroscedasticBase):
    r"""Heteroscedastic multi-class output layer (FA noise + softmax).

    Implements Collier et al. (2021): the logit covariance is
    input-dependent low-rank-plus-diagonal,

    $$
    \eta(x) = W_\mu x + b_\mu + \epsilon, \qquad
    \Sigma(x) = V(x) V(x)^\top + \operatorname{diag}\!\bigl(\sigma^2(x)\bigr),
    \;\; \epsilon \sim \mathcal{N}(0, \Sigma(x)),
    $$

    where $V(x) = \mathrm{reshape}(W_V x + b_V, [C, r])$ and
    $\sigma(x) = \exp(W_\sigma x + b_\sigma)$. Output is the
    Monte Carlo average of softmaxed perturbed logits

    $$
    \hat{p}(y = k \mid x) \approx
    \frac{1}{S}\sum_{s=1}^{S}
    \mathrm{softmax}_k\!\bigl(\eta(x) + \epsilon_s\bigr).
    $$

    All linear factors are deterministic ``pyrox_param`` sites — the
    layer is heteroscedastic but not Bayesian over its weights. Use it
    as a drop-in head for classification when label noise is known to
    be input-dependent (label disagreement, fine-grained categories).

    Plate semantics:
        Same as other ``pyrox_nn`` Bayesian dense layers — call
        outside ``numpyro.plate("data", ..., subsample_size=...)`` so
        the parameter sites are unscaled. The MC noise is drawn from
        ``numpyro.prng_key()``.

    Attributes:
        in_features: Input dimension $D_\mathrm{in}$.
        num_classes: Number of classes $C$.
        rank: Rank $r$ of the low-rank factor $V(x)$.
        num_mc_samples: Number of MC softmax samples $S$ per
            forward call.
        diag_init_bias: Initial value for the diagonal-scale bias
            ``b_diag`` (a small negative number keeps initial noise
            small).
        pyrox_name: Explicit scope name for NumPyro site registration.

    Examples:
        >>> import jax.random as jr
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> layer = MCSoftmaxDenseFA.init(
        ...     jr.PRNGKey(0), in_features=4, num_classes=3, rank=2,
        ... )
        >>> x = jnp.ones((5, 4))
        >>> with handlers.seed(rng_seed=0):
        ...     probs = layer(x)
        >>> probs.shape
        (5, 3)
        >>> bool(jnp.allclose(probs.sum(axis=-1), 1.0))
        True

    References:
        Collier, M., Mustafa, B., Kokiopoulou, E., Jenatton, R., &
        Berent, J. (2021). *Correlated Input-Dependent Label Noise in
        Large-Scale Image Classification.* CVPR.
    """

    @classmethod
    def _core_cls(cls) -> type[geonnax.HeteroscedasticHead]:
        return geonnax.MCSoftmaxDenseFA

    @pyrox_method
    def __call__(self, x: Float[Array, "N D_in"]) -> Float[Array, "N C"]:
        new_core = self._swap_core()
        logits = _batched_logits(new_core, x)
        return jnp.mean(jax.nn.softmax(logits, axis=-1), axis=0)

Bayesian Neural Field stack

Standardization

Bases: PyroxModule

Apply a fixed-coefficient affine standardization.

\[ \tilde x \;=\; \frac{x - \mu}{\sigma}. \]

Both mu and std are static (fit-time) constants, not learned. Use pyrox_nn.preprocessing.fit_standardization to construct from a pandas DataFrame.

Attributes:

Name	Type	Description
mu	Float[Array, ' D']	Per-feature mean, shape (D,).
std	Float[Array, ' D']	Per-feature standard deviation, shape (D,). Must be strictly positive — guard upstream.
pyrox_name	str | None	Optional override for the per-instance scope name.

Source code in packages/pyrox-nn/src/pyrox_nn/_bnf.py

class Standardization(PyroxModule):
    r"""Apply a fixed-coefficient affine standardization.

    $$
    \tilde x \;=\; \frac{x - \mu}{\sigma}.
    $$

    Both ``mu`` and ``std`` are static (fit-time) constants, not
    learned. Use `pyrox_nn.preprocessing.fit_standardization` to
    construct from a pandas DataFrame.

    Attributes:
        mu: Per-feature mean, shape ``(D,)``.
        std: Per-feature standard deviation, shape ``(D,)``. Must be
            strictly positive — guard upstream.
        pyrox_name: Optional override for the per-instance scope name.
    """

    mu: Float[Array, " D"]
    std: Float[Array, " D"]
    pyrox_name: str | None = None

    @pyrox_method
    def __call__(self, x: Float[Array, "N D"]) -> Float[Array, "N D"]:
        return (x - self.mu) / self.std

FourierFeatures

Bases: PyroxModule

Per-input dyadic-frequency Fourier basis.

For each input column, evaluates 2 * degree Fourier features at frequencies \(2\pi \cdot 2^d\) for \(d \in \{0, \dots, \text{degree} - 1\}\). Concatenated across all columns.

Wraps pyrox_nn._features.fourier_features per input dimension.

Attributes:

Name	Type	Description
degrees	tuple[int, ...]	Number of dyadic frequencies per input column, as a Python tuple[int, ...]. A column with degree = 0 contributes no features. Marked static so the loop over columns unrolls at trace time.
rescale	bool	If True, divide each (cos_d, sin_d) pair by d + 1 to bias the prior toward lower frequencies.
pyrox_name	str | None	Optional scope-name override.

Source code in packages/pyrox-nn/src/pyrox_nn/_bnf.py

class FourierFeatures(PyroxModule):
    r"""Per-input dyadic-frequency Fourier basis.

    For each input column, evaluates ``2 * degree`` Fourier features at
    frequencies $2\pi \cdot 2^d$ for $d \in \{0, \dots,
    \text{degree} - 1\}$. Concatenated across all columns.

    Wraps `pyrox_nn._features.fourier_features` per input
    dimension.

    Attributes:
        degrees: Number of dyadic frequencies per input column, as a
            Python ``tuple[int, ...]``. A column with ``degree = 0``
            contributes no features. Marked ``static`` so the loop
            over columns unrolls at trace time.
        rescale: If ``True``, divide each ``(cos_d, sin_d)`` pair by
            ``d + 1`` to bias the prior toward lower frequencies.
        pyrox_name: Optional scope-name override.
    """

    degrees: tuple[int, ...] = eqx.field(static=True)
    rescale: bool = eqx.field(static=True, default=False)
    pyrox_name: str | None = None

    @pyrox_method
    def __call__(self, x: Float[Array, "N D"]) -> Float[Array, "N F"]:
        feats = []
        for col_idx, d in enumerate(self.degrees):
            if d <= 0:
                continue
            feats.append(fourier_features(x[:, col_idx], d, rescale=self.rescale))
        if not feats:
            return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
        return jnp.concatenate(feats, axis=-1)

SeasonalFeatures

Bases: PyroxModule

Period-and-harmonic cos/sin basis on a scalar time axis.

For each period \(\tau_p\) with \(H_p\) harmonics, emits 2 * H_p cos/sin columns. Total output width is \(2 \sum_p H_p\).

Wraps pyrox_nn._features.seasonal_features. Periods and harmonics are kept as Python tuples (static) so the inner shape structure is known at trace time.

Attributes:

Name	Type	Description
periods	tuple[float, ...]	Period values, tuple[float, ...].
harmonics	tuple[int, ...]	Harmonics per period, tuple[int, ...].
rescale	bool	If True, divide each (cos, sin) pair by its within-period harmonic index.
pyrox_name	str | None	Optional scope-name override.

Source code in packages/pyrox-nn/src/pyrox_nn/_bnf.py

class SeasonalFeatures(PyroxModule):
    r"""Period-and-harmonic cos/sin basis on a scalar time axis.

    For each period $\tau_p$ with $H_p$ harmonics, emits
    ``2 * H_p`` cos/sin columns. Total output width is $2 \sum_p
    H_p$.

    Wraps `pyrox_nn._features.seasonal_features`. Periods and
    harmonics are kept as Python tuples (static) so the inner shape
    structure is known at trace time.

    Attributes:
        periods: Period values, ``tuple[float, ...]``.
        harmonics: Harmonics per period, ``tuple[int, ...]``.
        rescale: If ``True``, divide each ``(cos, sin)`` pair by its
            within-period harmonic index.
        pyrox_name: Optional scope-name override.
    """

    periods: tuple[float, ...] = eqx.field(static=True)
    harmonics: tuple[int, ...] = eqx.field(static=True)
    rescale: bool = eqx.field(static=True, default=False)
    pyrox_name: str | None = None

    @pyrox_method
    def __call__(self, t: Float[Array, " N"]) -> Float[Array, "N F"]:
        return seasonal_features(t, self.periods, self.harmonics, rescale=self.rescale)

InteractionFeatures

Bases: PyroxModule

Element-wise products on selected pairs of input columns.

Wraps pyrox_nn._features.interaction_features.

Attributes:

Name	Type	Description
pairs	tuple[tuple[int, int], ...]	Index pairs, tuple[tuple[int, int], ...]. Empty tuple produces an (N, 0) output. Static so the count K is known at trace time.
pyrox_name	str | None	Optional scope-name override.

Source code in packages/pyrox-nn/src/pyrox_nn/_bnf.py

class InteractionFeatures(PyroxModule):
    r"""Element-wise products on selected pairs of input columns.

    Wraps `pyrox_nn._features.interaction_features`.

    Attributes:
        pairs: Index pairs, ``tuple[tuple[int, int], ...]``. Empty
            tuple produces an ``(N, 0)`` output. Static so the count
            ``K`` is known at trace time.
        pyrox_name: Optional scope-name override.
    """

    pairs: tuple[tuple[int, int], ...] = eqx.field(static=True)
    pyrox_name: str | None = None

    @pyrox_method
    def __call__(self, x: Float[Array, "N D"]) -> Float[Array, "N K"]:
        if not self.pairs:
            return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
        return interaction_features(x, jnp.asarray(self.pairs, dtype=jnp.int32))

BayesianNeuralField

Bases: PyroxModule

The full Bayesian Neural Field architecture, wrapping geonnax.NeuralFieldMLP.

The deterministic network is geonnax.NeuralFieldMLP; this layer puts an independent \(\mathrm{Logistic}(0, 1)\) prior on every one of its parameters, registers them with PyroxModule.pyrox_sample, splices the draws into the core with eqx.tree_at and runs it over the batch. For one input \(x \in \mathbb{R}^D\) with fixed scales \(s\):

1. The input is rescaled, \(\tilde x = x \oslash (s \odot e^{a})\), with a learned log-scale adjustment \(a\).
2. Feature blocks \(F_k\) are built: \(\tilde x\), dyadic Fourier features of each column with a positive degree, seasonal features of the raw time column, interaction products of \(\tilde x\), and, when an encoder is given, the encoder's features of x[encoder_cols]. Each block is gated by \(\mathrm{softplus}(g_k)\); the concatenation is \(h_0\). The encoder's columns enter only through the encoder: they are left out of \(\tilde x\) (and so of the scales, Fourier and interaction features). A spherical-harmonic encoder on lon/lat therefore gives a field that is continuous across the dateline and regular at the poles by construction.
3. Hidden layers \(h_{\ell+1} = \sigma_\alpha\bigl(\mathrm{softplus}(c_\ell)\, (W_\ell^\top h_\ell / \sqrt{\lvert h_\ell \rvert} + b_\ell)\bigr)\), with \(\sigma_\alpha = \alpha\,\mathrm{elu} + (1 - \alpha)\tanh\) and \(\alpha = \mathrm{sigmoid}(\beta)\).
4. A readout \(f(x) = \mathrm{softplus}(c_{\text{out}})\, (W_{\text{out}}^\top h_L / \sqrt{\lvert h_L \rvert} + b_{\text{out}})\).

The \(1/\sqrt{\text{fan-in}}\) factor is the NTK scaling of BayesNF: it makes the layer-wise prior predictive a fan-in-independent Gaussian process in the infinite-width limit (Lee et al., 2018).

Sites (scope pyrox_name, else BayesianNeuralField), each with a \(\mathrm{Logistic}(0, 1)\) prior at full shape, in registration order: log_scale_adjustment \((D,)\) (\(D\) minus the encoder's columns with an encoder); feature_gain_{b} (scalar, one per feature block, the encoder block last); logit_activation_weight (scalar); layer_{l}_W \((I_\ell, \text{width})\), layer_{l}_b \((\text{width},)\) and layer_{l}_gain (scalar) for each hidden layer; output_W \((\text{width}, 1)\); output_b \((1,)\); output_gain (scalar). The names and shapes are those of the original inline implementation, so saved posteriors still load.

The fixed scales follow the input's dtype (the inline implementation cast them to float32), so float64 inputs get float64 scales.

Attributes:

Name	Type	Description
input_scales	tuple[float, ...]	Per-input fixed scale (typically training-data inter-quartile range). Static tuple[float, ...].
fourier_degrees	tuple[int, ...]	Per-input number of dyadic Fourier frequencies. Static tuple[int, ...]; use 0 to skip a column.
interactions	tuple[tuple[int, int], ...]	Pair-index list for interaction features. Static tuple[tuple[int, int], ...]; empty for none.
seasonality_periods	tuple[float, ...]	Periods for seasonal features. Static tuple[float, ...]. The time variable is taken from input column time_col.
num_seasonal_harmonics	tuple[int, ...]	Harmonics per period. Static tuple[int, ...].
width	int	Hidden layer width.
depth	int	Number of hidden MLP layers.
time_col	int	Index of the time column inside x used for seasonal features (default 0).
pyrox_name	str | None	Optional scope-name override.
encoder	Callable[[Float[Array, ' C']], Float[Array, ' E']] | None	Optional single-example encoder (C,) -> (E,) applied to x[encoder_cols], whose output enters as one extra gated feature block (for example geonnax.SphericalHarmonicEncoder(l_max=10, input_mode="lonlat") on a lon/lat pair in radians). It is deterministic: its arrays, if any, are not sample sites.
encoder_cols	tuple[int, ...] | None	Input columns fed to encoder, in order; required exactly when encoder is given. They feed only the encoder, so they may not have Fourier degrees, appear in interactions or be the seasonal time column, and at least one column must stay outside them.
core	NeuralFieldMLP	The wrapped geonnax.NeuralFieldMLP, built from the fields above. Its arrays are placeholders: every one is replaced by a sample site on each call.

Raises:

Type	Description
ValueError	If encoder and encoder_cols are not given together, an encoder column is out of range or also used by another feature block, no column is left outside them, the encoder does not return a vector, or the core rejects the configuration.

Examples:

>>> import jax.numpy as jnp
>>> from numpyro import handlers
>>> from geonnax import SphericalHarmonicEncoder
>>> bnf = BayesianNeuralField(
...     input_scales=(1.0, 1.0, 1.0), fourier_degrees=(0, 0, 0),
...     interactions=(), seasonality_periods=(365.25,),
...     num_seasonal_harmonics=(2,), width=16, depth=2,
...     encoder=SphericalHarmonicEncoder(l_max=3, input_mode="lonlat"),
...     encoder_cols=(1, 2), pyrox_name="bnf",
... )
>>> with handlers.seed(rng_seed=0):
...     f = bnf(jnp.zeros((5, 3)))
>>> f.shape
(5,)

References

Saad, F. A. et al. (2024). Scalable spatiotemporal prediction with Bayesian neural fields. Nature Communications 15, 7942. Lee, J. et al. (2018). Deep neural networks as Gaussian processes. ICLR.

Source code in packages/pyrox-nn/src/pyrox_nn/_bnf.py

class BayesianNeuralField(PyroxModule):
    r"""The full Bayesian Neural Field architecture, wrapping `geonnax.NeuralFieldMLP`.

    The deterministic network is `geonnax.NeuralFieldMLP`; this layer puts
    an independent $\mathrm{Logistic}(0, 1)$ prior on every one of its
    parameters, registers them with `PyroxModule.pyrox_sample`, splices the
    draws into the core with `eqx.tree_at` and runs it over the batch. For
    one input $x \in \mathbb{R}^D$ with fixed scales $s$:

    1. The input is rescaled, $\tilde x = x \oslash (s \odot e^{a})$, with a
       learned log-scale adjustment $a$.
    2. Feature blocks $F_k$ are built: $\tilde x$, dyadic Fourier features of
       each column with a positive degree, seasonal features of the raw time
       column, interaction products of $\tilde x$, and, when an ``encoder``
       is given, the encoder's features of ``x[encoder_cols]``. Each block is
       gated by $\mathrm{softplus}(g_k)$; the concatenation is $h_0$. The
       encoder's columns enter **only** through the encoder: they are left
       out of $\tilde x$ (and so of the scales, Fourier and interaction
       features). A spherical-harmonic encoder on lon/lat therefore gives a
       field that is continuous across the dateline and regular at the
       poles by construction.
    3. Hidden layers
       $h_{\ell+1} = \sigma_\alpha\bigl(\mathrm{softplus}(c_\ell)\,
       (W_\ell^\top h_\ell / \sqrt{\lvert h_\ell \rvert} + b_\ell)\bigr)$,
       with $\sigma_\alpha = \alpha\,\mathrm{elu} + (1 - \alpha)\tanh$ and
       $\alpha = \mathrm{sigmoid}(\beta)$.
    4. A readout
       $f(x) = \mathrm{softplus}(c_{\text{out}})\,
       (W_{\text{out}}^\top h_L / \sqrt{\lvert h_L \rvert} + b_{\text{out}})$.

    The $1/\sqrt{\text{fan-in}}$ factor is the NTK scaling of BayesNF: it
    makes the layer-wise prior predictive a fan-in-independent Gaussian
    process in the infinite-width limit (Lee et al., 2018).

    Sites (scope ``pyrox_name``, else ``BayesianNeuralField``), each with a
    $\mathrm{Logistic}(0, 1)$ prior at full shape, in registration order:
    ``log_scale_adjustment`` $(D,)$ ($D$ minus the encoder's columns with
    an encoder); ``feature_gain_{b}`` (scalar, one per
    feature block, the encoder block last); ``logit_activation_weight``
    (scalar); ``layer_{l}_W`` $(I_\ell, \text{width})$, ``layer_{l}_b``
    $(\text{width},)$ and ``layer_{l}_gain`` (scalar) for each hidden layer;
    ``output_W`` $(\text{width}, 1)$; ``output_b`` $(1,)$; ``output_gain``
    (scalar). The names and shapes are those of the original inline
    implementation, so saved posteriors still load.

    The fixed scales follow the input's dtype (the inline implementation
    cast them to float32), so float64 inputs get float64 scales.

    Attributes:
        input_scales: Per-input fixed scale (typically training-data
            inter-quartile range). Static ``tuple[float, ...]``.
        fourier_degrees: Per-input number of dyadic Fourier
            frequencies. Static ``tuple[int, ...]``; use ``0`` to skip
            a column.
        interactions: Pair-index list for interaction features. Static
            ``tuple[tuple[int, int], ...]``; empty for none.
        seasonality_periods: Periods for seasonal features. Static
            ``tuple[float, ...]``. The time variable is taken from
            input column ``time_col``.
        num_seasonal_harmonics: Harmonics per period. Static
            ``tuple[int, ...]``.
        width: Hidden layer width.
        depth: Number of hidden MLP layers.
        time_col: Index of the time column inside ``x`` used for
            seasonal features (default 0).
        pyrox_name: Optional scope-name override.
        encoder: Optional single-example encoder ``(C,) -> (E,)`` applied
            to ``x[encoder_cols]``, whose output enters as one extra gated
            feature block (for example
            ``geonnax.SphericalHarmonicEncoder(l_max=10, input_mode="lonlat")``
            on a lon/lat pair in radians). It is deterministic: its arrays,
            if any, are not sample sites.
        encoder_cols: Input columns fed to ``encoder``, in order; required
            exactly when ``encoder`` is given. They feed only the encoder,
            so they may not have Fourier degrees, appear in
            ``interactions`` or be the seasonal time column, and at least
            one column must stay outside them.
        core: The wrapped `geonnax.NeuralFieldMLP`, built from the fields
            above. Its arrays are placeholders: every one is replaced by a
            sample site on each call.

    Raises:
        ValueError: If ``encoder`` and ``encoder_cols`` are not given
            together, an encoder column is out of range or also used by
            another feature block, no column is left outside them, the
            encoder does not return a vector, or the core rejects the
            configuration.

    Examples:
        >>> import jax.numpy as jnp
        >>> from numpyro import handlers
        >>> from geonnax import SphericalHarmonicEncoder
        >>> bnf = BayesianNeuralField(
        ...     input_scales=(1.0, 1.0, 1.0), fourier_degrees=(0, 0, 0),
        ...     interactions=(), seasonality_periods=(365.25,),
        ...     num_seasonal_harmonics=(2,), width=16, depth=2,
        ...     encoder=SphericalHarmonicEncoder(l_max=3, input_mode="lonlat"),
        ...     encoder_cols=(1, 2), pyrox_name="bnf",
        ... )
        >>> with handlers.seed(rng_seed=0):
        ...     f = bnf(jnp.zeros((5, 3)))
        >>> f.shape
        (5,)

    References:
        Saad, F. A. *et al.* (2024). Scalable spatiotemporal prediction with
        Bayesian neural fields. *Nature Communications* 15, 7942.
        Lee, J. *et al.* (2018). Deep neural networks as Gaussian processes.
        *ICLR*.
    """

    input_scales: tuple[float, ...] = eqx.field(static=True)
    fourier_degrees: tuple[int, ...] = eqx.field(static=True)
    interactions: tuple[tuple[int, int], ...] = eqx.field(static=True)
    seasonality_periods: tuple[float, ...] = eqx.field(static=True)
    num_seasonal_harmonics: tuple[int, ...] = eqx.field(static=True)
    width: int = eqx.field(static=True)
    depth: int = eqx.field(static=True)
    time_col: int = eqx.field(static=True)
    pyrox_name: str | None
    encoder: Callable[[Float[Array, " C"]], Float[Array, " E"]] | None
    encoder_cols: tuple[int, ...] | None = eqx.field(static=True)
    core: NeuralFieldMLP

    def __init__(
        self,
        input_scales: tuple[float, ...],
        fourier_degrees: tuple[int, ...],
        interactions: tuple[tuple[int, int], ...],
        seasonality_periods: tuple[float, ...],
        num_seasonal_harmonics: tuple[int, ...],
        width: int,
        depth: int,
        time_col: int = 0,
        pyrox_name: str | None = None,
        encoder: Callable[[Float[Array, " C"]], Float[Array, " E"]] | None = None,
        encoder_cols: tuple[int, ...] | None = None,
    ) -> None:
        d_in = len(input_scales)
        extra_features = 0
        if (encoder is None) != (encoder_cols is None):
            raise ValueError(
                "encoder and encoder_cols must be given together; got "
                f"encoder={encoder!r}, encoder_cols={encoder_cols!r}."
            )
        if encoder is not None and encoder_cols is not None:
            encoder_cols = tuple(int(c) for c in encoder_cols)
            if not encoder_cols or any(not 0 <= c < d_in for c in encoder_cols):
                raise ValueError(
                    f"encoder_cols must be a non-empty tuple of indices in "
                    f"[0, {d_in}); got {encoder_cols!r}."
                )
            out = jax.eval_shape(
                encoder, jax.ShapeDtypeStruct((len(encoder_cols),), jnp.float32)
            )
            if len(out.shape) != 1:
                raise ValueError(
                    "encoder must map a (C,) vector to an (E,) vector; got output "
                    f"shape {out.shape}."
                )
            extra_features = out.shape[0]
        core_cols = _core_columns(
            d_in,
            encoder_cols or (),
            fourier_degrees,
            interactions,
            seasonality_periods,
            num_seasonal_harmonics,
            time_col,
        )
        remap = {c: i for i, c in enumerate(core_cols)}
        self.input_scales = input_scales
        self.fourier_degrees = fourier_degrees
        self.interactions = interactions
        self.seasonality_periods = seasonality_periods
        self.num_seasonal_harmonics = num_seasonal_harmonics
        self.width = width
        self.depth = depth
        self.time_col = time_col
        self.pyrox_name = pyrox_name
        self.encoder = encoder
        self.encoder_cols = encoder_cols
        # The core sees only the columns the encoder does not consume. Its
        # arrays are placeholders: `__call__` replaces every one of them
        # (all nine parameter fields) with a sample site before running the
        # core, so their values never reach an output and a fixed key is
        # harmless. Real arrays (rather than `eqx.filter_eval_shape`
        # structs) keep the module an ordinary pytree of arrays under
        # `jit`, `vmap` and serialisation.
        self.core = NeuralFieldMLP.init(
            tuple(input_scales[c] for c in core_cols),
            width=width,
            depth=depth,
            fourier_degrees=tuple(fourier_degrees[c] for c in core_cols)
            if fourier_degrees
            else (),
            interactions=tuple((remap[i], remap[j]) for i, j in interactions),
            seasonality_periods=seasonality_periods,
            num_seasonal_harmonics=num_seasonal_harmonics,
            time_col=remap.get(time_col, 0),
            extra_features=extra_features,
            key=jr.PRNGKey(0),
        )

    @staticmethod
    def _logistic_prior(shape: tuple[int, ...]) -> dist.Distribution:
        """Independent Logistic(0, 1) prior over an array of given shape."""
        if not shape:
            return dist.Logistic(0.0, 1.0)
        return dist.Logistic(jnp.zeros(shape), 1.0).to_event(len(shape))

    @pyrox_method
    def __call__(self, x: Float[Array, "*batch D"]) -> Float[Array, " *batch"]:
        core = self.core
        prior = self._logistic_prior
        # Sites are registered in the order of the original inline
        # implementation, so seeded draws are unchanged too.
        log_scale_adjustment = self.pyrox_sample(
            "log_scale_adjustment", prior(core.log_scale_adjustment.shape)
        )
        feature_gain = jnp.stack(
            [
                self.pyrox_sample(f"feature_gain_{b}", prior(()))
                for b in range(core.feature_gain.shape[0])
            ]
        )
        logit_activation_weight = self.pyrox_sample(
            "logit_activation_weight", prior(())
        )
        layers_W, layers_b, layers_gain = [], [], []
        for layer_idx in range(self.depth):
            layers_W.append(
                self.pyrox_sample(
                    f"layer_{layer_idx}_W", prior(core.layers_W[layer_idx].shape)
                )
            )
            layers_b.append(
                self.pyrox_sample(
                    f"layer_{layer_idx}_b", prior(core.layers_b[layer_idx].shape)
                )
            )
            layers_gain.append(self.pyrox_sample(f"layer_{layer_idx}_gain", prior(())))
        output_W = self.pyrox_sample("output_W", prior(core.output_W.shape))
        output_b = self.pyrox_sample("output_b", prior(core.output_b.shape))
        output_gain = self.pyrox_sample("output_gain", prior(()))

        # With depth 0 the layer fields are empty (an empty tuple has no
        # identity `eqx.tree_at` can find), so they are only swapped when
        # there are hidden layers.
        def where(c: NeuralFieldMLP) -> tuple[Any, ...]:
            nodes = (
                c.log_scale_adjustment,
                c.feature_gain,
                c.logit_activation_weight,
                c.output_W,
                c.output_b,
                c.output_gain,
            )
            if self.depth > 0:
                nodes = (*nodes, c.layers_W, c.layers_b, c.layers_gain)
            return nodes

        values: tuple[Any, ...] = (
            log_scale_adjustment,
            feature_gain,
            logit_activation_weight,
            output_W,
            output_b,
            output_gain,
        )
        if self.depth > 0:
            values = (
                *values,
                tuple(layers_W),
                tuple(layers_b),
                jnp.stack(layers_gain),
            )
        core = eqx.tree_at(where, core, values)

        encoder, encoder_cols = self.encoder, self.encoder_cols
        if encoder is None or encoder_cols is None:
            fn = core
        else:
            enc_idx = jnp.asarray(encoder_cols, dtype=jnp.int32)
            core_idx = jnp.asarray(
                [c for c in range(len(self.input_scales)) if c not in encoder_cols],
                dtype=jnp.int32,
            )

            def fn(xi: Float[Array, " D"]) -> Float[Array, " 1"]:
                return core(xi[core_idx], encoder(xi[enc_idx]))

        return vmap_over_flat_batch(fn, x)[..., 0]

Pure-JAX feature helpers

fourier_features(x: Float[Array, ' N'], max_degree: int, *, rescale: bool = False) -> Float[Array, 'N two_max_degree']

Cos/sin Fourier basis at dyadic frequencies.

For each input element and each degree \(d \in \{0, \dots, D-1\}\), evaluates

\[ \phi_{d, \cos}(x) = \cos(2\pi \cdot 2^d \cdot x), \qquad \phi_{d, \sin}(x) = \sin(2\pi \cdot 2^d \cdot x). \]

Returns the columns concatenated as [cos_0, ..., cos_{D-1}, sin_0, ..., sin_{D-1}], matching Google's bayesnf layout.

Parameters:

Name	Type	Description	Default
x	Float[Array, ' N']	Length-N input vector.	required
max_degree	int	Number of dyadic frequencies D. Output has 2 * max_degree columns.	required
rescale	bool	If True, divide each (cos_d, sin_d) pair by d + 1 to bias the prior toward lower-frequency basis functions.	False

Returns:

Type	Description
Float[Array, 'N two_max_degree']	Array of shape (N, 2 * max_degree).

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import fourier_features
>>> fourier_features(jnp.linspace(0.0, 1.0, 5), max_degree=3).shape
(5, 6)

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def fourier_features(
    x: Float[Array, " N"],
    max_degree: int,
    *,
    rescale: bool = False,
) -> Float[Array, "N two_max_degree"]:
    r"""Cos/sin Fourier basis at dyadic frequencies.

    For each input element and each degree $d \in \{0, \dots,
    D-1\}$, evaluates

    $$
    \phi_{d, \cos}(x) = \cos(2\pi \cdot 2^d \cdot x), \qquad
    \phi_{d, \sin}(x) = \sin(2\pi \cdot 2^d \cdot x).
    $$


    Returns the columns concatenated as ``[cos_0, ..., cos_{D-1},
    sin_0, ..., sin_{D-1}]``, matching Google's bayesnf layout.

    Args:
        x: Length-``N`` input vector.
        max_degree: Number of dyadic frequencies ``D``. Output has
            ``2 * max_degree`` columns.
        rescale: If ``True``, divide each ``(cos_d, sin_d)`` pair by
            ``d + 1`` to bias the prior toward lower-frequency basis
            functions.

    Returns:
        Array of shape ``(N, 2 * max_degree)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import fourier_features
        >>> fourier_features(jnp.linspace(0.0, 1.0, 5), max_degree=3).shape
        (5, 6)
    """
    x = _promote_to_floating(x)
    # Degrees in x's dtype so a float32 input stays float32 under x64.
    degrees = jnp.arange(max_degree, dtype=x.dtype)
    # Broadcast x to (N, D) frequencies without an explicit reshape.
    z = einx.id("n -> n d", x, d=max_degree) * (2.0 * jnp.pi * 2.0**degrees)
    feats = jnp.concatenate([jnp.cos(z), jnp.sin(z)], axis=-1)
    if rescale:
        denom = jnp.concatenate([degrees + 1, degrees + 1])
        feats = feats / denom
    return feats

seasonal_features(x: Float[Array, ' N'], periods: Sequence[float], harmonics: Sequence[int], *, rescale: bool = False) -> Float[Array, 'N two_F']

Cos/sin features at multiples of \(2\pi / \tau_p\).

For each period \(\tau_p\) with \(H_p\) harmonics, evaluates

\[ \phi_{p, h, \cos}(x) = \cos(2\pi h x / \tau_p), \qquad \phi_{p, h, \sin}(x) = \sin(2\pi h x / \tau_p), \]

for \(h = 1, \dots, H_p\). Returns the cos columns concatenated with the sin columns, length \(F = \sum_p H_p\) each.

periods and harmonics are Python sequences (tuples, lists, or 0-d JAX arrays wrapped at the call site). Keeping them as Python values lets the function run cleanly under jax.jit and lax.scan without triggering a concretization error.

Parameters:

Name	Type	Description	Default
x	Float[Array, ' N']	Time/index input, shape (N,).	required
periods	Sequence[float]	Period values.	required
harmonics	Sequence[int]	Harmonics per period.	required
rescale	bool	If True, divide each (cos, sin) pair by its within-period harmonic index, biasing the prior toward longer-wavelength modes within each period.	False

Returns:

Type	Description
Float[Array, 'N two_F']	Array of shape (N, 2 * F).

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import seasonal_features
>>> x = jnp.linspace(0.0, 10.0, 4)
>>> # F = 1 + 2 = 3 frequencies -> 2 * F = 6 columns
>>> seasonal_features(x, [7.0, 365.0], [1, 2]).shape
(4, 6)

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def seasonal_features(
    x: Float[Array, " N"],
    periods: Sequence[float],
    harmonics: Sequence[int],
    *,
    rescale: bool = False,
) -> Float[Array, "N two_F"]:
    r"""Cos/sin features at multiples of $2\pi / \tau_p$.

    For each period $\tau_p$ with $H_p$ harmonics, evaluates

    $$
    \phi_{p, h, \cos}(x) = \cos(2\pi h x / \tau_p), \qquad
    \phi_{p, h, \sin}(x) = \sin(2\pi h x / \tau_p),
    $$


    for $h = 1, \dots, H_p$. Returns the cos columns concatenated
    with the sin columns, length $F = \sum_p H_p$ each.

    ``periods`` and ``harmonics`` are **Python sequences** (tuples,
    lists, or 0-d JAX arrays wrapped at the call site). Keeping them as
    Python values lets the function run cleanly under ``jax.jit`` and
    ``lax.scan`` without triggering a concretization error.

    Args:
        x: Time/index input, shape ``(N,)``.
        periods: Period values.
        harmonics: Harmonics per period.
        rescale: If ``True``, divide each ``(cos, sin)`` pair by its
            within-period harmonic index, biasing the prior toward
            longer-wavelength modes within each period.

    Returns:
        Array of shape ``(N, 2 * F)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import seasonal_features
        >>> x = jnp.linspace(0.0, 10.0, 4)
        >>> # F = 1 + 2 = 3 frequencies -> 2 * F = 6 columns
        >>> seasonal_features(x, [7.0, 365.0], [1, 2]).shape
        (4, 6)
    """
    x = _promote_to_floating(x)
    _, freq_list = seasonal_frequencies(periods, harmonics)
    if not freq_list:
        return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
    # Frequencies in x's dtype: a float64 time axis (e.g. seconds since an
    # epoch) keeps a float64 phase instead of a float32-rounded 1/period.
    freqs = jnp.asarray(freq_list, dtype=x.dtype)
    z = einx.id("n -> n f", x, f=freqs.shape[0]) * (2.0 * jnp.pi * freqs)
    feats = jnp.concatenate([jnp.cos(z), jnp.sin(z)], axis=-1)
    if rescale:
        # Rescale by within-period harmonic index (1, 2, ..., H_p).
        h_within_list: list[float] = []
        for n_h in harmonics:
            h_within_list.extend(range(1, int(n_h) + 1))
        h_within = jnp.asarray(h_within_list, dtype=x.dtype)
        denom = jnp.concatenate([h_within, h_within])
        feats = feats / denom
    return feats

seasonal_frequencies(periods: Sequence[float], harmonics: Sequence[int]) -> tuple[list[int], list[float]]

Flatten (period, harmonic_count) pairs into Python lists.

For each period \(\tau_p\) with \(H_p\) harmonics, emits frequencies \(f_{p, h} = h / \tau_p\) for \(h = 1, \dots, H_p\). The total length is \(F = \sum_p H_p\).

Inputs are Python sequences, not JAX arrays, so this helper runs at trace time and never triggers a concretization error under jax.jit. Most callers won't use it directly; it's exposed for symmetry with seasonal_features.

Parameters:

Name	Type	Description	Default
periods	Sequence[float]	Period values.	required
harmonics	Sequence[int]	Number of harmonics per period.	required

Returns:

Type	Description
list[int]	(period_index, frequency): two Python lists of length
list[float]	\(F = \sum_p H_p\).

Examples:

>>> from geonnax.basis import seasonal_frequencies
>>> # periods (7, 365) with (1, 2) harmonics -> F = 1 + 2 = 3 freqs
>>> idx, freqs = seasonal_frequencies([7.0, 365.0], [1, 2])
>>> idx
[0, 1, 1]
>>> len(freqs)
3

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def seasonal_frequencies(
    periods: Sequence[float],
    harmonics: Sequence[int],
) -> tuple[list[int], list[float]]:
    r"""Flatten ``(period, harmonic_count)`` pairs into Python lists.

    For each period $\tau_p$ with $H_p$ harmonics, emits
    frequencies $f_{p, h} = h / \tau_p$ for $h = 1, \dots,
    H_p$. The total length is $F = \sum_p H_p$.

    Inputs are **Python sequences**, not JAX arrays, so this helper
    runs at trace time and never triggers a concretization error under
    ``jax.jit``. Most callers won't use it directly; it's exposed for
    symmetry with `seasonal_features`.

    Args:
        periods: Period values.
        harmonics: Number of harmonics per period.

    Returns:
        ``(period_index, frequency)``: two Python lists of length
        $F = \sum_p H_p$.

    Examples:
        >>> from geonnax.basis import seasonal_frequencies
        >>> # periods (7, 365) with (1, 2) harmonics -> F = 1 + 2 = 3 freqs
        >>> idx, freqs = seasonal_frequencies([7.0, 365.0], [1, 2])
        >>> idx
        [0, 1, 1]
        >>> len(freqs)
        3
    """
    period_index: list[int] = []
    freqs: list[float] = []
    for p_idx, (period, n_h) in enumerate(zip(periods, harmonics, strict=True)):
        for h in range(1, int(n_h) + 1):
            period_index.append(p_idx)
            freqs.append(float(h) / float(period))
    return period_index, freqs

interaction_features(x: Float[Array, 'N D'], pairs: Int[Array, 'K 2']) -> Float[Array, 'N K']

Element-wise products on selected pairs of input columns.

For each pair \((i, j)\) and each row \(n\), computes \(x_{n, i} \cdot x_{n, j}\).

Parameters:

Name	Type	Description	Default
x	Float[Array, 'N D']	Input matrix, shape (N, D).	required
pairs	Int[Array, 'K 2']	Index pairs, shape (K, 2). Empty pairs yield an (N, 0) output.	required

Returns:

Type	Description
Float[Array, 'N K']	Array of shape (N, K) of pairwise products.

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import interaction_features
>>> x = jnp.arange(6.0).reshape(2, 3)  # (N=2, D=3)
>>> pairs = jnp.array([[0, 1], [0, 2]])  # (K=2, 2)
>>> interaction_features(x, pairs).shape
(2, 2)

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def interaction_features(
    x: Float[Array, "N D"],
    pairs: Int[Array, "K 2"],
) -> Float[Array, "N K"]:
    r"""Element-wise products on selected pairs of input columns.

    For each pair $(i, j)$ and each row $n$, computes
    $x_{n, i} \cdot x_{n, j}$.

    Args:
        x: Input matrix, shape ``(N, D)``.
        pairs: Index pairs, shape ``(K, 2)``. Empty pairs yield an
            ``(N, 0)`` output.

    Returns:
        Array of shape ``(N, K)`` of pairwise products.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import interaction_features
        >>> x = jnp.arange(6.0).reshape(2, 3)  # (N=2, D=3)
        >>> pairs = jnp.array([[0, 1], [0, 2]])  # (K=2, 2)
        >>> interaction_features(x, pairs).shape
        (2, 2)
    """
    if pairs.shape[0] == 0:
        return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
    # x[:, pairs] has shape (N, K, 2); reduce the paired axis with prod.
    selected = x[:, pairs]
    return einx.prod("n k [two]", selected)

standardize(x: Float[Array, '*shape'], mu: Float[Array, '*shape'], std: Float[Array, '*shape']) -> Float[Array, '*shape']

Affine standardize: (x - mu) / std.

Broadcasts mu and std against x per the JAX broadcasting rules. std is not clamped; pass a positive value or guard upstream.

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import standardize
>>> x = jnp.array([1.0, 3.0, 5.0])
>>> bool(jnp.allclose(standardize(x, 3.0, 2.0), jnp.array([-1, 0, 1])))
True

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def standardize(
    x: Float[Array, "*shape"],
    mu: Float[Array, "*shape"],
    std: Float[Array, "*shape"],
) -> Float[Array, "*shape"]:
    """Affine standardize: ``(x - mu) / std``.

    Broadcasts ``mu`` and ``std`` against ``x`` per the JAX broadcasting
    rules. ``std`` is *not* clamped; pass a positive value or guard
    upstream.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import standardize
        >>> x = jnp.array([1.0, 3.0, 5.0])
        >>> bool(jnp.allclose(standardize(x, 3.0, 2.0), jnp.array([-1, 0, 1])))
        True
    """
    return (x - mu) / std

unstandardize(z: Float[Array, '*shape'], mu: Float[Array, '*shape'], std: Float[Array, '*shape']) -> Float[Array, '*shape']

Inverse of standardize: z * std + mu.

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import standardize, unstandardize
>>> x = jnp.array([1.0, 3.0, 5.0])
>>> # unstandardize undoes standardize for the same (mu, std).
>>> z = standardize(x, 3.0, 2.0)
>>> bool(jnp.allclose(unstandardize(z, 3.0, 2.0), x))
True

Source code in .venv/lib/python3.12/site-packages/geonnax/basis.py

def unstandardize(
    z: Float[Array, "*shape"],
    mu: Float[Array, "*shape"],
    std: Float[Array, "*shape"],
) -> Float[Array, "*shape"]:
    """Inverse of `standardize`: ``z * std + mu``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import standardize, unstandardize
        >>> x = jnp.array([1.0, 3.0, 5.0])
        >>> # unstandardize undoes standardize for the same (mu, std).
        >>> z = standardize(x, 3.0, 2.0)
        >>> bool(jnp.allclose(unstandardize(z, 3.0, 2.0), x))
        True
    """
    return z * std + mu

Neural operators and basis transforms (re-exported from geonnax)

Deterministic operator cores from geonnax, importable from pyrox_nn so that one import surface covers the stack (from pyrox_nn import GeoFNO).

GeoFNO

Bases: Module

Geo-FNO: a Fourier neural operator on per-sample point clouds.

Reads a field \(u\) at \(N\) points \(x_n\) with quadrature weights \(w_n\), maps the points onto the latent torus with the learned chart \(\xi = \varphi(x)\) (LearnedDeformation), runs FNO blocks on a uniform latent grid \(G\) (shape latent_shape, nodes \(j / s_a\) on \([0, 1)^d\)), and evaluates the result at \(M\) query points \(x'_m\) (the input points by default):

\[ \begin{aligned} h_0 &= \sigma\big(\mathcal{K}_0[P u;\ \xi, w] + b_0(G)\big) && \text{points} \to G, \\ h_\ell &= \sigma\big(\mathcal{K}_\ell h_{\ell-1} + W_\ell h_{\ell-1} + b_\ell(G)\big), \quad \ell = 1, \dots, L-2 && \text{on } G, \\ y &= Q_2\, \sigma\big(Q_1 (\mathcal{K}_{L-1} h_{L-2}\big|_{\xi'} + b_{L-1}(\xi'))\big), \quad \xi' = \varphi(x') && G \to \text{points}. \end{aligned} \]

\(P\), \(Q_1\), \(Q_2\) and \(W_\ell\) are pointwise linear maps; \(\mathcal{K}_0\) is SpectralConv.from_points (the weighted non-uniform Fourier transform \(\hat v_k = \sum_n w_n (P u)_n e^{-2\pi i k\cdot\xi_n}\), mixed and inverse-FFT-ed onto \(G\)), \(\mathcal{K}_\ell\) and \(W_\ell\) form a geonnax.fno.FNOBlock, and \(\mathcal{K}_{L-1}\) is SpectralConv.to_points (the truncated Fourier series evaluated at \(\xi'\)). The positional bias \(b_\ell\) is a 1×1 conv of the latent coordinates, added in every layer as in the reference (b0…b4 in elas_geofno.py): it tells the network where it is on the torus once \(\varphi\) moves points. With jacobian_weights=True the weights become the latent measure \(w^\xi_n = w_n |\det\nabla\varphi(x_n)| / \sum_m w_m |\det\nabla\varphi(x_m)|\) (off by default; jejjohnson/geonnax#85 decides).

Pseudocode of __call__ (H = hidden_channels):

ξ  = φ(x, code)                                         (N, d)
w  = weights                                            (N,)
if jacobian_weights: w = w·|det∇φ(x)| / Σ w·|det∇φ(x)|
h  = P u                                                (H, N)
h  = σ(encoder.from_points(h, ξ, w, G) + b_0(G))        (H, *G)
for ℓ, block in enumerate(blocks, 1):                   L - 2 blocks
    h = σ(block(h) + b_ℓ(G))   # block = K_ℓ h + W_ℓ h, no activation
ξ' = φ(x_out, code)                                     (M, d)
h  = decoder.to_points(h, ξ') + b_out(ξ')               (H, M)
y  = Q_2 σ(Q_1 h)                                       (C_out, M)

Ragged batches. Clouds of different sizes are batched by padding:

• pad every sample to \(N_{\max}\) points with any finite in-box coordinates, u = 0 and w = 0;
• jax.vmap(model)(u, x, w);
• mask the loss with w > 0.

A zero-weight point contributes exact zeros to the encoder, so the output at the real points does not depend on the padding. The latent grid and the modes are static: latent_shape and n_modes cannot vary per sample.

Deviations from the reference implementation (neuraloperator/Geo-FNO):

• quadrature weights: the reference encoder sums \(\sum_n u_n e^{-2\pi i k\cdot\xi_n}\) unweighted, which grows with \(N\); here \(w_n\) makes the coefficients a quadrature of the Fourier integral, so the operator converges as the cloud is refined;
• additive, zero-initialised chart: \(\varphi\) is exactly the box map at init for any bounding box (LearnedDeformation), where the reference's \(x + x\,\delta(x)\) is near the identity only near the origin;
• latent guard band: the box maps into \([0, b]^d\), \(b = 1/(1 + \mathrm{pad})\), so opposite edges of a non-periodic domain do not couple through the periodic latent FFT.

Cost per example: two dense non-uniform transforms, \(O(N\,|\mathcal{K}|\,H)\) with \(|\mathcal{K}|\) the retained modes, plus \(L - 2\) FFT blocks, \(O(H\, S \log S)\) with \(S = \prod_a s_a\). Memory is dominated by the \(N \times |\mathcal{K}|\) exponential table of the encoder; wrap the model in jax.checkpoint to trade it for compute.

Reference

Z. Li, D. Z. Huang, B. Liu and A. Anandkumar (2023), "Fourier Neural Operator with Learned Deformations for PDEs on General Geometries", Journal of Machine Learning Research 24(388), 1-26, arXiv:2207.05209.

Attributes:

Name	Type	Description
lifting	Conv	1-D 1×1 conv over the points, in_channels -> H.
encoder	SpectralConv	SpectralConv applied with SpectralConv.from_points.
blocks	list[FNOBlock]	The n_layers - 2 latent blocks (geonnax.fno.FNOBlock, built with use_activation=False: the activation follows the positional bias).
decoder	SpectralConv	SpectralConv applied with SpectralConv.to_points.
pos_bias	list[Conv] | None	One d-D 1×1 conv d -> H of the latent grid coordinates per latent layer (n_layers - 1), or None.
pos_bias_out	Conv | None	1-D 1×1 conv d -> H of \(\xi'\) at the query points, or None.
proj1,	proj2	1-D 1×1 convs H -> projection_channels -> out_channels.
deformation	LearnedDeformation	The chart \(\varphi\) (LearnedDeformation).
activation	Callable[[Array], Array]	Pointwise nonlinearity \(\sigma\).
latent_shape	tuple[int, ...]	Latent grid extents (s_1, …, s_d).
jacobian_weights	bool	Whether to reweight by \(|\det\nabla\varphi|\).
in_channels,	(out_channels, hidden_channels)	Channel counts.

Examples:

>>> import einx
>>> import jax
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> from geonnax.geofno import GeoFNO
>>> model = GeoFNO.init(
...     1, 2, (6, 6), (24, 24), key=jr.PRNGKey(0),
...     lower=[0.0, 0.0], upper=[10.0, 5.0], hidden_channels=8, n_layers=3,
...     deformation_kwargs={"width": 16, "depth": 1},
... )
>>> box = jnp.array([10.0, 5.0])
>>> x = einx.multiply("n d, d", jr.uniform(jr.PRNGKey(1), (300, 2)), box)
>>> u = einx.id("n 1 -> 1 n", jnp.sin(x[:, :1]))   # (C_in, N)
>>> w = jnp.full((300,), 1 / 300)                   # quadrature weights
>>> model(u, x, w).shape                            # (C_out, N)
(2, 300)
>>> model(u, x, w, x_out=x[:50]).shape              # (C_out, M)
(2, 50)

A ragged batch: clouds of 300 and 200 real points, zero-padded to 300:

>>> n_real = jnp.array([300, 200])
>>> real = einx.less("n, b -> b n", jnp.arange(300), n_real)  # (B, N)
>>> wb = einx.where("b n, b, -> b n", real, 1.0 / n_real, 0.0)
>>> ub = einx.where("b n, c n, -> b c n", real, u, 0.0)
>>> xb = jnp.stack([x, x])
>>> yb = jax.vmap(model)(ub, xb, wb)                # mask the loss by w > 0
>>> yb.shape
(2, 2, 300)

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

class GeoFNO(eqx.Module):
    r"""Geo-FNO: a Fourier neural operator on per-sample point clouds.

    Reads a field $u$ at $N$ points $x_n$ with quadrature weights $w_n$, maps
    the points onto the latent torus with the learned chart
    $\xi = \varphi(x)$ (`LearnedDeformation`), runs FNO blocks on a uniform
    latent grid $G$ (shape ``latent_shape``, nodes $j / s_a$ on $[0, 1)^d$),
    and evaluates the result at $M$ query points $x'_m$ (the input points by
    default):

    $$
    \begin{aligned}
    h_0 &= \sigma\big(\mathcal{K}_0[P u;\ \xi, w] + b_0(G)\big)
        && \text{points} \to G, \\
    h_\ell &= \sigma\big(\mathcal{K}_\ell h_{\ell-1} + W_\ell h_{\ell-1}
        + b_\ell(G)\big), \quad \ell = 1, \dots, L-2
        && \text{on } G, \\
    y &= Q_2\, \sigma\big(Q_1 (\mathcal{K}_{L-1} h_{L-2}\big|_{\xi'}
        + b_{L-1}(\xi'))\big), \quad \xi' = \varphi(x')
        && G \to \text{points}.
    \end{aligned}
    $$

    $P$, $Q_1$, $Q_2$ and $W_\ell$ are pointwise linear maps; $\mathcal{K}_0$
    is `SpectralConv.from_points` (the weighted non-uniform Fourier
    transform $\hat v_k = \sum_n w_n (P u)_n e^{-2\pi i k\cdot\xi_n}$,
    mixed and inverse-FFT-ed onto $G$), $\mathcal{K}_\ell$ and $W_\ell$ form
    a `geonnax.fno.FNOBlock`, and $\mathcal{K}_{L-1}$ is
    `SpectralConv.to_points` (the truncated Fourier series evaluated at
    $\xi'$). The positional bias $b_\ell$ is a 1×1 conv of the latent
    coordinates, added in every layer as in the reference
    (``b0``…``b4`` in ``elas_geofno.py``): it tells the network where it is
    on the torus once $\varphi$ moves points. With ``jacobian_weights=True``
    the weights become the latent measure
    $w^\xi_n = w_n |\det\nabla\varphi(x_n)| / \sum_m w_m
    |\det\nabla\varphi(x_m)|$ (off by default; jejjohnson/geonnax#85 decides).

    Pseudocode of `__call__` (``H = hidden_channels``):

    ```text
    ξ  = φ(x, code)                                         (N, d)
    w  = weights                                            (N,)
    if jacobian_weights: w = w·|det∇φ(x)| / Σ w·|det∇φ(x)|
    h  = P u                                                (H, N)
    h  = σ(encoder.from_points(h, ξ, w, G) + b_0(G))        (H, *G)
    for ℓ, block in enumerate(blocks, 1):                   L - 2 blocks
        h = σ(block(h) + b_ℓ(G))   # block = K_ℓ h + W_ℓ h, no activation
    ξ' = φ(x_out, code)                                     (M, d)
    h  = decoder.to_points(h, ξ') + b_out(ξ')               (H, M)
    y  = Q_2 σ(Q_1 h)                                       (C_out, M)
    ```

    **Ragged batches.** Clouds of different sizes are batched by padding:

    - pad every sample to $N_{\max}$ points with any finite in-box
      coordinates, ``u = 0`` and ``w = 0``;
    - ``jax.vmap(model)(u, x, w)``;
    - mask the loss with ``w > 0``.

    A zero-weight point contributes exact zeros to the encoder, so the
    output at the real points does not depend on the padding. The latent
    grid and the modes are static: ``latent_shape`` and ``n_modes`` cannot
    vary per sample.

    Deviations from the reference implementation (``neuraloperator/Geo-FNO``):

    - **quadrature weights**: the reference encoder sums
      $\sum_n u_n e^{-2\pi i k\cdot\xi_n}$ unweighted, which grows with $N$;
      here $w_n$ makes the coefficients a quadrature of the Fourier integral,
      so the operator converges as the cloud is refined;
    - **additive, zero-initialised chart**: $\varphi$ is exactly the box map
      at init for any bounding box (`LearnedDeformation`), where the
      reference's $x + x\,\delta(x)$ is near the identity only near the
      origin;
    - **latent guard band**: the box maps into $[0, b]^d$,
      $b = 1/(1 + \mathrm{pad})$, so opposite edges of a non-periodic domain
      do not couple through the periodic latent FFT.

    Cost per example: two dense non-uniform transforms,
    $O(N\,|\mathcal{K}|\,H)$ with $|\mathcal{K}|$ the retained modes, plus
    $L - 2$ FFT blocks, $O(H\, S \log S)$ with $S = \prod_a s_a$. Memory is
    dominated by the $N \times |\mathcal{K}|$ exponential table of the
    encoder; wrap the model in ``jax.checkpoint`` to trade it for compute.

    Reference:
        Z. Li, D. Z. Huang, B. Liu and A. Anandkumar (2023), "Fourier Neural
        Operator with Learned Deformations for PDEs on General Geometries",
        Journal of Machine Learning Research 24(388), 1-26, arXiv:2207.05209.

    Attributes:
        lifting: 1-D 1×1 conv over the points, ``in_channels -> H``.
        encoder: `SpectralConv` applied with `SpectralConv.from_points`.
        blocks: The ``n_layers - 2`` latent blocks (`geonnax.fno.FNOBlock`,
            built with ``use_activation=False``: the activation follows the
            positional bias).
        decoder: `SpectralConv` applied with `SpectralConv.to_points`.
        pos_bias: One ``d``-D 1×1 conv ``d -> H`` of the latent grid
            coordinates per latent layer (``n_layers - 1``), or ``None``.
        pos_bias_out: 1-D 1×1 conv ``d -> H`` of $\xi'$ at the query points,
            or ``None``.
        proj1, proj2: 1-D 1×1 convs ``H -> projection_channels -> out_channels``.
        deformation: The chart $\varphi$ (`LearnedDeformation`).
        activation: Pointwise nonlinearity $\sigma$.
        latent_shape: Latent grid extents ``(s_1, …, s_d)``.
        jacobian_weights: Whether to reweight by $|\det\nabla\varphi|$.
        in_channels, out_channels, hidden_channels: Channel counts.

    Examples:
        >>> import einx
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> from geonnax.geofno import GeoFNO
        >>> model = GeoFNO.init(
        ...     1, 2, (6, 6), (24, 24), key=jr.PRNGKey(0),
        ...     lower=[0.0, 0.0], upper=[10.0, 5.0], hidden_channels=8, n_layers=3,
        ...     deformation_kwargs={"width": 16, "depth": 1},
        ... )
        >>> box = jnp.array([10.0, 5.0])
        >>> x = einx.multiply("n d, d", jr.uniform(jr.PRNGKey(1), (300, 2)), box)
        >>> u = einx.id("n 1 -> 1 n", jnp.sin(x[:, :1]))   # (C_in, N)
        >>> w = jnp.full((300,), 1 / 300)                   # quadrature weights
        >>> model(u, x, w).shape                            # (C_out, N)
        (2, 300)
        >>> model(u, x, w, x_out=x[:50]).shape              # (C_out, M)
        (2, 50)

        A ragged batch: clouds of 300 and 200 real points, zero-padded to 300:

        >>> n_real = jnp.array([300, 200])
        >>> real = einx.less("n, b -> b n", jnp.arange(300), n_real)  # (B, N)
        >>> wb = einx.where("b n, b, -> b n", real, 1.0 / n_real, 0.0)
        >>> ub = einx.where("b n, c n, -> b c n", real, u, 0.0)
        >>> xb = jnp.stack([x, x])
        >>> yb = jax.vmap(model)(ub, xb, wb)                # mask the loss by w > 0
        >>> yb.shape
        (2, 2, 300)
    """

    lifting: eqx.nn.Conv
    encoder: SpectralConv
    blocks: list[FNOBlock]
    decoder: SpectralConv
    pos_bias: list[eqx.nn.Conv] | None
    pos_bias_out: eqx.nn.Conv | None
    proj1: eqx.nn.Conv
    proj2: eqx.nn.Conv
    deformation: LearnedDeformation
    activation: Callable[[Array], Array] = eqx.field(static=True)
    latent_shape: tuple[int, ...] = eqx.field(static=True)
    jacobian_weights: bool = eqx.field(static=True)
    in_channels: int = eqx.field(static=True)
    out_channels: int = eqx.field(static=True)
    hidden_channels: int = eqx.field(static=True)

    @classmethod
    def init(
        cls,
        in_channels: int,
        out_channels: int,
        n_modes: tuple[int, ...],
        latent_shape: tuple[int, ...],
        *,
        key: Array,
        lower: Sequence[float],
        upper: Sequence[float],
        hidden_channels: int = 32,
        n_layers: int = 4,
        projection_channels: int | None = None,
        deformation: DeformationKind = "mlp",
        deformation_kwargs: dict | None = None,
        positional_bias: bool = True,
        jacobian_weights: bool = False,
        factorization: Factorization = "dense",
        rank: float | int = 0.5,
        activation: Callable[[Array], Array] = jax.nn.gelu,
        code_dim: int = 0,
    ) -> GeoFNO:
        """Construct a Geo-FNO whose chart starts as the box map.

        Args:
            in_channels: Number of input channels ``C_in``.
            out_channels: Number of output channels ``C_out``.
            n_modes: Retained Fourier modes per latent axis; its length is the
                dimension ``d``.
            latent_shape: Latent grid extents, one per axis, each greater than
                ``2 * n_modes``.
            key: PRNG key.
            lower: Lower corner of the physical bounding box, length ``d``.
            upper: Upper corner of the physical bounding box, length ``d``.
            hidden_channels: Latent channel width ``H``.
            n_layers: Spectral layers, ``>= 2``: the encoder, ``n_layers - 2``
                latent blocks (`geonnax.fno.FNOBlock`) and the decoder.
            projection_channels: Hidden width of the projection (defaults to
                ``hidden_channels``).
            deformation: Kind of chart residual, ``"identity"``, ``"mlp"`` or
                ``"fourier"`` (see `LearnedDeformation`).
            deformation_kwargs: Extra keyword arguments for
                `LearnedDeformation.init` (``pad``, ``width``, ``depth``,
                ``n_frequencies``, ``n_basis``).
            positional_bias: Add the 1×1 conv of the latent coordinates in
                every layer (as the reference does).
            jacobian_weights: Reweight the quadrature by
                $|\\det\\nabla\\varphi|$ (jejjohnson/geonnax#85).
            factorization: Spectral-weight parameterisation (``"dense"`` /
                ``"cp"`` / ``"tucker"`` / ``"tt"``) of every spectral conv.
            rank: Rank for the factorised forms; ignored when dense.
            activation: Pointwise nonlinearity.
            code_dim: Length ``P`` of the per-sample chart code (``0``: none;
                needs ``deformation="mlp"``).

        Returns:
            A `GeoFNO`.

        Raises:
            ValueError: If ``n_layers < 2``; if ``n_modes`` is empty or
                ``n_modes``, ``latent_shape``, ``lower`` and ``upper`` differ
                in length; if some ``latent_shape[a] <= 2 * n_modes[a]``; or
                for an invalid chart configuration (`LearnedDeformation.init`).
        """
        if n_layers < 2:
            raise ValueError(f"n_layers must be >= 2, got {n_layers}.")
        n_modes = tuple(int(m) for m in n_modes)
        latent_shape = tuple(int(s) for s in latent_shape)
        d = len(n_modes)
        if d == 0 or not d == len(latent_shape) == len(lower) == len(upper):
            raise ValueError(
                "n_modes, latent_shape, lower and upper must be non-empty and of "
                f"equal length, got n_modes={n_modes}, latent_shape={latent_shape}, "
                f"lower={lower!r}, upper={upper!r}."
            )
        if any(2 * m >= s for m, s in zip(n_modes, latent_shape, strict=True)):
            raise ValueError(
                "each latent_shape extent must exceed 2 * n_modes = "
                f"{tuple(2 * m for m in n_modes)}, got latent_shape={latent_shape}."
            )
        projection_channels = projection_channels or hidden_channels

        (k_lift, k_enc, k_blocks, k_dec, k_pos, k_pos_out, k_p1, k_p2, k_def) = (
            jax.random.split(key, 9)
        )
        deform = LearnedDeformation.init(
            lower,
            upper,
            key=k_def,
            kind=deformation,
            code_dim=code_dim,
            **(deformation_kwargs or {}),
        )
        lifting = pointwise(1, in_channels, hidden_channels, key=k_lift)
        encoder, decoder = (
            SpectralConv.init(
                hidden_channels,
                hidden_channels,
                n_modes,
                key=k,
                factorization=factorization,
                rank=rank,
            )
            for k in (k_enc, k_dec)
        )
        blocks = [
            FNOBlock.init(
                hidden_channels,
                n_modes,
                key=k,
                factorization=factorization,
                rank=rank,
                activation=activation,
                use_activation=False,
            )
            for k in jax.random.split(k_blocks, n_layers - 2)
        ]
        pos_bias = pos_bias_out = None
        if positional_bias:
            pos_bias = [
                pointwise(d, d, hidden_channels, key=k)
                for k in jax.random.split(k_pos, n_layers - 1)
            ]
            pos_bias_out = pointwise(1, d, hidden_channels, key=k_pos_out)
        proj1 = pointwise(1, hidden_channels, projection_channels, key=k_p1)
        proj2 = pointwise(1, projection_channels, out_channels, key=k_p2)
        return cls(
            lifting=lifting,
            encoder=encoder,
            blocks=blocks,
            decoder=decoder,
            pos_bias=pos_bias,
            pos_bias_out=pos_bias_out,
            proj1=proj1,
            proj2=proj2,
            deformation=deform,
            activation=activation,
            latent_shape=latent_shape,
            jacobian_weights=jacobian_weights,
            in_channels=in_channels,
            out_channels=out_channels,
            hidden_channels=hidden_channels,
        )

    def _latent_coords(self, dtype: jnp.dtype) -> Float[Array, "d *G"]:
        """The latent grid nodes ``j / s_a`` on [0, 1)^d, channel-first."""
        axes = [jnp.arange(s, dtype=dtype) / s for s in self.latent_shape]
        return jnp.stack(jnp.meshgrid(*axes, indexing="ij"))

    def _add_pos(
        self, h: Float[Array, "H *G"], layer: int, grid: Float[Array, "d *G"]
    ) -> Float[Array, "H *G"]:
        """Add the positional bias $b_\\ell(G)$ of latent layer ``layer``, if any."""
        if self.pos_bias is None:
            return h
        return h + self.pos_bias[layer](grid)

    def __call__(
        self,
        u: Float[Array, "C_in N"],
        x: Float[Array, "N d"],
        weights: Float[Array, " N"],
        *,
        x_out: Float[Array, "M d"] | None = None,
        code: Float[Array, " P"] | None = None,
    ) -> Float[Array, "C_out M"]:
        """Apply the operator to one point cloud.

        Args:
            u: Field values at the points, ``(C_in, N)``; ``0`` at padding.
            x: Physical point coordinates, ``(N, d)``; finite (in-box) at
                padding.
            weights: Quadrature weights, ``(N,)``, approximating cell volumes
                ($w_n \\ge 0$, $\\sum_n w_n \\approx 1$); ``0`` marks padding.
            x_out: Physical query points, ``(M, d)``; ``None`` evaluates at
                ``x``.
            code: Per-sample chart code, ``(P,)``, exactly when
                ``code_dim > 0``.

        Returns:
            The output field at the query points, ``(C_out, M)``.

        Raises:
            ValueError: If the shapes of ``u``, ``x``, ``weights`` or
                ``x_out`` disagree with each other or with the model, or
                ``code`` does not match ``code_dim``.
        """
        d = len(self.latent_shape)
        if x.ndim != 2 or x.shape[1] != d:
            raise ValueError(f"x must have shape (N, {d}), got {x.shape}.")
        n = x.shape[0]
        if u.shape != (self.in_channels, n):
            raise ValueError(
                f"u must have shape (in_channels, N) = ({self.in_channels}, {n}), "
                f"got {u.shape}."
            )
        if weights.shape != (n,):
            raise ValueError(f"weights must have shape ({n},), got {weights.shape}.")
        if x_out is not None and (x_out.ndim != 2 or x_out.shape[1] != d):
            raise ValueError(f"x_out must have shape (M, {d}), got {x_out.shape}.")

        xi = self.deformation(x, code)  # (N, d)
        w = weights
        if self.jacobian_weights:
            w = w * jnp.abs(self.deformation.jacobian_det(x, code))
            w = w / jnp.maximum(einx.sum("[n]", w), jnp.finfo(w.dtype).tiny)
        h = self.lifting(u)  # (H, N)
        h = self.encoder.from_points(h, xi, w, self.latent_shape)  # (H, *G)
        grid = self._latent_coords(h.dtype)  # (d, *G)
        h = self.activation(self._add_pos(h, 0, grid))
        for i, block in enumerate(self.blocks, start=1):
            h = self.activation(self._add_pos(block(h), i, grid))  # K h + W h + b
        xi_out = xi if x_out is None else self.deformation(x_out, code)  # (M, d)
        h = self.decoder.to_points(h, xi_out)  # (H, M)
        if self.pos_bias_out is not None:
            h = h + self.pos_bias_out(einx.id("m d -> d m", xi_out))
        return self.proj2(self.activation(self.proj1(h)))  # (C_out, M)

init(in_channels: int, out_channels: int, n_modes: tuple[int, ...], latent_shape: tuple[int, ...], *, key: Array, lower: Sequence[float], upper: Sequence[float], hidden_channels: int = 32, n_layers: int = 4, projection_channels: int | None = None, deformation: DeformationKind = 'mlp', deformation_kwargs: dict | None = None, positional_bias: bool = True, jacobian_weights: bool = False, factorization: Factorization = 'dense', rank: float | int = 0.5, activation: Callable[[Array], Array] = jax.nn.gelu, code_dim: int = 0) -> GeoFNO classmethod

Construct a Geo-FNO whose chart starts as the box map.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels C_in.	required
out_channels	int	Number of output channels C_out.	required
n_modes	tuple[int, ...]	Retained Fourier modes per latent axis; its length is the dimension d.	required
latent_shape	tuple[int, ...]	Latent grid extents, one per axis, each greater than 2 * n_modes.	required
key	Array	PRNG key.	required
lower	Sequence[float]	Lower corner of the physical bounding box, length d.	required
upper	Sequence[float]	Upper corner of the physical bounding box, length d.	required
hidden_channels	int	Latent channel width H.	32
n_layers	int	Spectral layers, >= 2: the encoder, n_layers - 2 latent blocks (geonnax.fno.FNOBlock) and the decoder.	4
projection_channels	int | None	Hidden width of the projection (defaults to hidden_channels).	None
deformation	DeformationKind	Kind of chart residual, "identity", "mlp" or "fourier" (see LearnedDeformation).	'mlp'
deformation_kwargs	dict | None	Extra keyword arguments for LearnedDeformation.init (pad, width, depth, n_frequencies, n_basis).	None
positional_bias	bool	Add the 1×1 conv of the latent coordinates in every layer (as the reference does).	True
jacobian_weights	bool	Reweight the quadrature by \(|\det\nabla\varphi|\) (jejjohnson/geonnax#85).	False
factorization	Factorization	Spectral-weight parameterisation ("dense" / "cp" / "tucker" / "tt") of every spectral conv.	'dense'
rank	float | int	Rank for the factorised forms; ignored when dense.	0.5
activation	Callable[[Array], Array]	Pointwise nonlinearity.	gelu
code_dim	int	Length P of the per-sample chart code (0: none; needs deformation="mlp").	0

Returns:

Type	Description
GeoFNO	A GeoFNO.

Raises:

Type	Description
ValueError	If n_layers < 2; if n_modes is empty or n_modes, latent_shape, lower and upper differ in length; if some latent_shape[a] <= 2 * n_modes[a]; or for an invalid chart configuration (LearnedDeformation.init).

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

@classmethod
def init(
    cls,
    in_channels: int,
    out_channels: int,
    n_modes: tuple[int, ...],
    latent_shape: tuple[int, ...],
    *,
    key: Array,
    lower: Sequence[float],
    upper: Sequence[float],
    hidden_channels: int = 32,
    n_layers: int = 4,
    projection_channels: int | None = None,
    deformation: DeformationKind = "mlp",
    deformation_kwargs: dict | None = None,
    positional_bias: bool = True,
    jacobian_weights: bool = False,
    factorization: Factorization = "dense",
    rank: float | int = 0.5,
    activation: Callable[[Array], Array] = jax.nn.gelu,
    code_dim: int = 0,
) -> GeoFNO:
    """Construct a Geo-FNO whose chart starts as the box map.

    Args:
        in_channels: Number of input channels ``C_in``.
        out_channels: Number of output channels ``C_out``.
        n_modes: Retained Fourier modes per latent axis; its length is the
            dimension ``d``.
        latent_shape: Latent grid extents, one per axis, each greater than
            ``2 * n_modes``.
        key: PRNG key.
        lower: Lower corner of the physical bounding box, length ``d``.
        upper: Upper corner of the physical bounding box, length ``d``.
        hidden_channels: Latent channel width ``H``.
        n_layers: Spectral layers, ``>= 2``: the encoder, ``n_layers - 2``
            latent blocks (`geonnax.fno.FNOBlock`) and the decoder.
        projection_channels: Hidden width of the projection (defaults to
            ``hidden_channels``).
        deformation: Kind of chart residual, ``"identity"``, ``"mlp"`` or
            ``"fourier"`` (see `LearnedDeformation`).
        deformation_kwargs: Extra keyword arguments for
            `LearnedDeformation.init` (``pad``, ``width``, ``depth``,
            ``n_frequencies``, ``n_basis``).
        positional_bias: Add the 1×1 conv of the latent coordinates in
            every layer (as the reference does).
        jacobian_weights: Reweight the quadrature by
            $|\\det\\nabla\\varphi|$ (jejjohnson/geonnax#85).
        factorization: Spectral-weight parameterisation (``"dense"`` /
            ``"cp"`` / ``"tucker"`` / ``"tt"``) of every spectral conv.
        rank: Rank for the factorised forms; ignored when dense.
        activation: Pointwise nonlinearity.
        code_dim: Length ``P`` of the per-sample chart code (``0``: none;
            needs ``deformation="mlp"``).

    Returns:
        A `GeoFNO`.

    Raises:
        ValueError: If ``n_layers < 2``; if ``n_modes`` is empty or
            ``n_modes``, ``latent_shape``, ``lower`` and ``upper`` differ
            in length; if some ``latent_shape[a] <= 2 * n_modes[a]``; or
            for an invalid chart configuration (`LearnedDeformation.init`).
    """
    if n_layers < 2:
        raise ValueError(f"n_layers must be >= 2, got {n_layers}.")
    n_modes = tuple(int(m) for m in n_modes)
    latent_shape = tuple(int(s) for s in latent_shape)
    d = len(n_modes)
    if d == 0 or not d == len(latent_shape) == len(lower) == len(upper):
        raise ValueError(
            "n_modes, latent_shape, lower and upper must be non-empty and of "
            f"equal length, got n_modes={n_modes}, latent_shape={latent_shape}, "
            f"lower={lower!r}, upper={upper!r}."
        )
    if any(2 * m >= s for m, s in zip(n_modes, latent_shape, strict=True)):
        raise ValueError(
            "each latent_shape extent must exceed 2 * n_modes = "
            f"{tuple(2 * m for m in n_modes)}, got latent_shape={latent_shape}."
        )
    projection_channels = projection_channels or hidden_channels

    (k_lift, k_enc, k_blocks, k_dec, k_pos, k_pos_out, k_p1, k_p2, k_def) = (
        jax.random.split(key, 9)
    )
    deform = LearnedDeformation.init(
        lower,
        upper,
        key=k_def,
        kind=deformation,
        code_dim=code_dim,
        **(deformation_kwargs or {}),
    )
    lifting = pointwise(1, in_channels, hidden_channels, key=k_lift)
    encoder, decoder = (
        SpectralConv.init(
            hidden_channels,
            hidden_channels,
            n_modes,
            key=k,
            factorization=factorization,
            rank=rank,
        )
        for k in (k_enc, k_dec)
    )
    blocks = [
        FNOBlock.init(
            hidden_channels,
            n_modes,
            key=k,
            factorization=factorization,
            rank=rank,
            activation=activation,
            use_activation=False,
        )
        for k in jax.random.split(k_blocks, n_layers - 2)
    ]
    pos_bias = pos_bias_out = None
    if positional_bias:
        pos_bias = [
            pointwise(d, d, hidden_channels, key=k)
            for k in jax.random.split(k_pos, n_layers - 1)
        ]
        pos_bias_out = pointwise(1, d, hidden_channels, key=k_pos_out)
    proj1 = pointwise(1, hidden_channels, projection_channels, key=k_p1)
    proj2 = pointwise(1, projection_channels, out_channels, key=k_p2)
    return cls(
        lifting=lifting,
        encoder=encoder,
        blocks=blocks,
        decoder=decoder,
        pos_bias=pos_bias,
        pos_bias_out=pos_bias_out,
        proj1=proj1,
        proj2=proj2,
        deformation=deform,
        activation=activation,
        latent_shape=latent_shape,
        jacobian_weights=jacobian_weights,
        in_channels=in_channels,
        out_channels=out_channels,
        hidden_channels=hidden_channels,
    )

LearnedDeformation

Bases: Module

A learned chart from the physical domain onto the latent torus (Geo-FNO).

The chart is an affine box map plus a residual. With the physical bounding box \([\ell, h] \subset \mathbb{R}^d\) (lower, upper) and the latent extent \(b = 1/(1 + \mathrm{pad})\),

\[ s(x) = b\,\frac{x - \ell}{h - \ell} \in [0, b]^d, \qquad \xi = \varphi(x) = s(x) + \delta\big(s(x)\big). \]

The band \((b, 1)\) is a guard band: the latent FFT is periodic on \([0, 1)^d\), and the band keeps opposite edges of a non-periodic domain from coupling through the torus. The residual \(\delta\) depends on kind:

• "identity": \(\delta \equiv 0\).
• "mlp": \(\delta(s) = \mathrm{MLP}\big([s, \sin(Bs), \cos(Bs), c]\big)\), with NeRF features at the frozen frequencies \(B_j = \pi\,2^j\), \(j = 0, \dots, F - 1\), an optional per-sample code \(c\), and the final linear layer initialised to zero.
• "fourier": \(\delta(s) = \Phi(s - b/2)\,\beta\), where \(\Phi \in \mathbb{R}^{N \times M}\) holds the tensor-product Dirichlet Laplacian eigenfunctions on \([-b/2, b/2]^d\) (geonnax.basis.fourier_basis, \(M = \prod_a\) n_basis[a]) and \(\beta\) = coeffs \(\in \mathbb{R}^{M \times d}\) starts at zero. The Dirichlet condition gives \(\delta = 0\) on the box boundary, so the edges of the latent box stay fixed and the chart cannot push boundary points into the guard band.

The residual is additive and zero at init, so every kind starts as the exact box map: \(\varphi = s\), \(\det \nabla\varphi = \prod_a \mathrm{scale}_a\). Unlike the reference IPHI (a multiplicative residual \(x + x\,\delta(x)\), close to the identity only near the origin), this holds for any bounding box, e.g. lon/lat. As in the reference the chart maps physical to latent coordinates only, is never inverted and is not constrained to be invertible; jacobian_det is the diagnostic (\(\det \nabla\varphi \le 0\) somewhere means a fold-over).

lower, scale and frequencies are frozen leaves (wrapped in jax.lax.stop_gradient at use), so a gradient never moves the box or the feature frequencies. The field name coeffs is API: pyrox swaps it for a sample site with a spectral prior over the eigenvalues \(\lambda\) returned by fourier_basis (jejjohnson/pyrox#298).

Pseudocode of __call__:

s  = (x - lower) * scale                        (N, d) in [0, b]^d
identity: δ = 0
mlp:      f = [s, sin(s ⊗ B), cos(s ⊗ B), code]  (N, d + 2dF + P)
          δ = vmap(MLP)(f)                      (N, d)
fourier:  Φ, λ = fourier_basis(s - b/2, n_basis, L=b/2, "dirichlet")
          δ = Φ @ coeffs                        (N, M) @ (M, d)
ξ  = s + δ                                      (N, d)

Points outside the bounding box are mapped by the same formulas (the box map and the bases extrapolate); nothing clips them.

Attributes:

Name	Type	Description
lower	Float[Array, ' d']	Lower corner \(\ell\) of the physical box, (d,); frozen.
scale	Float[Array, ' d']	\(b / (h - \ell)\), (d,); frozen.
mlp	MLP | None	The residual MLP (kind="mlp"), else None.
frequencies	Float[Array, ' F'] | None	The NeRF frequencies \(\pi 2^j\), (F,) (kind="mlp"), else None; frozen.
coeffs	Float[Array, 'M d'] | None	Eigenbasis coefficients \(\beta\), (M, d) (kind="fourier"), else None. Named API.
kind	DeformationKind	"identity", "mlp" or "fourier".
n_basis	tuple[int, ...]	Eigenfunctions per axis (kind="fourier"), else ().
code_dim	int	Length P of the per-sample code (0: no code).
extent	float	The latent extent \(b = 1/(1 + \mathrm{pad})\).

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from geonnax.geofno import LearnedDeformation
>>> x = jax.random.uniform(jax.random.PRNGKey(1), (300, 2)) * 10.0
>>> phi = LearnedDeformation.init(
...     [0.0, 0.0], [10.0, 10.0], key=jax.random.PRNGKey(0), width=32, depth=2
... )
>>> xi = phi(x)                          # (N, d) -> (N, d)
>>> xi.shape
(300, 2)
>>> bool(jnp.all((xi >= 0.0) & (xi <= phi.extent)))  # inside [0, b]^d
True
>>> phi.jacobian_det(x[:3]).round(5)     # prod(scale) = (b / 10)^2 at init
Array([0.00826, 0.00826, 0.00826], dtype=float32)
>>> fphi = LearnedDeformation.init(
...     [0.0, 0.0], [10.0, 10.0], kind="fourier", n_basis=(4, 6)
... )
>>> fphi.coeffs.shape                    # (M, d), M = 4 * 6
(24, 2)
>>> cphi = LearnedDeformation.init(
...     [0.0, 0.0], [10.0, 10.0], key=jax.random.PRNGKey(0), width=32, depth=2,
...     code_dim=5,
... )
>>> cphi(x, jnp.ones(5)).shape           # per-sample code (P,)
(300, 2)
>>> batch = jnp.stack([x, x])            # (B, N, d)
>>> jax.vmap(phi)(batch).shape
(2, 300, 2)

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

class LearnedDeformation(eqx.Module):
    r"""A learned chart from the physical domain onto the latent torus (Geo-FNO).

    The chart is an affine box map plus a residual. With the physical bounding
    box $[\ell, h] \subset \mathbb{R}^d$ (``lower``, ``upper``) and the latent
    extent $b = 1/(1 + \mathrm{pad})$,

    $$
    s(x) = b\,\frac{x - \ell}{h - \ell} \in [0, b]^d,
    \qquad
    \xi = \varphi(x) = s(x) + \delta\big(s(x)\big).
    $$

    The band $(b, 1)$ is a guard band: the latent FFT is periodic on
    $[0, 1)^d$, and the band keeps opposite edges of a non-periodic domain from
    coupling through the torus. The residual $\delta$ depends on ``kind``:

    - ``"identity"``: $\delta \equiv 0$.
    - ``"mlp"``: $\delta(s) = \mathrm{MLP}\big([s, \sin(Bs), \cos(Bs), c]\big)$,
      with NeRF features at the frozen frequencies $B_j = \pi\,2^j$,
      $j = 0, \dots, F - 1$, an optional per-sample code $c$, and the final
      linear layer initialised to zero.
    - ``"fourier"``: $\delta(s) = \Phi(s - b/2)\,\beta$, where
      $\Phi \in \mathbb{R}^{N \times M}$ holds the tensor-product Dirichlet
      Laplacian eigenfunctions on $[-b/2, b/2]^d$
      (`geonnax.basis.fourier_basis`, $M = \prod_a$ ``n_basis[a]``) and
      $\beta$ = ``coeffs`` $\in \mathbb{R}^{M \times d}$ starts at zero. The
      Dirichlet condition gives $\delta = 0$ on the box boundary, so the
      edges of the latent box stay fixed and the chart cannot push boundary
      points into the guard band.

    The residual is additive and zero at init, so every kind starts as the
    exact box map: $\varphi = s$, $\det \nabla\varphi = \prod_a
    \mathrm{scale}_a$. Unlike the reference ``IPHI`` (a multiplicative
    residual $x + x\,\delta(x)$, close to the identity only near the origin),
    this holds for any bounding box, e.g. lon/lat. As in the reference the
    chart maps physical to latent coordinates only, is never inverted and is
    not constrained to be invertible; `jacobian_det` is the diagnostic
    ($\det \nabla\varphi \le 0$ somewhere means a fold-over).

    ``lower``, ``scale`` and ``frequencies`` are frozen leaves (wrapped in
    `jax.lax.stop_gradient` at use), so a gradient never moves the box or
    the feature frequencies. The field name ``coeffs`` is API: pyrox swaps it
    for a sample site with a spectral prior over the eigenvalues
    $\lambda$ returned by `fourier_basis` (jejjohnson/pyrox#298).

    Pseudocode of `__call__`:

    ```text
    s  = (x - lower) * scale                        (N, d) in [0, b]^d
    identity: δ = 0
    mlp:      f = [s, sin(s ⊗ B), cos(s ⊗ B), code]  (N, d + 2dF + P)
              δ = vmap(MLP)(f)                      (N, d)
    fourier:  Φ, λ = fourier_basis(s - b/2, n_basis, L=b/2, "dirichlet")
              δ = Φ @ coeffs                        (N, M) @ (M, d)
    ξ  = s + δ                                      (N, d)
    ```

    Points outside the bounding box are mapped by the same formulas (the box
    map and the bases extrapolate); nothing clips them.

    Attributes:
        lower: Lower corner $\ell$ of the physical box, ``(d,)``; frozen.
        scale: $b / (h - \ell)$, ``(d,)``; frozen.
        mlp: The residual MLP (``kind="mlp"``), else ``None``.
        frequencies: The NeRF frequencies $\pi 2^j$, ``(F,)``
            (``kind="mlp"``), else ``None``; frozen.
        coeffs: Eigenbasis coefficients $\beta$, ``(M, d)``
            (``kind="fourier"``), else ``None``. Named API.
        kind: ``"identity"``, ``"mlp"`` or ``"fourier"``.
        n_basis: Eigenfunctions per axis (``kind="fourier"``), else ``()``.
        code_dim: Length ``P`` of the per-sample code (``0``: no code).
        extent: The latent extent $b = 1/(1 + \mathrm{pad})$.

    Examples:
        >>> import jax
        >>> import jax.numpy as jnp
        >>> from geonnax.geofno import LearnedDeformation
        >>> x = jax.random.uniform(jax.random.PRNGKey(1), (300, 2)) * 10.0
        >>> phi = LearnedDeformation.init(
        ...     [0.0, 0.0], [10.0, 10.0], key=jax.random.PRNGKey(0), width=32, depth=2
        ... )
        >>> xi = phi(x)                          # (N, d) -> (N, d)
        >>> xi.shape
        (300, 2)
        >>> bool(jnp.all((xi >= 0.0) & (xi <= phi.extent)))  # inside [0, b]^d
        True
        >>> phi.jacobian_det(x[:3]).round(5)     # prod(scale) = (b / 10)^2 at init
        Array([0.00826, 0.00826, 0.00826], dtype=float32)
        >>> fphi = LearnedDeformation.init(
        ...     [0.0, 0.0], [10.0, 10.0], kind="fourier", n_basis=(4, 6)
        ... )
        >>> fphi.coeffs.shape                    # (M, d), M = 4 * 6
        (24, 2)
        >>> cphi = LearnedDeformation.init(
        ...     [0.0, 0.0], [10.0, 10.0], key=jax.random.PRNGKey(0), width=32, depth=2,
        ...     code_dim=5,
        ... )
        >>> cphi(x, jnp.ones(5)).shape           # per-sample code (P,)
        (300, 2)
        >>> batch = jnp.stack([x, x])            # (B, N, d)
        >>> jax.vmap(phi)(batch).shape
        (2, 300, 2)
    """

    lower: Float[Array, " d"]
    scale: Float[Array, " d"]
    mlp: eqx.nn.MLP | None
    frequencies: Float[Array, " F"] | None
    coeffs: Float[Array, "M d"] | None
    kind: DeformationKind = eqx.field(static=True)
    n_basis: tuple[int, ...] = eqx.field(static=True)
    code_dim: int = eqx.field(static=True)
    extent: float = eqx.field(static=True)

    @classmethod
    def init(
        cls,
        lower: Sequence[float],
        upper: Sequence[float],
        *,
        key: Array | None = None,
        kind: DeformationKind = "mlp",
        pad: float = 0.1,
        width: int = 128,
        depth: int = 4,
        n_frequencies: int = 8,
        n_basis: int | tuple[int, ...] = 8,
        code_dim: int = 0,
    ) -> LearnedDeformation:
        """Build a chart that is exactly the box map onto $[0, b]^d$.

        Args:
            lower: Lower corner of the physical bounding box, length ``d``.
            upper: Upper corner, length ``d``, with ``upper > lower``.
            key: PRNG key for the MLP weights; required for ``kind="mlp"``,
                unused otherwise (the other kinds draw nothing).
            kind: ``"identity"``, ``"mlp"`` or ``"fourier"``.
            pad: Guard-band fraction; the latent extent is
                ``b = 1 / (1 + pad)``. ``pad >= 0``.
            width: Hidden width of the MLP (``kind="mlp"``).
            depth: Number of hidden layers of the MLP (``kind="mlp"``).
            n_frequencies: Number ``F >= 1`` of NeRF frequencies $\\pi 2^j$
                (``kind="mlp"``).
            n_basis: Dirichlet eigenfunctions per axis (``kind="fourier"``);
                an ``int`` is broadcast to all ``d`` axes.
            code_dim: Length ``P`` of the per-sample code fed to the MLP;
                ``0`` for none. Only ``kind="mlp"`` takes a code.

        Returns:
            A `LearnedDeformation` whose residual is zero.

        Raises:
            ValueError: If ``lower`` / ``upper`` are empty, differ in length
                or have ``upper <= lower`` on some axis; if ``pad < 0``; if
                ``kind`` is unknown; if ``key`` is missing for
                ``kind="mlp"``; if ``width < 1``, ``depth < 0``,
                ``n_frequencies < 1``, ``code_dim < 0`` or an ``n_basis``
                entry is ``< 1`` or has the wrong length; or if
                ``code_dim > 0`` with a kind other than ``"mlp"``.

        Examples:
            >>> from geonnax.geofno import LearnedDeformation
            >>> phi = LearnedDeformation.init(
            ...     [-6.0, 35.0], [10.0, 45.0], kind="identity"
            ... )
            >>> phi.scale.round(4)                # b / (upper - lower)
            Array([0.0568, 0.0909], dtype=float32)
        """
        lo = [float(v) for v in lower]
        hi = [float(v) for v in upper]
        if not lo or len(lo) != len(hi):
            raise ValueError(
                "lower and upper must be non-empty and of equal length, got "
                f"lower={lower!r}, upper={upper!r}."
            )
        bad = [a for a, (l_a, h_a) in enumerate(zip(lo, hi, strict=True)) if h_a <= l_a]
        if bad:
            raise ValueError(
                f"upper must exceed lower on every axis; got lower={lo}, upper={hi} "
                f"(axes {bad})."
            )
        if not pad >= 0.0:
            raise ValueError(f"pad must be >= 0, got {pad}.")
        if kind not in _KINDS:
            raise ValueError(f"kind must be one of {_KINDS}, got {kind!r}.")
        if code_dim < 0:
            raise ValueError(f"code_dim must be >= 0, got {code_dim}.")
        if code_dim > 0 and kind != "mlp":
            raise ValueError(
                f"code_dim={code_dim} needs kind='mlp'; kind={kind!r} takes no code."
            )
        d = len(lo)
        b = 1.0 / (1.0 + float(pad))
        lower_arr = jnp.asarray(lo, dtype=float)
        scale = b / (jnp.asarray(hi, dtype=float) - lower_arr)

        mlp = frequencies = coeffs = None
        n_basis_t: tuple[int, ...] = ()
        if kind == "mlp":
            if key is None:
                raise ValueError("kind='mlp' needs a PRNG key, got key=None.")
            if width < 1:
                raise ValueError(f"width must be >= 1, got {width}.")
            if depth < 0:
                raise ValueError(f"depth must be >= 0, got {depth}.")
            if n_frequencies < 1:
                raise ValueError(f"n_frequencies must be >= 1, got {n_frequencies}.")
            frequencies = jnp.pi * 2.0 ** jnp.arange(n_frequencies, dtype=float)
            mlp = eqx.nn.MLP(
                in_size=d * (1 + 2 * n_frequencies) + code_dim,
                out_size=d,
                width_size=width,
                depth=depth,
                activation=jnp.tanh,
                key=key,
            )
            # Zero the readout so δ ≡ 0 and φ is the box map at init.
            readout = mlp.layers[-1]
            assert readout.bias is not None
            mlp = eqx.tree_at(
                lambda m: (m.layers[-1].weight, m.layers[-1].bias),
                mlp,
                (jnp.zeros_like(readout.weight), jnp.zeros_like(readout.bias)),
            )
        elif kind == "fourier":
            n_basis_t = (n_basis,) * d if isinstance(n_basis, int) else tuple(n_basis)
            if len(n_basis_t) != d or any(m < 1 for m in n_basis_t):
                raise ValueError(
                    f"n_basis must be an int >= 1 or {d} ints >= 1, got {n_basis!r}."
                )
            coeffs = jnp.zeros((math.prod(n_basis_t), d), dtype=float)

        return cls(
            lower=lower_arr,
            scale=scale,
            mlp=mlp,
            frequencies=frequencies,
            coeffs=coeffs,
            kind=kind,
            n_basis=n_basis_t,
            code_dim=code_dim,
            extent=b,
        )

    def box(self, x: Float[Array, "N d"]) -> Float[Array, "N d"]:
        """Apply the frozen affine box map $s(x) = (x - \\ell)\\,\\mathrm{scale}$.

        Args:
            x: Physical points, ``(N, d)``.

        Returns:
            Box-normalised points, ``(N, d)``, in $[0, b]^d$ for ``x`` in the box.
        """
        lower = jax.lax.stop_gradient(self.lower)
        scale = jax.lax.stop_gradient(self.scale)
        return einx.multiply("n d, d", einx.subtract("n d, d", x, lower), scale)

    def __call__(
        self, x: Float[Array, "N d"], code: Float[Array, " P"] | None = None
    ) -> Float[Array, "N d"]:
        """Map physical points to latent coordinates, $\\xi = s(x) + \\delta(s(x))$.

        Args:
            x: Physical points of one example, ``(N, d)``.
            code: Per-sample code, ``(P,)`` with ``P = code_dim``; must be
                ``None`` exactly when ``code_dim == 0``.

        Returns:
            Latent coordinates ``(N, d)``.

        Raises:
            ValueError: If ``x`` is not ``(N, d)`` or ``code`` does not match
                ``code_dim``.
        """
        d = self.lower.shape[0]
        if x.ndim != 2 or x.shape[1] != d:
            raise ValueError(f"x must have shape (N, {d}), got {x.shape}.")
        if (code is None) != (self.code_dim == 0):
            raise ValueError(
                "code must be given exactly when code_dim > 0; "
                f"code_dim={self.code_dim}, "
                f"code={'None' if code is None else code.shape}."
            )
        if code is not None and code.shape != (self.code_dim,):
            raise ValueError(
                f"code must have shape ({self.code_dim},), got {code.shape}."
            )
        s = self.box(x)
        if self.kind == "identity":
            return s
        if self.kind == "fourier":
            assert self.coeffs is not None
            half = 0.5 * self.extent
            phi, _ = fourier_basis(s - half, self.n_basis, half, boundary="dirichlet")
            return s + einx.dot("n m, m d -> n d", phi, self.coeffs)
        assert self.mlp is not None and self.frequencies is not None
        freqs = jax.lax.stop_gradient(self.frequencies)
        bs = einx.multiply("n d, f -> n (d f)", s, freqs)
        if code is None:
            feats = einx.id(
                "n d, n e, n e -> n (d + e + e)", s, jnp.sin(bs), jnp.cos(bs)
            )
        else:
            feats = einx.id(
                "n d, n e, n e, p -> n (d + e + e + p)",
                s,
                jnp.sin(bs),
                jnp.cos(bs),
                code,
            )
        return s + jax.vmap(self.mlp)(feats)

    def jacobian_det(
        self, x: Float[Array, "N d"], code: Float[Array, " P"] | None = None
    ) -> Float[Array, " N"]:
        """Per-point Jacobian determinant $\\det \\nabla\\varphi(x_n)$.

        A diagnostic: at init it equals $\\prod_a \\mathrm{scale}_a$ at every
        point, and $\\det \\nabla\\varphi \\le 0$ anywhere means the chart
        folds over there (no claim is made that it never does). It is also
        the change-of-measure factor $d\\xi = |\\det \\nabla\\varphi|\\,dx$ for
        latent quadrature weights (#85).

        Computed as ``jax.vmap(jax.jacfwd(φ))`` over single points: ``d``
        forward-mode JVPs of a one-point forward pass per point, plus an
        ``O(d^3)`` determinant, i.e. about ``d`` times the cost of `__call__`
        and ``(N, d, d)`` memory for the Jacobians.

        Args:
            x: Physical points, ``(N, d)``.
            code: Per-sample code, as for `__call__`.

        Returns:
            Determinants, ``(N,)``.
        """

        def point(xn: Float[Array, " d"]) -> Float[Array, " d"]:
            return self(einx.id("d -> 1 d", xn), code)[0]

        return jnp.linalg.det(jax.vmap(jax.jacfwd(point))(x))

init(lower: Sequence[float], upper: Sequence[float], *, key: Array | None = None, kind: DeformationKind = 'mlp', pad: float = 0.1, width: int = 128, depth: int = 4, n_frequencies: int = 8, n_basis: int | tuple[int, ...] = 8, code_dim: int = 0) -> LearnedDeformation classmethod

Build a chart that is exactly the box map onto \([0, b]^d\).

Parameters:

Name	Type	Description	Default
lower	Sequence[float]	Lower corner of the physical bounding box, length d.	required
upper	Sequence[float]	Upper corner, length d, with upper > lower.	required
key	Array | None	PRNG key for the MLP weights; required for kind="mlp", unused otherwise (the other kinds draw nothing).	None
kind	DeformationKind	"identity", "mlp" or "fourier".	'mlp'
pad	float	Guard-band fraction; the latent extent is b = 1 / (1 + pad). pad >= 0.	0.1
width	int	Hidden width of the MLP (kind="mlp").	128
depth	int	Number of hidden layers of the MLP (kind="mlp").	4
n_frequencies	int	Number F >= 1 of NeRF frequencies \(\pi 2^j\) (kind="mlp").	8
n_basis	int | tuple[int, ...]	Dirichlet eigenfunctions per axis (kind="fourier"); an int is broadcast to all d axes.	8
code_dim	int	Length P of the per-sample code fed to the MLP; 0 for none. Only kind="mlp" takes a code.	0

Returns:

Type	Description
LearnedDeformation	A LearnedDeformation whose residual is zero.

Raises:

Type	Description
ValueError	If lower / upper are empty, differ in length or have upper <= lower on some axis; if pad < 0; if kind is unknown; if key is missing for kind="mlp"; if width < 1, depth < 0, n_frequencies < 1, code_dim < 0 or an n_basis entry is < 1 or has the wrong length; or if code_dim > 0 with a kind other than "mlp".

Examples:

>>> from geonnax.geofno import LearnedDeformation
>>> phi = LearnedDeformation.init(
...     [-6.0, 35.0], [10.0, 45.0], kind="identity"
... )
>>> phi.scale.round(4)                # b / (upper - lower)
Array([0.0568, 0.0909], dtype=float32)

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

@classmethod
def init(
    cls,
    lower: Sequence[float],
    upper: Sequence[float],
    *,
    key: Array | None = None,
    kind: DeformationKind = "mlp",
    pad: float = 0.1,
    width: int = 128,
    depth: int = 4,
    n_frequencies: int = 8,
    n_basis: int | tuple[int, ...] = 8,
    code_dim: int = 0,
) -> LearnedDeformation:
    """Build a chart that is exactly the box map onto $[0, b]^d$.

    Args:
        lower: Lower corner of the physical bounding box, length ``d``.
        upper: Upper corner, length ``d``, with ``upper > lower``.
        key: PRNG key for the MLP weights; required for ``kind="mlp"``,
            unused otherwise (the other kinds draw nothing).
        kind: ``"identity"``, ``"mlp"`` or ``"fourier"``.
        pad: Guard-band fraction; the latent extent is
            ``b = 1 / (1 + pad)``. ``pad >= 0``.
        width: Hidden width of the MLP (``kind="mlp"``).
        depth: Number of hidden layers of the MLP (``kind="mlp"``).
        n_frequencies: Number ``F >= 1`` of NeRF frequencies $\\pi 2^j$
            (``kind="mlp"``).
        n_basis: Dirichlet eigenfunctions per axis (``kind="fourier"``);
            an ``int`` is broadcast to all ``d`` axes.
        code_dim: Length ``P`` of the per-sample code fed to the MLP;
            ``0`` for none. Only ``kind="mlp"`` takes a code.

    Returns:
        A `LearnedDeformation` whose residual is zero.

    Raises:
        ValueError: If ``lower`` / ``upper`` are empty, differ in length
            or have ``upper <= lower`` on some axis; if ``pad < 0``; if
            ``kind`` is unknown; if ``key`` is missing for
            ``kind="mlp"``; if ``width < 1``, ``depth < 0``,
            ``n_frequencies < 1``, ``code_dim < 0`` or an ``n_basis``
            entry is ``< 1`` or has the wrong length; or if
            ``code_dim > 0`` with a kind other than ``"mlp"``.

    Examples:
        >>> from geonnax.geofno import LearnedDeformation
        >>> phi = LearnedDeformation.init(
        ...     [-6.0, 35.0], [10.0, 45.0], kind="identity"
        ... )
        >>> phi.scale.round(4)                # b / (upper - lower)
        Array([0.0568, 0.0909], dtype=float32)
    """
    lo = [float(v) for v in lower]
    hi = [float(v) for v in upper]
    if not lo or len(lo) != len(hi):
        raise ValueError(
            "lower and upper must be non-empty and of equal length, got "
            f"lower={lower!r}, upper={upper!r}."
        )
    bad = [a for a, (l_a, h_a) in enumerate(zip(lo, hi, strict=True)) if h_a <= l_a]
    if bad:
        raise ValueError(
            f"upper must exceed lower on every axis; got lower={lo}, upper={hi} "
            f"(axes {bad})."
        )
    if not pad >= 0.0:
        raise ValueError(f"pad must be >= 0, got {pad}.")
    if kind not in _KINDS:
        raise ValueError(f"kind must be one of {_KINDS}, got {kind!r}.")
    if code_dim < 0:
        raise ValueError(f"code_dim must be >= 0, got {code_dim}.")
    if code_dim > 0 and kind != "mlp":
        raise ValueError(
            f"code_dim={code_dim} needs kind='mlp'; kind={kind!r} takes no code."
        )
    d = len(lo)
    b = 1.0 / (1.0 + float(pad))
    lower_arr = jnp.asarray(lo, dtype=float)
    scale = b / (jnp.asarray(hi, dtype=float) - lower_arr)

    mlp = frequencies = coeffs = None
    n_basis_t: tuple[int, ...] = ()
    if kind == "mlp":
        if key is None:
            raise ValueError("kind='mlp' needs a PRNG key, got key=None.")
        if width < 1:
            raise ValueError(f"width must be >= 1, got {width}.")
        if depth < 0:
            raise ValueError(f"depth must be >= 0, got {depth}.")
        if n_frequencies < 1:
            raise ValueError(f"n_frequencies must be >= 1, got {n_frequencies}.")
        frequencies = jnp.pi * 2.0 ** jnp.arange(n_frequencies, dtype=float)
        mlp = eqx.nn.MLP(
            in_size=d * (1 + 2 * n_frequencies) + code_dim,
            out_size=d,
            width_size=width,
            depth=depth,
            activation=jnp.tanh,
            key=key,
        )
        # Zero the readout so δ ≡ 0 and φ is the box map at init.
        readout = mlp.layers[-1]
        assert readout.bias is not None
        mlp = eqx.tree_at(
            lambda m: (m.layers[-1].weight, m.layers[-1].bias),
            mlp,
            (jnp.zeros_like(readout.weight), jnp.zeros_like(readout.bias)),
        )
    elif kind == "fourier":
        n_basis_t = (n_basis,) * d if isinstance(n_basis, int) else tuple(n_basis)
        if len(n_basis_t) != d or any(m < 1 for m in n_basis_t):
            raise ValueError(
                f"n_basis must be an int >= 1 or {d} ints >= 1, got {n_basis!r}."
            )
        coeffs = jnp.zeros((math.prod(n_basis_t), d), dtype=float)

    return cls(
        lower=lower_arr,
        scale=scale,
        mlp=mlp,
        frequencies=frequencies,
        coeffs=coeffs,
        kind=kind,
        n_basis=n_basis_t,
        code_dim=code_dim,
        extent=b,
    )

box(x: Float[Array, 'N d']) -> Float[Array, 'N d']

Apply the frozen affine box map \(s(x) = (x - \ell)\,\mathrm{scale}\).

Parameters:

Name	Type	Description	Default
x	Float[Array, 'N d']	Physical points, (N, d).	required

Returns:

Type	Description
Float[Array, 'N d']	Box-normalised points, (N, d), in \([0, b]^d\) for x in the box.

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

def box(self, x: Float[Array, "N d"]) -> Float[Array, "N d"]:
    """Apply the frozen affine box map $s(x) = (x - \\ell)\\,\\mathrm{scale}$.

    Args:
        x: Physical points, ``(N, d)``.

    Returns:
        Box-normalised points, ``(N, d)``, in $[0, b]^d$ for ``x`` in the box.
    """
    lower = jax.lax.stop_gradient(self.lower)
    scale = jax.lax.stop_gradient(self.scale)
    return einx.multiply("n d, d", einx.subtract("n d, d", x, lower), scale)

jacobian_det(x: Float[Array, 'N d'], code: Float[Array, ' P'] | None = None) -> Float[Array, ' N']

Per-point Jacobian determinant \(\det \nabla\varphi(x_n)\).

A diagnostic: at init it equals \(\prod_a \mathrm{scale}_a\) at every point, and \(\det \nabla\varphi \le 0\) anywhere means the chart folds over there (no claim is made that it never does). It is also the change-of-measure factor \(d\xi = |\det \nabla\varphi|\,dx\) for latent quadrature weights (#85).

Computed as jax.vmap(jax.jacfwd(φ)) over single points: d forward-mode JVPs of a one-point forward pass per point, plus an O(d^3) determinant, i.e. about d times the cost of __call__ and (N, d, d) memory for the Jacobians.

Parameters:

Name	Type	Description	Default
x	Float[Array, 'N d']	Physical points, (N, d).	required
code	Float[Array, ' P'] | None	Per-sample code, as for __call__.	None

Returns:

Type	Description
Float[Array, ' N']	Determinants, (N,).

Source code in .venv/lib/python3.12/site-packages/geonnax/geofno.py

def jacobian_det(
    self, x: Float[Array, "N d"], code: Float[Array, " P"] | None = None
) -> Float[Array, " N"]:
    """Per-point Jacobian determinant $\\det \\nabla\\varphi(x_n)$.

    A diagnostic: at init it equals $\\prod_a \\mathrm{scale}_a$ at every
    point, and $\\det \\nabla\\varphi \\le 0$ anywhere means the chart
    folds over there (no claim is made that it never does). It is also
    the change-of-measure factor $d\\xi = |\\det \\nabla\\varphi|\\,dx$ for
    latent quadrature weights (#85).

    Computed as ``jax.vmap(jax.jacfwd(φ))`` over single points: ``d``
    forward-mode JVPs of a one-point forward pass per point, plus an
    ``O(d^3)`` determinant, i.e. about ``d`` times the cost of `__call__`
    and ``(N, d, d)`` memory for the Jacobians.

    Args:
        x: Physical points, ``(N, d)``.
        code: Per-sample code, as for `__call__`.

    Returns:
        Determinants, ``(N,)``.
    """

    def point(xn: Float[Array, " d"]) -> Float[Array, " d"]:
        return self(einx.id("d -> 1 d", xn), code)[0]

    return jnp.linalg.det(jax.vmap(jax.jacfwd(point))(x))

ManifoldNO

Bases: Module

Manifold Neural Operator: lifting → basis blocks → projection.

With \(P\) the lifting, \(\mathcal K_\ell\) the BasisSpectralConv and \(W_\ell\) the 1×1 skip of block \(\ell\), \(Q_1, Q_2\) the projection and \(\sigma\) the activation,

\[ h_0 = P u, \qquad h_\ell = \sigma\big(\mathcal K_\ell(h_{\ell-1};\,T) + W_\ell h_{\ell-1}\big), \qquad y = Q_2\,\sigma(Q_1 h_L), \]

where the last block is linear (no \(\sigma\)), as in FNO. Each \(\mathcal K_\ell\) analyses into the \(M\) modes of the shared transform \(T\), mixes channels with one matrix per mode and synthesises back at the same \(N\) points (the NORM layer of Chen et al. 2024 in any basis).

Pseudocode:

h = lifting(u)                              # (H, N)
for block in blocks:                        # last block linear
    h = block(h, transform)                 # (H, N)
y = proj2(activation(proj1(h)))             # (C_out, N)

Operates on (in_channels, N) fields at the transform's N points; jax.vmap over a batch. The transform's tables are read through jax.lax.stop_gradient, so they never train (their gradient is zero).

Attributes:

Name	Type	Description
lifting	Conv	Pointwise map in_channels -> hidden_channels.
blocks	list[ManifoldNOBlock]	The ManifoldNOBlock layers.
proj1,	proj2	Two-layer pointwise projection hidden -> proj -> out.
transform	BasisTransform	The shared, frozen BasisTransform (N points in and out).
activation	Callable[[Array], Array]	Pointwise nonlinearity (static).
in_channels,	(out_channels, hidden_channels)	Static config.

Examples:

On the Laplace–Beltrami eigenbasis of a mesh (NORM), here the 162 vertices of an icosphere with the 9 modes of degree l <= 2:

>>> import jax, jax.numpy as jnp, jax.random as jr
>>> from geonnax.basis import icosphere
>>> from geonnax.layers import BasisTransform
>>> from geonnax.manifold_no import ManifoldNO
>>> transform = BasisTransform.from_mesh(*icosphere(2), num_basis=9)
>>> op = ManifoldNO.init(3, 1, transform, key=jr.PRNGKey(0),
...                      hidden_channels=16, n_layers=3)
>>> op(jnp.ones((3, 162))).shape
(1, 162)
>>> jax.vmap(op)(jnp.ones((4, 3, 162))).shape   # batched
(4, 1, 162)

On a region Slepian basis of a 60° polar cap:

>>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
>>> points, weights = gauss_legendre_grid(12)
>>> cap = points[:, 2] > 0.5
>>> basis = slepian_region_basis(5, points[cap], weights[cap])
>>> transform = BasisTransform.from_slepian(basis, points[cap], weights[cap])
>>> op = ManifoldNO.init(2, 2, transform, key=jr.PRNGKey(0),
...                      hidden_channels=8, n_layers=2)
>>> op(jnp.ones((2, 96))).shape
(2, 96)

Source code in .venv/lib/python3.12/site-packages/geonnax/manifold_no.py

class ManifoldNO(eqx.Module):
    r"""Manifold Neural Operator: lifting → basis blocks → projection.

    With $P$ the lifting, $\mathcal K_\ell$ the `BasisSpectralConv` and $W_\ell$
    the ``1×1`` skip of block $\ell$, $Q_1, Q_2$ the projection and $\sigma$ the
    activation,

    $$
    h_0 = P u, \qquad
    h_\ell = \sigma\big(\mathcal K_\ell(h_{\ell-1};\,T) + W_\ell h_{\ell-1}\big),
    \qquad y = Q_2\,\sigma(Q_1 h_L),
    $$

    where the last block is linear (no $\sigma$), as in `FNO`. Each
    $\mathcal K_\ell$ analyses into the $M$ modes of the shared transform $T$,
    mixes channels with one matrix per mode and synthesises back at the same
    $N$ points (the NORM layer of Chen et al. 2024 in any basis).

    Pseudocode:

    ```text
    h = lifting(u)                              # (H, N)
    for block in blocks:                        # last block linear
        h = block(h, transform)                 # (H, N)
    y = proj2(activation(proj1(h)))             # (C_out, N)
    ```

    Operates on ``(in_channels, N)`` fields at the transform's ``N`` points;
    ``jax.vmap`` over a batch. The transform's tables are read through
    `jax.lax.stop_gradient`, so they never train (their gradient is zero).

    Attributes:
        lifting: Pointwise map ``in_channels -> hidden_channels``.
        blocks: The `ManifoldNOBlock` layers.
        proj1, proj2: Two-layer pointwise projection ``hidden -> proj -> out``.
        transform: The shared, frozen `BasisTransform` (``N`` points in and
            out).
        activation: Pointwise nonlinearity (static).
        in_channels, out_channels, hidden_channels: Static config.

    Examples:
        On the Laplace–Beltrami eigenbasis of a mesh (NORM), here the 162
        vertices of an icosphere with the 9 modes of degree ``l <= 2``:

        >>> import jax, jax.numpy as jnp, jax.random as jr
        >>> from geonnax.basis import icosphere
        >>> from geonnax.layers import BasisTransform
        >>> from geonnax.manifold_no import ManifoldNO
        >>> transform = BasisTransform.from_mesh(*icosphere(2), num_basis=9)
        >>> op = ManifoldNO.init(3, 1, transform, key=jr.PRNGKey(0),
        ...                      hidden_channels=16, n_layers=3)
        >>> op(jnp.ones((3, 162))).shape
        (1, 162)
        >>> jax.vmap(op)(jnp.ones((4, 3, 162))).shape   # batched
        (4, 1, 162)

        On a region Slepian basis of a 60° polar cap:

        >>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
        >>> points, weights = gauss_legendre_grid(12)
        >>> cap = points[:, 2] > 0.5
        >>> basis = slepian_region_basis(5, points[cap], weights[cap])
        >>> transform = BasisTransform.from_slepian(basis, points[cap], weights[cap])
        >>> op = ManifoldNO.init(2, 2, transform, key=jr.PRNGKey(0),
        ...                      hidden_channels=8, n_layers=2)
        >>> op(jnp.ones((2, 96))).shape
        (2, 96)
    """

    lifting: eqx.nn.Conv
    blocks: list[ManifoldNOBlock]
    proj1: eqx.nn.Conv
    proj2: eqx.nn.Conv
    transform: BasisTransform
    activation: Callable[[Array], Array] = eqx.field(static=True)
    in_channels: int = eqx.field(static=True)
    out_channels: int = eqx.field(static=True)
    hidden_channels: int = eqx.field(static=True)

    @classmethod
    def init(
        cls,
        in_channels: int,
        out_channels: int,
        transform: BasisTransform,
        *,
        key: Array,
        hidden_channels: int = 32,
        n_layers: int = 4,
        projection_channels: int | None = None,
        activation: Callable[[Array], Array] = jax.nn.gelu,
    ) -> ManifoldNO:
        """Construct a manifold neural operator on ``transform``'s domain.

        Args:
            in_channels: Number of input channels.
            out_channels: Number of output channels.
            transform: The shared `BasisTransform`; its synthesis must target
                the analysis points (``N_out == N``) so the blocks stack.
            key: PRNG key.
            hidden_channels: Channel width ``H`` of the blocks.
            n_layers: Number of blocks ``L``.
            projection_channels: Hidden width of the projection MLP (defaults
                to ``hidden_channels``).
            activation: Pointwise nonlinearity.

        Returns:
            A `ManifoldNO`.

        Raises:
            ValueError: If ``n_layers < 1`` or the transform's output points
                differ from its input points.

        Examples:
            >>> import jax.numpy as jnp, jax.random as jr
            >>> from geonnax.layers import BasisTransform
            >>> from geonnax.manifold_no import ManifoldNO
            >>> phi, _ = jnp.linalg.qr(jr.normal(jr.PRNGKey(0), (50, 6)))
            >>> transform = BasisTransform.from_matrix(phi, jnp.ones(50))
            >>> op = ManifoldNO.init(1, 1, transform, key=jr.PRNGKey(1), n_layers=2)
            >>> len(op.blocks), op.hidden_channels
            (2, 32)
        """
        if n_layers < 1:
            raise ValueError(f"n_layers must be >= 1, got {n_layers}.")
        n_in = transform.analysis.shape[1]
        n_out = transform.synthesis.shape[0]
        if n_out != n_in:
            raise ValueError(
                "transform must synthesise at its analysis points so the blocks "
                f"stack (N_out == N); got N = {n_in}, N_out = {n_out}."
            )
        projection_channels = projection_channels or hidden_channels

        k_lift, k_p1, k_p2, k_blocks = jax.random.split(key, 4)
        lifting = pointwise(1, in_channels, hidden_channels, key=k_lift)
        bkeys = jax.random.split(k_blocks, n_layers)
        blocks = [
            ManifoldNOBlock.init(
                hidden_channels,
                transform.n_modes,
                key=bkeys[i],
                activation=activation,
                # The last block stays linear before projection, as in FNO.
                use_activation=i < n_layers - 1,
            )
            for i in range(n_layers)
        ]
        proj1 = pointwise(1, hidden_channels, projection_channels, key=k_p1)
        proj2 = pointwise(1, projection_channels, out_channels, key=k_p2)
        return cls(
            lifting=lifting,
            blocks=blocks,
            proj1=proj1,
            proj2=proj2,
            transform=transform,
            activation=activation,
            in_channels=in_channels,
            out_channels=out_channels,
            hidden_channels=hidden_channels,
        )

    def __call__(self, u: Float[Array, "C_in N"]) -> Float[Array, "C_out N"]:
        """Apply the operator to one ``(C_in, N)`` example."""
        h = self.lifting(u)
        for block in self.blocks:
            h = block(h, self.transform)
        h = self.activation(self.proj1(h))
        return self.proj2(h)

init(in_channels: int, out_channels: int, transform: BasisTransform, *, key: Array, hidden_channels: int = 32, n_layers: int = 4, projection_channels: int | None = None, activation: Callable[[Array], Array] = jax.nn.gelu) -> ManifoldNO classmethod

Construct a manifold neural operator on transform's domain.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required
transform	BasisTransform	The shared BasisTransform; its synthesis must target the analysis points (N_out == N) so the blocks stack.	required
key	Array	PRNG key.	required
hidden_channels	int	Channel width H of the blocks.	32
n_layers	int	Number of blocks L.	4
projection_channels	int | None	Hidden width of the projection MLP (defaults to hidden_channels).	None
activation	Callable[[Array], Array]	Pointwise nonlinearity.	gelu

Returns:

Type	Description
ManifoldNO	A ManifoldNO.

Raises:

Type	Description
ValueError	If n_layers < 1 or the transform's output points differ from its input points.

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from geonnax.layers import BasisTransform
>>> from geonnax.manifold_no import ManifoldNO
>>> phi, _ = jnp.linalg.qr(jr.normal(jr.PRNGKey(0), (50, 6)))
>>> transform = BasisTransform.from_matrix(phi, jnp.ones(50))
>>> op = ManifoldNO.init(1, 1, transform, key=jr.PRNGKey(1), n_layers=2)
>>> len(op.blocks), op.hidden_channels
(2, 32)

Source code in .venv/lib/python3.12/site-packages/geonnax/manifold_no.py

@classmethod
def init(
    cls,
    in_channels: int,
    out_channels: int,
    transform: BasisTransform,
    *,
    key: Array,
    hidden_channels: int = 32,
    n_layers: int = 4,
    projection_channels: int | None = None,
    activation: Callable[[Array], Array] = jax.nn.gelu,
) -> ManifoldNO:
    """Construct a manifold neural operator on ``transform``'s domain.

    Args:
        in_channels: Number of input channels.
        out_channels: Number of output channels.
        transform: The shared `BasisTransform`; its synthesis must target
            the analysis points (``N_out == N``) so the blocks stack.
        key: PRNG key.
        hidden_channels: Channel width ``H`` of the blocks.
        n_layers: Number of blocks ``L``.
        projection_channels: Hidden width of the projection MLP (defaults
            to ``hidden_channels``).
        activation: Pointwise nonlinearity.

    Returns:
        A `ManifoldNO`.

    Raises:
        ValueError: If ``n_layers < 1`` or the transform's output points
            differ from its input points.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from geonnax.layers import BasisTransform
        >>> from geonnax.manifold_no import ManifoldNO
        >>> phi, _ = jnp.linalg.qr(jr.normal(jr.PRNGKey(0), (50, 6)))
        >>> transform = BasisTransform.from_matrix(phi, jnp.ones(50))
        >>> op = ManifoldNO.init(1, 1, transform, key=jr.PRNGKey(1), n_layers=2)
        >>> len(op.blocks), op.hidden_channels
        (2, 32)
    """
    if n_layers < 1:
        raise ValueError(f"n_layers must be >= 1, got {n_layers}.")
    n_in = transform.analysis.shape[1]
    n_out = transform.synthesis.shape[0]
    if n_out != n_in:
        raise ValueError(
            "transform must synthesise at its analysis points so the blocks "
            f"stack (N_out == N); got N = {n_in}, N_out = {n_out}."
        )
    projection_channels = projection_channels or hidden_channels

    k_lift, k_p1, k_p2, k_blocks = jax.random.split(key, 4)
    lifting = pointwise(1, in_channels, hidden_channels, key=k_lift)
    bkeys = jax.random.split(k_blocks, n_layers)
    blocks = [
        ManifoldNOBlock.init(
            hidden_channels,
            transform.n_modes,
            key=bkeys[i],
            activation=activation,
            # The last block stays linear before projection, as in FNO.
            use_activation=i < n_layers - 1,
        )
        for i in range(n_layers)
    ]
    proj1 = pointwise(1, hidden_channels, projection_channels, key=k_p1)
    proj2 = pointwise(1, projection_channels, out_channels, key=k_p2)
    return cls(
        lifting=lifting,
        blocks=blocks,
        proj1=proj1,
        proj2=proj2,
        transform=transform,
        activation=activation,
        in_channels=in_channels,
        out_channels=out_channels,
        hidden_channels=hidden_channels,
    )

BasisTransform

Bases: Module

Frozen analysis / synthesis pair for a precomputed basis.

Holds the analysis matrix \(A = \Phi^\top \operatorname{diag}(w)\), shape (M, N), and the synthesis matrix \(\Phi_{\text{out}}\), shape (N_out, M). Both are array leaves read through jax.lax.stop_gradient, so they never train; build one transform and share it across every BasisSpectralConv of a model.

For a \(w\)-orthonormal basis (\(\Phi^\top \operatorname{diag}(w)\,\Phi = I\)) with \(\Phi_{\text{out}} = \Phi\), inverse(forward(u)) is the orthogonal projection of u onto \(\operatorname{span}(\Phi)\), so it reproduces any field already in that span.

Pseudocode:

forward(u):  â = u Aᵀ            # (C, N) → (C, M),     â_m ≈ ⟨u, φ_m⟩_w
inverse(â):  y = â Φ_outᵀ        # (C, M) → (C, N_out)

Attributes:

Name	Type	Description
analysis	Float[Array, 'M N']	(M, N) analysis matrix \(\Phi^\top \operatorname{diag}(w)\) (see from_slepian for the non-orthonormal Slepian case).
synthesis	Float[Array, 'N_out M']	(N_out, M) basis evaluated at the output points.
n_modes	int	Number of modes M (static).

Examples:

A \(w\)-orthonormal basis reproduces fields in its span:

>>> import jax.numpy as jnp, jax.random as jr
>>> from geonnax.layers import BasisTransform
>>> phi, _ = jnp.linalg.qr(jr.normal(jr.PRNGKey(0), (50, 6)))
>>> transform = BasisTransform.from_matrix(phi, jnp.ones(50))
>>> u = jr.normal(jr.PRNGKey(1), (2, 6)) @ phi.T   # (C, N) in span(Φ)
>>> transform.forward(u).shape
(2, 6)
>>> bool(jnp.allclose(transform.inverse(transform.forward(u)), u, atol=1e-5))
True

A Slepian basis of a 60° polar cap, analysed over the cap and synthesised anywhere on the sphere:

>>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
>>> points, weights = gauss_legendre_grid(12)
>>> cap = points[:, 2] > 0.5
>>> basis = slepian_region_basis(5, points[cap], weights[cap])
>>> transform = BasisTransform.from_slepian(
...     basis, points[cap], weights[cap], points_out_xyz=points)
>>> transform.analysis.shape, transform.synthesis.shape
((9, 96), (288, 9))

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

class BasisTransform(eqx.Module):
    r"""Frozen analysis / synthesis pair for a precomputed basis.

    Holds the analysis matrix $A = \Phi^\top \operatorname{diag}(w)$, shape
    ``(M, N)``, and the synthesis matrix $\Phi_{\text{out}}$, shape
    ``(N_out, M)``. Both are array leaves read through
    `jax.lax.stop_gradient`, so they never train; build one transform and
    share it across every `BasisSpectralConv` of a model.

    For a $w$-orthonormal basis ($\Phi^\top \operatorname{diag}(w)\,\Phi = I$)
    with $\Phi_{\text{out}} = \Phi$, ``inverse(forward(u))`` is the
    orthogonal projection of ``u`` onto $\operatorname{span}(\Phi)$, so it
    reproduces any field already in that span.

    Pseudocode:

    ```text
    forward(u):  â = u Aᵀ            # (C, N) → (C, M),     â_m ≈ ⟨u, φ_m⟩_w
    inverse(â):  y = â Φ_outᵀ        # (C, M) → (C, N_out)
    ```

    Attributes:
        analysis: ``(M, N)`` analysis matrix $\Phi^\top \operatorname{diag}(w)$
            (see `from_slepian` for the non-orthonormal Slepian case).
        synthesis: ``(N_out, M)`` basis evaluated at the output points.
        n_modes: Number of modes ``M`` (static).

    Examples:
        A $w$-orthonormal basis reproduces fields in its span:

        >>> import jax.numpy as jnp, jax.random as jr
        >>> from geonnax.layers import BasisTransform
        >>> phi, _ = jnp.linalg.qr(jr.normal(jr.PRNGKey(0), (50, 6)))
        >>> transform = BasisTransform.from_matrix(phi, jnp.ones(50))
        >>> u = jr.normal(jr.PRNGKey(1), (2, 6)) @ phi.T   # (C, N) in span(Φ)
        >>> transform.forward(u).shape
        (2, 6)
        >>> bool(jnp.allclose(transform.inverse(transform.forward(u)), u, atol=1e-5))
        True

        A Slepian basis of a 60° polar cap, analysed over the cap and
        synthesised anywhere on the sphere:

        >>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
        >>> points, weights = gauss_legendre_grid(12)
        >>> cap = points[:, 2] > 0.5
        >>> basis = slepian_region_basis(5, points[cap], weights[cap])
        >>> transform = BasisTransform.from_slepian(
        ...     basis, points[cap], weights[cap], points_out_xyz=points)
        >>> transform.analysis.shape, transform.synthesis.shape
        ((9, 96), (288, 9))
    """

    analysis: Float[Array, "M N"]
    synthesis: Float[Array, "N_out M"]
    n_modes: int = eqx.field(static=True)

    @classmethod
    def from_matrix(
        cls,
        phi: Float[Array, "N M"],
        weights: Float[Array, " N"],
        phi_out: Float[Array, "N_out M"] | None = None,
    ) -> BasisTransform:
        r"""Build the transform from a basis table and quadrature weights.

        $A = \Phi^\top \operatorname{diag}(w)$ and $\Phi_{\text{out}}$ =
        ``phi_out`` (default ``phi``, i.e. output at the input points). The
        analysis is the quadrature $\hat a_m = \sum_n w_n \phi_m(x_n) u(x_n)
        \approx \langle u, \phi_m \rangle$, so it inverts the synthesis only
        when $\Phi$ is $w$-orthonormal; this constructor does not
        orthonormalise.

        Args:
            phi: ``(N, M)`` basis at the ``N`` input points.
            weights: ``(N,)`` quadrature (mass) weights,
                $\sum_n w_n g(x_n) \approx \int g$.
            phi_out: ``(N_out, M)`` the same basis at the output points;
                defaults to ``phi``.

        Returns:
            A `BasisTransform` with ``M`` modes.

        Raises:
            ValueError: If ``phi`` is not rank 2, ``weights`` is not ``(N,)``,
                or ``phi_out`` does not have ``M`` columns.

        Examples:
            >>> import jax.numpy as jnp
            >>> from geonnax.layers import BasisTransform
            >>> transform = BasisTransform.from_matrix(jnp.eye(4)[:, :3], jnp.ones(4))
            >>> transform.n_modes, transform.synthesis.shape
            (3, (4, 3))
        """
        phi = jnp.asarray(phi)
        weights = jnp.asarray(weights)
        _check_table("phi", phi)
        if weights.shape != phi.shape[:1]:
            raise ValueError(f"weights must be ({phi.shape[0]},); got {weights.shape}.")
        n_modes = int(phi.shape[1])
        phi_out = phi if phi_out is None else jnp.asarray(phi_out)
        _check_table("phi_out", phi_out, n_modes)
        analysis = einx.multiply("n m, n -> m n", phi, weights)
        return cls(analysis=analysis, synthesis=phi_out, n_modes=n_modes)

    @classmethod
    def from_slepian(
        cls,
        basis: SlepianBasis,
        points_xyz: Float[Array, "N 3"],
        weights: Float[Array, " N"],
        points_out_xyz: Float[Array, "N_out 3"] | None = None,
    ) -> BasisTransform:
        r"""Build the transform from a region Slepian basis.

        ``points_xyz`` / ``weights`` should be a quadrature over the
        **region** (e.g. the one passed to `slepian_region_basis`). Slepian
        functions are orthonormal over the whole sphere but only orthogonal
        over the region, with $\sum_n w_n g_k(x_n) g_j(x_n) = \lambda_k
        \delta_{kj}$ (their concentrations; Simons et al. 2006). The analysis
        is the plain region-weighted projection,

        $$
        A = \Phi^\top \operatorname{diag}(w),
        \qquad
        A\Phi = \operatorname{diag}(\lambda),
        $$

        with no per-mode normalisation. A round trip ``inverse(forward(u))``
        therefore maps coefficients $c_k$ to $\lambda_k c_k$: it attenuates
        weakly concentrated modes by $\lambda_k$, which acts as a soft
        concentration taper, and is not the identity on the Slepian span. For
        an exact round trip, normalise $\Phi$ yourself (for example, divide
        each column by $\sqrt{\lambda_k}$) and build the transform with
        `from_matrix`.

        The synthesis evaluates the basis with ``basis.evaluate`` at
        ``points_out_xyz`` (default ``points_xyz``), which may be any points on
        the sphere.

        Args:
            basis: A `SlepianBasis` (anything with ``evaluate(xyz) -> (N, K)``).
            points_xyz: ``(N, 3)`` unit vectors of the region quadrature.
            weights: ``(N,)`` region quadrature weights.
            points_out_xyz: ``(N_out, 3)`` unit vectors for the output.

        Returns:
            A `BasisTransform` with ``K = basis.num_modes`` modes.

        Raises:
            ValueError: If the points are not ``(N, 3)`` or ``weights`` is not
                ``(N,)``.

        Examples:
            >>> import jax.numpy as jnp
            >>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
            >>> from geonnax.layers import BasisTransform
            >>> points, weights = gauss_legendre_grid(12)
            >>> cap = points[:, 2] > 0.5   # a cap of half-angle 60°
            >>> basis = slepian_region_basis(5, points[cap], weights[cap])
            >>> transform = BasisTransform.from_slepian(
            ...     basis, points[cap], weights[cap])
            >>> transform.n_modes, transform.synthesis.shape
            (9, (96, 9))
        """
        points = jnp.asarray(points_xyz)
        weights = jnp.asarray(weights)
        if points.ndim != 2 or points.shape[-1] != 3:
            raise ValueError(f"points_xyz must be (N, 3); got {points.shape}.")
        if weights.shape != points.shape[:1]:
            raise ValueError(
                f"weights must be ({points.shape[0]},); got {weights.shape}."
            )
        phi = basis.evaluate(points)
        if points_out_xyz is None:
            phi_out = phi
        else:
            out = jnp.asarray(points_out_xyz)
            if out.ndim != 2 or out.shape[-1] != 3:
                raise ValueError(f"points_out_xyz must be (N_out, 3); got {out.shape}.")
            phi_out = basis.evaluate(out)
        return cls.from_matrix(phi, weights, phi_out)

    @classmethod
    def from_mesh(
        cls,
        vertices: Float[Array, "V 3"],
        triangles: Int[Array, "T 3"],
        num_basis: int,
    ) -> BasisTransform:
        r"""Build the transform from a mesh's Laplace–Beltrami eigenbasis.

        $\Phi$ holds the ``num_basis`` lowest eigenvectors of the cotangent
        Laplacian (`cotangent_laplacian_eigpairs`), which are orthonormal
        under the lumped vertex-area mass, $\Phi^\top \operatorname{diag}(a)
        \,\Phi = I$. So $A = \Phi^\top \operatorname{diag}(a)$ with the
        vertex areas $a$ as quadrature weights, and the synthesis is $\Phi$:
        input and output are both the ``V`` vertices, and
        ``inverse(forward(u))`` is the $L^2$-orthogonal projection of a vertex
        field onto the low-frequency span. This is the NORM basis (Chen et al.
        2024).

        The eigenvectors come from a dense ``eigh``: their signs, and the basis
        chosen within a degenerate eigenspace (e.g. the $2\ell + 1$ degree-$\ell$
        modes of a sphere), are whatever the solver returns. Per-mode weights
        are tied to that order, so pick ``num_basis`` at the end of a
        degenerate cluster (9 = degrees $\ell \le 2$ on a sphere) and reuse
        the same transform when sharing weights. The construction is eager
        (NumPy, not traceable under ``jit``).

        Args:
            vertices: ``(V, 3)`` vertex positions.
            triangles: ``(T, 3)`` integer vertex indices.
            num_basis: Number of modes ``M``, ``1 <= M <= V``.

        Returns:
            A `BasisTransform` with ``M`` modes, ``(M, V)`` analysis and
            ``(V, M)`` synthesis.

        Raises:
            ValueError: As `cotangent_laplacian_eigpairs` (bad mesh or
                ``num_basis``).

        Examples:
            >>> from geonnax.basis import icosphere
            >>> from geonnax.layers import BasisTransform
            >>> vertices, triangles = icosphere(2)
            >>> transform = BasisTransform.from_mesh(vertices, triangles, num_basis=9)
            >>> transform.analysis.shape, transform.synthesis.shape
            ((9, 162), (162, 9))
        """
        _, phi, areas = cotangent_laplacian_eigpairs(vertices, triangles, num_basis)
        return cls.from_matrix(phi, areas)

    def forward(self, u: Float[Array, "C N"]) -> Float[Array, "C M"]:
        r"""Analyse ``(C, N)`` point values into ``(C, M)`` coefficients.

        Computes $\hat a = u A^\top$.
        """
        analysis = jax.lax.stop_gradient(self.analysis)
        return einx.dot("m n, c n -> c m", analysis, u)

    def inverse(self, a: Float[Array, "C M"]) -> Float[Array, "C N_out"]:
        r"""Synthesise ``(C, M)`` coefficients at the ``N_out`` output points.

        Computes $y = a\,\Phi_{\text{out}}^\top$.
        """
        synthesis = jax.lax.stop_gradient(self.synthesis)
        return einx.dot("p m, c m -> c p", synthesis, a)

from_matrix(phi: Float[Array, 'N M'], weights: Float[Array, ' N'], phi_out: Float[Array, 'N_out M'] | None = None) -> BasisTransform classmethod

Build the transform from a basis table and quadrature weights.

\(A = \Phi^\top \operatorname{diag}(w)\) and \(\Phi_{\text{out}}\) = phi_out (default phi, i.e. output at the input points). The analysis is the quadrature \(\hat a_m = \sum_n w_n \phi_m(x_n) u(x_n) \approx \langle u, \phi_m \rangle\), so it inverts the synthesis only when \(\Phi\) is \(w\)-orthonormal; this constructor does not orthonormalise.

Parameters:

Name	Type	Description	Default
phi	Float[Array, 'N M']	(N, M) basis at the N input points.	required
weights	Float[Array, ' N']	(N,) quadrature (mass) weights, \(\sum_n w_n g(x_n) \approx \int g\).	required
phi_out	Float[Array, 'N_out M'] | None	(N_out, M) the same basis at the output points; defaults to phi.	None

Returns:

Type	Description
BasisTransform	A BasisTransform with M modes.

Raises:

Type	Description
ValueError	If phi is not rank 2, weights is not (N,), or phi_out does not have M columns.

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.layers import BasisTransform
>>> transform = BasisTransform.from_matrix(jnp.eye(4)[:, :3], jnp.ones(4))
>>> transform.n_modes, transform.synthesis.shape
(3, (4, 3))

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

@classmethod
def from_matrix(
    cls,
    phi: Float[Array, "N M"],
    weights: Float[Array, " N"],
    phi_out: Float[Array, "N_out M"] | None = None,
) -> BasisTransform:
    r"""Build the transform from a basis table and quadrature weights.

    $A = \Phi^\top \operatorname{diag}(w)$ and $\Phi_{\text{out}}$ =
    ``phi_out`` (default ``phi``, i.e. output at the input points). The
    analysis is the quadrature $\hat a_m = \sum_n w_n \phi_m(x_n) u(x_n)
    \approx \langle u, \phi_m \rangle$, so it inverts the synthesis only
    when $\Phi$ is $w$-orthonormal; this constructor does not
    orthonormalise.

    Args:
        phi: ``(N, M)`` basis at the ``N`` input points.
        weights: ``(N,)`` quadrature (mass) weights,
            $\sum_n w_n g(x_n) \approx \int g$.
        phi_out: ``(N_out, M)`` the same basis at the output points;
            defaults to ``phi``.

    Returns:
        A `BasisTransform` with ``M`` modes.

    Raises:
        ValueError: If ``phi`` is not rank 2, ``weights`` is not ``(N,)``,
            or ``phi_out`` does not have ``M`` columns.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.layers import BasisTransform
        >>> transform = BasisTransform.from_matrix(jnp.eye(4)[:, :3], jnp.ones(4))
        >>> transform.n_modes, transform.synthesis.shape
        (3, (4, 3))
    """
    phi = jnp.asarray(phi)
    weights = jnp.asarray(weights)
    _check_table("phi", phi)
    if weights.shape != phi.shape[:1]:
        raise ValueError(f"weights must be ({phi.shape[0]},); got {weights.shape}.")
    n_modes = int(phi.shape[1])
    phi_out = phi if phi_out is None else jnp.asarray(phi_out)
    _check_table("phi_out", phi_out, n_modes)
    analysis = einx.multiply("n m, n -> m n", phi, weights)
    return cls(analysis=analysis, synthesis=phi_out, n_modes=n_modes)

from_slepian(basis: SlepianBasis, points_xyz: Float[Array, 'N 3'], weights: Float[Array, ' N'], points_out_xyz: Float[Array, 'N_out 3'] | None = None) -> BasisTransform classmethod

Build the transform from a region Slepian basis.

points_xyz / weights should be a quadrature over the region (e.g. the one passed to slepian_region_basis). Slepian functions are orthonormal over the whole sphere but only orthogonal over the region, with \(\sum_n w_n g_k(x_n) g_j(x_n) = \lambda_k \delta_{kj}\) (their concentrations; Simons et al. 2006). The analysis is the plain region-weighted projection,

\[ A = \Phi^\top \operatorname{diag}(w), \qquad A\Phi = \operatorname{diag}(\lambda), \]

with no per-mode normalisation. A round trip inverse(forward(u)) therefore maps coefficients \(c_k\) to \(\lambda_k c_k\): it attenuates weakly concentrated modes by \(\lambda_k\), which acts as a soft concentration taper, and is not the identity on the Slepian span. For an exact round trip, normalise \(\Phi\) yourself (for example, divide each column by \(\sqrt{\lambda_k}\)) and build the transform with from_matrix.

The synthesis evaluates the basis with basis.evaluate at points_out_xyz (default points_xyz), which may be any points on the sphere.

Parameters:

Name	Type	Description	Default
basis	SlepianBasis	A SlepianBasis (anything with evaluate(xyz) -> (N, K)).	required
points_xyz	Float[Array, 'N 3']	(N, 3) unit vectors of the region quadrature.	required
weights	Float[Array, ' N']	(N,) region quadrature weights.	required
points_out_xyz	Float[Array, 'N_out 3'] | None	(N_out, 3) unit vectors for the output.	None

Returns:

Type	Description
BasisTransform	A BasisTransform with K = basis.num_modes modes.

Raises:

Type	Description
ValueError	If the points are not (N, 3) or weights is not (N,).

Examples:

>>> import jax.numpy as jnp
>>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
>>> from geonnax.layers import BasisTransform
>>> points, weights = gauss_legendre_grid(12)
>>> cap = points[:, 2] > 0.5   # a cap of half-angle 60°
>>> basis = slepian_region_basis(5, points[cap], weights[cap])
>>> transform = BasisTransform.from_slepian(
...     basis, points[cap], weights[cap])
>>> transform.n_modes, transform.synthesis.shape
(9, (96, 9))

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

@classmethod
def from_slepian(
    cls,
    basis: SlepianBasis,
    points_xyz: Float[Array, "N 3"],
    weights: Float[Array, " N"],
    points_out_xyz: Float[Array, "N_out 3"] | None = None,
) -> BasisTransform:
    r"""Build the transform from a region Slepian basis.

    ``points_xyz`` / ``weights`` should be a quadrature over the
    **region** (e.g. the one passed to `slepian_region_basis`). Slepian
    functions are orthonormal over the whole sphere but only orthogonal
    over the region, with $\sum_n w_n g_k(x_n) g_j(x_n) = \lambda_k
    \delta_{kj}$ (their concentrations; Simons et al. 2006). The analysis
    is the plain region-weighted projection,

    $$
    A = \Phi^\top \operatorname{diag}(w),
    \qquad
    A\Phi = \operatorname{diag}(\lambda),
    $$

    with no per-mode normalisation. A round trip ``inverse(forward(u))``
    therefore maps coefficients $c_k$ to $\lambda_k c_k$: it attenuates
    weakly concentrated modes by $\lambda_k$, which acts as a soft
    concentration taper, and is not the identity on the Slepian span. For
    an exact round trip, normalise $\Phi$ yourself (for example, divide
    each column by $\sqrt{\lambda_k}$) and build the transform with
    `from_matrix`.

    The synthesis evaluates the basis with ``basis.evaluate`` at
    ``points_out_xyz`` (default ``points_xyz``), which may be any points on
    the sphere.

    Args:
        basis: A `SlepianBasis` (anything with ``evaluate(xyz) -> (N, K)``).
        points_xyz: ``(N, 3)`` unit vectors of the region quadrature.
        weights: ``(N,)`` region quadrature weights.
        points_out_xyz: ``(N_out, 3)`` unit vectors for the output.

    Returns:
        A `BasisTransform` with ``K = basis.num_modes`` modes.

    Raises:
        ValueError: If the points are not ``(N, 3)`` or ``weights`` is not
            ``(N,)``.

    Examples:
        >>> import jax.numpy as jnp
        >>> from geonnax.basis import gauss_legendre_grid, slepian_region_basis
        >>> from geonnax.layers import BasisTransform
        >>> points, weights = gauss_legendre_grid(12)
        >>> cap = points[:, 2] > 0.5   # a cap of half-angle 60°
        >>> basis = slepian_region_basis(5, points[cap], weights[cap])
        >>> transform = BasisTransform.from_slepian(
        ...     basis, points[cap], weights[cap])
        >>> transform.n_modes, transform.synthesis.shape
        (9, (96, 9))
    """
    points = jnp.asarray(points_xyz)
    weights = jnp.asarray(weights)
    if points.ndim != 2 or points.shape[-1] != 3:
        raise ValueError(f"points_xyz must be (N, 3); got {points.shape}.")
    if weights.shape != points.shape[:1]:
        raise ValueError(
            f"weights must be ({points.shape[0]},); got {weights.shape}."
        )
    phi = basis.evaluate(points)
    if points_out_xyz is None:
        phi_out = phi
    else:
        out = jnp.asarray(points_out_xyz)
        if out.ndim != 2 or out.shape[-1] != 3:
            raise ValueError(f"points_out_xyz must be (N_out, 3); got {out.shape}.")
        phi_out = basis.evaluate(out)
    return cls.from_matrix(phi, weights, phi_out)

from_mesh(vertices: Float[Array, 'V 3'], triangles: Int[Array, 'T 3'], num_basis: int) -> BasisTransform classmethod

Build the transform from a mesh's Laplace–Beltrami eigenbasis.

\(\Phi\) holds the num_basis lowest eigenvectors of the cotangent Laplacian (cotangent_laplacian_eigpairs), which are orthonormal under the lumped vertex-area mass, \(\Phi^\top \operatorname{diag}(a) \,\Phi = I\). So \(A = \Phi^\top \operatorname{diag}(a)\) with the vertex areas \(a\) as quadrature weights, and the synthesis is \(\Phi\): input and output are both the V vertices, and inverse(forward(u)) is the \(L^2\)-orthogonal projection of a vertex field onto the low-frequency span. This is the NORM basis (Chen et al. 2024).

The eigenvectors come from a dense eigh: their signs, and the basis chosen within a degenerate eigenspace (e.g. the \(2\ell + 1\) degree-\(\ell\) modes of a sphere), are whatever the solver returns. Per-mode weights are tied to that order, so pick num_basis at the end of a degenerate cluster (9 = degrees \(\ell \le 2\) on a sphere) and reuse the same transform when sharing weights. The construction is eager (NumPy, not traceable under jit).

Parameters:

Name	Type	Description	Default
vertices	Float[Array, 'V 3']	(V, 3) vertex positions.	required
triangles	Int[Array, 'T 3']	(T, 3) integer vertex indices.	required
num_basis	int	Number of modes M, 1 <= M <= V.	required

Returns:

Type	Description
BasisTransform	A BasisTransform with M modes, (M, V) analysis and
BasisTransform	(V, M) synthesis.

Raises:

Type	Description
ValueError	As cotangent_laplacian_eigpairs (bad mesh or num_basis).

Examples:

>>> from geonnax.basis import icosphere
>>> from geonnax.layers import BasisTransform
>>> vertices, triangles = icosphere(2)
>>> transform = BasisTransform.from_mesh(vertices, triangles, num_basis=9)
>>> transform.analysis.shape, transform.synthesis.shape
((9, 162), (162, 9))

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

@classmethod
def from_mesh(
    cls,
    vertices: Float[Array, "V 3"],
    triangles: Int[Array, "T 3"],
    num_basis: int,
) -> BasisTransform:
    r"""Build the transform from a mesh's Laplace–Beltrami eigenbasis.

    $\Phi$ holds the ``num_basis`` lowest eigenvectors of the cotangent
    Laplacian (`cotangent_laplacian_eigpairs`), which are orthonormal
    under the lumped vertex-area mass, $\Phi^\top \operatorname{diag}(a)
    \,\Phi = I$. So $A = \Phi^\top \operatorname{diag}(a)$ with the
    vertex areas $a$ as quadrature weights, and the synthesis is $\Phi$:
    input and output are both the ``V`` vertices, and
    ``inverse(forward(u))`` is the $L^2$-orthogonal projection of a vertex
    field onto the low-frequency span. This is the NORM basis (Chen et al.
    2024).

    The eigenvectors come from a dense ``eigh``: their signs, and the basis
    chosen within a degenerate eigenspace (e.g. the $2\ell + 1$ degree-$\ell$
    modes of a sphere), are whatever the solver returns. Per-mode weights
    are tied to that order, so pick ``num_basis`` at the end of a
    degenerate cluster (9 = degrees $\ell \le 2$ on a sphere) and reuse
    the same transform when sharing weights. The construction is eager
    (NumPy, not traceable under ``jit``).

    Args:
        vertices: ``(V, 3)`` vertex positions.
        triangles: ``(T, 3)`` integer vertex indices.
        num_basis: Number of modes ``M``, ``1 <= M <= V``.

    Returns:
        A `BasisTransform` with ``M`` modes, ``(M, V)`` analysis and
        ``(V, M)`` synthesis.

    Raises:
        ValueError: As `cotangent_laplacian_eigpairs` (bad mesh or
            ``num_basis``).

    Examples:
        >>> from geonnax.basis import icosphere
        >>> from geonnax.layers import BasisTransform
        >>> vertices, triangles = icosphere(2)
        >>> transform = BasisTransform.from_mesh(vertices, triangles, num_basis=9)
        >>> transform.analysis.shape, transform.synthesis.shape
        ((9, 162), (162, 9))
    """
    _, phi, areas = cotangent_laplacian_eigpairs(vertices, triangles, num_basis)
    return cls.from_matrix(phi, areas)

forward(u: Float[Array, 'C N']) -> Float[Array, 'C M']

Analyse (C, N) point values into (C, M) coefficients.

Computes \(\hat a = u A^\top\).

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

def forward(self, u: Float[Array, "C N"]) -> Float[Array, "C M"]:
    r"""Analyse ``(C, N)`` point values into ``(C, M)`` coefficients.

    Computes $\hat a = u A^\top$.
    """
    analysis = jax.lax.stop_gradient(self.analysis)
    return einx.dot("m n, c n -> c m", analysis, u)

inverse(a: Float[Array, 'C M']) -> Float[Array, 'C N_out']

Synthesise (C, M) coefficients at the N_out output points.

Computes \(y = a\,\Phi_{\text{out}}^\top\).

Source code in .venv/lib/python3.12/site-packages/geonnax/layers/_basis_conv.py

def inverse(self, a: Float[Array, "C M"]) -> Float[Array, "C N_out"]:
    r"""Synthesise ``(C, M)`` coefficients at the ``N_out`` output points.

    Computes $y = a\,\Phi_{\text{out}}^\top$.
    """
    synthesis = jax.lax.stop_gradient(self.synthesis)
    return einx.dot("p m, c m -> c p", synthesis, a)