NN API

The pyrox.nn subpackage ships uncertainty-aware neural network layers in four families:

1. Geographic / spherical encoders (pyrox.nn._geo plus deterministic wrappers in pyrox.nn._layers) — degree/radian, lon/lat, cyclic, and spherical-harmonic preprocessing for geophysical inputs.
2. Dense / Bayesian-linear layers (pyrox.nn._layers) — twelve layers covering reparameterization, Flipout, NCP, MC-Dropout, and several random-feature variants.
3. Bayesian Neural Field stack (pyrox.nn._bnf) — five layers that together implement the BNF architecture (Saad et al., Nat. Comms. 2024).
4. Pure-JAX feature helpers (pyrox.nn._features) — 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

pyrox.nn.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.

.. math::

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 src/pyrox/nn/_layers.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.

    .. math::

        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
    out_features: int
    bias: bool = True
    prior_scale: float = 1.0
    pyrox_name: str | None = 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 = 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

pyrox.nn.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 :class: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 src/pyrox/nn/_layers.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
    :class:`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
    out_features: int
    bias: bool = True
    prior_scale: float = 1.0
    pyrox_name: str | None = 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 = 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

pyrox.nn.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 src/pyrox/nn/_layers.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
    out_features: int
    make_prior: Callable[..., Any] = eqx.field(static=True)
    bias: bool = True
    pyrox_name: str | None = 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 = 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

pyrox.nn.MCDropout

Bases: Module

Always-on dropout for Monte Carlo uncertainty estimation.

Unlike standard dropout, :class:MCDropout stays active at inference time — repeated forward passes with different keys produce a distribution of outputs whose spread approximates predictive uncertainty (Gal & Ghahramani, 2016).

Not a :class:PyroxModule — no NumPyro sites are registered. The stochasticity comes from the explicit PRNG key argument.

Attributes:

Name	Type	Description
rate	float	Dropout probability in :math:(0, 1).

Source code in src/pyrox/nn/_layers.py

class MCDropout(eqx.Module):
    """Always-on dropout for Monte Carlo uncertainty estimation.

    Unlike standard dropout, :class:`MCDropout` stays active at
    inference time — repeated forward passes with different keys
    produce a distribution of outputs whose spread approximates
    predictive uncertainty (Gal & Ghahramani, 2016).

    Not a :class:`PyroxModule` — no NumPyro sites are registered.
    The stochasticity comes from the explicit PRNG ``key`` argument.

    Attributes:
        rate: Dropout probability in :math:`(0, 1)`.
    """

    rate: float = 0.5

    def __post_init__(self) -> None:
        if not 0.0 <= self.rate < 1.0:
            raise ValueError(f"rate must be in [0, 1), got {self.rate}.")

    def __call__(
        self,
        x: Float[Array, ...],
        *,
        key: Array,
    ) -> Float[Array, ...]:
        """Apply dropout, scaling survivors by ``1 / (1 - rate)``."""
        keep = jax.random.bernoulli(key, 1.0 - self.rate, x.shape)
        return jnp.where(keep, x / (1.0 - self.rate), 0.0)

__call__(x, *, key)

Apply dropout, scaling survivors by 1 / (1 - rate).

Source code in src/pyrox/nn/_layers.py

def __call__(
    self,
    x: Float[Array, ...],
    *,
    key: Array,
) -> Float[Array, ...]:
    """Apply dropout, scaling survivors by ``1 / (1 - rate)``."""
    keep = jax.random.bernoulli(key, 1.0 - self.rate, x.shape)
    return jnp.where(keep, x / (1.0 - self.rate), 0.0)

pyrox.nn.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:

.. math::

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 :math:\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 :math:\sigma.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in src/pyrox/nn/_layers.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:

    .. math::

        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 :math:`\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
            :math:`\sigma`.
        pyrox_name: Explicit scope name for NumPyro site registration.
    """

    in_features: int
    out_features: int
    init_scale: float = 1.0
    pyrox_name: str | None = 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 = 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 * (x @ W_s + b_s)

        return det + stoch

pyrox.nn.NCPContinuousPerturb

Bases: Module

Input perturbation for the Noise Contrastive Prior pattern.

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

.. math::

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

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

Not a :class:PyroxModule — stochasticity comes from the explicit PRNG key.

Attributes:

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

Source code in src/pyrox/nn/_layers.py

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

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

    .. math::

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

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

    Not a :class:`PyroxModule` — stochasticity comes from the
    explicit PRNG ``key``.

    Attributes:
        scale: Perturbation scale :math:`\sigma`.
    """

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

    def __call__(
        self,
        x: Float[Array, "*batch D"],
        *,
        key: Array,
    ) -> Float[Array, "*batch D"]:
        eps = jax.random.normal(key, x.shape, dtype=x.dtype)
        return x + self.scale * eps

pyrox.nn.RBFFourierFeatures

Bases: PyroxModule

SSGP-style RFF layer with RBF spectral density.

Both the spectral frequencies :math:W and the lengthscale :math:\ell are pyrox_sample sites — :math:W has a standard normal prior (the RBF spectral density) and :math:\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 src/pyrox/nn/_layers.py

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

    Both the spectral frequencies :math:`W` and the lengthscale
    :math:`\ell` are ``pyrox_sample`` sites — :math:`W` has a
    standard normal prior (the RBF spectral density) and :math:`\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 _rff_forward(W, ls, self.n_features, x)

pyrox.nn.RBFCosineFeatures

Bases: PyroxModule

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

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

.. math::

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

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

All parameters (:math:W, :math:b, :math:\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 src/pyrox/nn/_layers.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:

    .. math::

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

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

    All parameters (:math:`W`, :math:`b`, :math:`\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:
        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 _rff_cosine_forward(W, b, ls, self.n_features, x)

pyrox.nn.MaternFourierFeatures

Bases: PyroxModule

SSGP-style RFF layer with Matern spectral density.

Spectral frequencies :math:W have a StudentT(df=2\nu) prior (the Matern spectral density). The smoothness :math:\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 :math:\nu.
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in src/pyrox/nn/_layers.py

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

    Spectral frequencies :math:`W` have a ``StudentT(df=2\nu)`` prior
    (the Matern spectral density). The smoothness :math:`\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 :math:`\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 _rff_forward(W, ls, self.n_features, x)

pyrox.nn.MaternCosineFeatures

Bases: PyroxModule

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

Single-cosine analogue of :class:MaternFourierFeatures:

.. math::

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

where :math:W \sim \mathrm{StudentT}(2\nu) (the Matern spectral density) and :math: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 :class:MaternFourierFeatures in expectation but with higher variance per draw — see Sutherland & Schneider (2015).

All parameters (:math:W, :math:b, :math:\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 :math:\nu.
init_lengthscale	float	Prior location for the lengthscale.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in src/pyrox/nn/_layers.py

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

    Single-cosine analogue of :class:`MaternFourierFeatures`:

    .. math::

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

    where :math:`W \sim \mathrm{StudentT}(2\nu)` (the Matern spectral
    density) and :math:`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
    :class:`MaternFourierFeatures` in expectation but with higher
    variance per draw — see Sutherland & Schneider (2015).

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

    Attributes:
        in_features: Input dimension.
        n_features: Number of random features (= output dimension).
        nu: Smoothness parameter :math:`\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 _rff_cosine_forward(W, b, ls, self.n_features, x)

pyrox.nn.LaplaceFourierFeatures

Bases: PyroxModule

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

Spectral frequencies :math: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 src/pyrox/nn/_layers.py

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

    Spectral frequencies :math:`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:
        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 _rff_forward(W, ls, self.n_features, x)

pyrox.nn.LaplaceCosineFeatures

Bases: PyroxModule

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

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

.. math::

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

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

All parameters (:math:W, :math:b, :math:\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 src/pyrox/nn/_layers.py

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

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

    .. math::

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

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

    All parameters (:math:`W`, :math:`b`, :math:`\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 _rff_cosine_forward(W, b, ls, self.n_features, x)

pyrox.nn.ArcCosineFourierFeatures

Bases: PyroxModule

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

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

.. math::

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

where :math: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 src/pyrox/nn/_layers.py

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

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

    .. math::

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

    where :math:`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 = 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

pyrox.nn.RandomKitchenSinks

Bases: PyroxModule

Random Kitchen Sinks: RFF + a learned linear head.

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

.. math::

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 src/pyrox/nn/_layers.py

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

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

    .. math::

        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 phi @ beta + bias

init(rff, out_features) classmethod

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

Source code in src/pyrox/nn/_layers.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)

pyrox.nn.VariationalFourierFeatures

Bases: PyroxModule

VSSGP — RFF with a learnable variational posterior over frequencies.

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

Prior: :math: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 :class:~numpyro.infer.autoguide.AutoNormal to learn the posterior on W; under prior-only seeds, behaves identically to :class:RBFFourierFeatures.

Attributes:

Name	Type	Description
in_features	int	Input dimension :math: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 src/pyrox/nn/_layers.py

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

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

    Prior: :math:`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 :class:`~numpyro.infer.autoguide.AutoNormal` to
    learn the posterior on ``W``; under prior-only seeds, behaves
    identically to :class:`RBFFourierFeatures`.

    Attributes:
        in_features: Input dimension :math:`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 _rff_forward(W, ls, self.n_features, x)

pyrox.nn.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 :class: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 :math: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, 2 * n_features // 2).

Note

For learnable lengthscale or full Bayesian treatment of the frequencies, prefer :class:VariationalFourierFeatures.

Source code in src/pyrox/nn/_layers.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 :class:`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 :math:`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, 2 * n_features // 2)``.

    Note:
        For learnable lengthscale or full Bayesian treatment of the
        frequencies, prefer :class:`VariationalFourierFeatures`.
    """

    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:
        if lengthscale <= 0:
            raise ValueError(f"lengthscale must be > 0, got {lengthscale}.")
        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, "*batch D_in"]) -> Float[Array, "*batch D_rff"]:
        return _rff_forward(self.W, self.lengthscale, self.n_features, x)

pyrox.nn.HSGPFeatures

Bases: PyroxModule

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

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

.. math::

\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 :class: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 :math: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 :mod:pyrox.gp whose spectral density supplies the per-basis prior variance. Currently :class:pyrox.gp.RBF and :class:pyrox.gp.Matern are supported by :func:pyrox._basis.spectral_density.
pyrox_name	str | None	Explicit scope name for NumPyro site registration.

Source code in src/pyrox/nn/_layers.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
    :math:`[-L, L]^D` plus learnable per-basis amplitudes with a
    kernel-spectral-density prior:

    .. math::

        \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 :class:`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 :math:`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 :mod:`pyrox.gp` whose spectral
            density supplies the per-basis prior variance. Currently
            :class:`pyrox.gp.RBF` and :class:`pyrox.gp.Matern` are
            supported by :func:`pyrox._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 jnp.einsum("nm,m->n", Phi, sqrt_S * alpha)

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.

pyrox.nn.SirenDense

Bases: Module

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

Single primitive of a SIREN network. 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).

Example

import jax.random as jr layer = SirenDense.init(3, 16, key=jr.PRNGKey(0), layer_type="first") import jax.numpy as jnp y = layer(jnp.ones((5, 3))) y.shape (5, 16)

Source code in src/pyrox/nn/_layers.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.  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).

    Example:
        >>> import jax.random as jr
        >>> layer = SirenDense.init(3, 16, key=jr.PRNGKey(0), layer_type="first")
        >>> import jax.numpy as jnp
        >>> y = layer(jnp.ones((5, 3)))
        >>> y.shape
        (5, 16)
    """

    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.

        Args:
            in_features: Input dimension.
            out_features: Output dimension.
            key: JAX PRNG key for weight and bias sampling.
            omega: Frequency multiplier (default 30.0, as in Sitzmann et al.).
            layer_type: Init regime — ``"first"``, ``"hidden"``, or ``"last"``.
            c: Theorem-1 constant (default 6.0).

        Returns:
            Initialised :class:`SirenDense`.

        Raises:
            ValueError: If ``layer_type`` is not a valid regime name, or if any of
                ``in_features``, ``out_features``, ``omega``, or ``c`` is
                non-positive.
        """
        _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, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        """Apply the sine-activated linear transform.

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

        Returns:
            ``sin(ω · (x W + b))`` for ``layer_type`` in ``{"first", "hidden"}``,
            or ``x W + b`` for ``layer_type == "last"``.
        """
        pre = x @ self.W + self.b
        if self.layer_type == "last":
            return pre
        return jnp.sin(self.omega * pre)

__call__(x)

Apply the sine-activated linear transform.

Parameters:

Name	Type	Description	Default
x	Float[Array, '*batch D_in']	Input tensor of shape (*batch, in_features).	required

Returns:

Type	Description
Float[Array, '*batch D_out']	sin(ω · (x W + b)) for layer_type in {"first", "hidden"},
Float[Array, '*batch D_out']	or x W + b for layer_type == "last".

Source code in src/pyrox/nn/_layers.py

def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
    """Apply the sine-activated linear transform.

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

    Returns:
        ``sin(ω · (x W + b))`` for ``layer_type`` in ``{"first", "hidden"}``,
        or ``x W + b`` for ``layer_type == "last"``.
    """
    pre = x @ self.W + self.b
    if self.layer_type == "last":
        return pre
    return jnp.sin(self.omega * pre)

init(in_features, out_features, *, key, omega=30.0, layer_type='hidden', c=6.0) classmethod

Construct a SirenDense with Sitzmann-regime weight initialisation.

Parameters:

Name	Type	Description	Default
in_features	int	Input dimension.	required
out_features	int	Output dimension.	required
key	Array	JAX PRNG key for weight and bias sampling.	required
omega	float	Frequency multiplier (default 30.0, as in Sitzmann et al.).	30.0
layer_type	_SirenLayerType	Init regime — "first", "hidden", or "last".	'hidden'
c	float	Theorem-1 constant (default 6.0).	6.0

Returns:

Name	Type	Description
Initialised	SirenDense	class:SirenDense.

Raises:

Type	Description
ValueError	If layer_type is not a valid regime name, or if any of in_features, out_features, omega, or c is non-positive.

Source code in src/pyrox/nn/_layers.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.

    Args:
        in_features: Input dimension.
        out_features: Output dimension.
        key: JAX PRNG key for weight and bias sampling.
        omega: Frequency multiplier (default 30.0, as in Sitzmann et al.).
        layer_type: Init regime — ``"first"``, ``"hidden"``, or ``"last"``.
        c: Theorem-1 constant (default 6.0).

    Returns:
        Initialised :class:`SirenDense`.

    Raises:
        ValueError: If ``layer_type`` is not a valid regime name, or if any of
            ``in_features``, ``out_features``, ``omega``, or ``c`` is
            non-positive.
    """
    _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,
    )

pyrox.nn.SIREN

Bases: Module

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

Topology:

.. math::

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.

Each layer uses the corresponding Sitzmann Theorem 1 init regime (:class: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.

Attributes:

Name	Type	Description
layers	list[SirenDense]	List of :class:SirenDense of length depth.
in_features	int	Input dimension.
hidden_features	int	Hidden dimension (all intermediate layers).
out_features	int	Output dimension.
depth	int	Total number of layers (including readout). Must be ≥ 2.
first_omega	float	Frequency multiplier for the first layer (default 30.0).
hidden_omega	float	Frequency multiplier for hidden layers (default 30.0).

Example

import jax.random as jr, jax.numpy as jnp net = SIREN.init(2, 64, 1, depth=5, key=jr.PRNGKey(0)) net(jnp.zeros((10, 2))).shape (10, 1)

Source code in src/pyrox/nn/_layers.py

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

    Topology:

    .. math::

        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.

    Each layer uses the corresponding Sitzmann Theorem 1 init regime
    (:class:`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.

    Attributes:
        layers: List of :class:`SirenDense` of length ``depth``.
        in_features: Input dimension.
        hidden_features: Hidden dimension (all intermediate layers).
        out_features: Output dimension.
        depth: Total number of layers (including readout).  Must be ≥ 2.
        first_omega: Frequency multiplier for the first layer (default 30.0).
        hidden_omega: Frequency multiplier for hidden layers (default 30.0).

    Example:
        >>> import jax.random as jr, jax.numpy as jnp
        >>> net = SIREN.init(2, 64, 1, depth=5, key=jr.PRNGKey(0))
        >>> net(jnp.zeros((10, 2))).shape
        (10, 1)
    """

    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.

        Args:
            in_features: Input dimension.
            hidden_features: Hidden dimension for all intermediate layers.
            out_features: Output dimension.
            depth: Total layers including readout.  Must be ≥ 2.
            key: JAX PRNG key.
            first_omega: Frequency for the first layer (default 30.0).
            hidden_omega: Frequency for hidden layers (default 30.0).
            c: Theorem-1 constant passed to each :class:`SirenDense`.

        Returns:
            Initialised :class:`SIREN`.

        Raises:
            ValueError: If ``depth < 2`` or any of the feature dimensions,
                omegas, or ``c`` 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,
        )
        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, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
        """Run the forward pass through all SIREN layers.

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

        Returns:
            Output tensor of shape ``(*batch, out_features)``.
        """
        z = x
        for layer in self.layers:
            z = layer(z)
        return z

__call__(x)

Run the forward pass through all SIREN layers.

Parameters:

Name	Type	Description	Default
x	Float[Array, '*batch D_in']	Input tensor of shape (*batch, in_features).	required

Returns:

Type	Description
Float[Array, '*batch D_out']	Output tensor of shape (*batch, out_features).

Source code in src/pyrox/nn/_layers.py

def __call__(self, x: Float[Array, "*batch D_in"]) -> Float[Array, "*batch D_out"]:
    """Run the forward pass through all SIREN layers.

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

    Returns:
        Output tensor of shape ``(*batch, out_features)``.
    """
    z = x
    for layer in self.layers:
        z = layer(z)
    return z

init(in_features, hidden_features, out_features, *, depth, key, first_omega=30.0, hidden_omega=30.0, c=6.0) classmethod

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

Parameters:

Name	Type	Description	Default
in_features	int	Input dimension.	required
hidden_features	int	Hidden dimension for all intermediate layers.	required
out_features	int	Output dimension.	required
depth	int	Total layers including readout. Must be ≥ 2.	required
key	Array	JAX PRNG key.	required
first_omega	float	Frequency for the first layer (default 30.0).	30.0
hidden_omega	float	Frequency for hidden layers (default 30.0).	30.0
c	float	Theorem-1 constant passed to each :class:SirenDense.	6.0

Returns:

Name	Type	Description
Initialised	SIREN	class:SIREN.

Raises:

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

Source code in src/pyrox/nn/_layers.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.

    Args:
        in_features: Input dimension.
        hidden_features: Hidden dimension for all intermediate layers.
        out_features: Output dimension.
        depth: Total layers including readout.  Must be ≥ 2.
        key: JAX PRNG key.
        first_omega: Frequency for the first layer (default 30.0).
        hidden_omega: Frequency for hidden layers (default 30.0).
        c: Theorem-1 constant passed to each :class:`SirenDense`.

    Returns:
        Initialised :class:`SIREN`.

    Raises:
        ValueError: If ``depth < 2`` or any of the feature dimensions,
            omegas, or ``c`` 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,
    )
    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,
    )

pyrox.nn.BayesianSIREN

Bases: PyroxModule

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

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

.. math::

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 :math:\sigma_0 is prior_std and :math:d_i is the input dimension of layer :math:i. The :math:a_i / \sqrt{3} factor makes :math:\operatorname{Var}(W_i) equal to the variance of Sitzmann's :math:\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 :math:\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.

Example

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 src/pyrox/nn/_layers.py

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

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

    .. math::

        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 :math:`\sigma_0` is ``prior_std`` and :math:`d_i` is the input
    dimension of layer :math:`i`.  The :math:`a_i / \sqrt{3}` factor makes
    :math:`\operatorname{Var}(W_i)` equal to the variance of Sitzmann's
    :math:`\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 :math:`\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.

    Example:
        >>> 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 :class:`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 :class:`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 = z @ W + b
            z = pre if spec.layer_type == "last" else jnp.sin(spec.omega * pre)
        return z

__call__(x)

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).

Parameters:

Name	Type	Description	Default
x	Float[Array, '*batch D_in']	Input tensor of shape (*batch, in_features).	required

Returns:

Type	Description
Float[Array, '*batch D_out']	Output tensor of shape (*batch, out_features).

Source code in src/pyrox/nn/_layers.py

@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 = z @ W + b
        z = pre if spec.layer_type == "last" else jnp.sin(spec.omega * pre)
    return z

init(in_features, hidden_features, out_features, *, depth, first_omega=30.0, hidden_omega=30.0, c=6.0, prior_std=1.0, pyrox_name=None) classmethod

Construct a :class: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:

Name	Type	Description
Initialised	BayesianSIREN	class: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 src/pyrox/nn/_layers.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 :class:`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 :class:`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,
    )

Bayesian Neural Field stack

pyrox.nn.Standardization

Bases: PyroxModule

Apply a fixed-coefficient affine standardization.

.. math::

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

Both mu and std are static (fit-time) constants, not learned. Use :func:pyrox.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 src/pyrox/nn/_bnf.py

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

    .. math::

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

    Both ``mu`` and ``std`` are static (fit-time) constants, not
    learned. Use :func:`pyrox.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

pyrox.nn.FourierFeatures

Bases: PyroxModule

Per-input dyadic-frequency Fourier basis.

For each input column, evaluates 2 * degree Fourier features at frequencies :math:2\pi \cdot 2^d for :math:d \in \{0, \dots, \text{degree} - 1\}. Concatenated across all columns.

Wraps :func: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 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 :math:`2\pi \cdot 2^d` for :math:`d \in \{0, \dots,
    \text{degree} - 1\}`. Concatenated across all columns.

    Wraps :func:`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)

pyrox.nn.SeasonalFeatures

Bases: PyroxModule

Period-and-harmonic cos/sin basis on a scalar time axis.

For each period :math:\tau_p with :math:H_p harmonics, emits 2 * H_p cos/sin columns. Total output width is :math:2 \sum_p H_p.

Wraps :func: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 src/pyrox/nn/_bnf.py

class SeasonalFeatures(PyroxModule):
    r"""Period-and-harmonic cos/sin basis on a scalar time axis.

    For each period :math:`\tau_p` with :math:`H_p` harmonics, emits
    ``2 * H_p`` cos/sin columns. Total output width is :math:`2 \sum_p
    H_p`.

    Wraps :func:`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)

pyrox.nn.InteractionFeatures

Bases: PyroxModule

Element-wise products on selected pairs of input columns.

Wraps :func: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 src/pyrox/nn/_bnf.py

class InteractionFeatures(PyroxModule):
    r"""Element-wise products on selected pairs of input columns.

    Wraps :func:`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))

pyrox.nn.BayesianNeuralField

Bases: PyroxModule

The full Bayesian Neural Field architecture.

A spatiotemporal MLP with:

1. A learned per-input log-scale adjustment (Logistic(0, 1) prior).
2. Four feature blocks concatenated into h_0: rescaled inputs, Fourier features, seasonal features, interaction products.
3. Per-block softplus(feature_gain) modulation.
4. A depth-L MLP whose layers are :math:h_{\ell+1} = \sigma_\alpha\bigl(g_\ell \cdot W_\ell\, h_\ell / \sqrt{\lvert h_\ell \rvert}\bigr), where :math:\sigma_\alpha = \mathrm{sig}(\beta) \cdot \mathrm{elu} + (1 - \mathrm{sig}(\beta)) \cdot \mathrm{tanh} is a learned mixed activation.
5. A final linear layer scaled by softplus(output_gain).

All weights, biases, gains, scales, and the activation logit carry independent :math:\mathrm{Logistic}(0, 1) priors registered via :meth:PyroxModule.pyrox_sample.

The :math:1/\sqrt{\text{fan-in}} pre-normalization is the standard NTK-scaling trick — it makes the layer-wise prior predictive a fan-in-independent Gaussian process in the infinite-width limit (Lee et al., 2018).

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.

Source code in src/pyrox/nn/_bnf.py

class BayesianNeuralField(PyroxModule):
    r"""The full Bayesian Neural Field architecture.

    A spatiotemporal MLP with:

    1. A learned per-input log-scale adjustment (Logistic(0, 1) prior).
    2. Four feature blocks concatenated into ``h_0``: rescaled inputs,
       Fourier features, seasonal features, interaction products.
    3. Per-block ``softplus(feature_gain)`` modulation.
    4. A depth-``L`` MLP whose layers are
       :math:`h_{\ell+1} = \sigma_\alpha\bigl(g_\ell \cdot W_\ell\, h_\ell
       / \sqrt{\lvert h_\ell \rvert}\bigr)`, where :math:`\sigma_\alpha
       = \mathrm{sig}(\beta) \cdot \mathrm{elu} + (1 - \mathrm{sig}(\beta))
       \cdot \mathrm{tanh}` is a learned mixed activation.
    5. A final linear layer scaled by ``softplus(output_gain)``.

    All weights, biases, gains, scales, and the activation logit carry
    independent :math:`\mathrm{Logistic}(0, 1)` priors registered via
    :meth:`PyroxModule.pyrox_sample`.

    The :math:`1/\sqrt{\text{fan-in}}` pre-normalization is the
    standard NTK-scaling trick — it makes the layer-wise prior
    predictive a fan-in-independent Gaussian process in the
    infinite-width limit (Lee et al., 2018).

    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.
    """

    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, default=0)
    pyrox_name: str | None = None

    @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, "N D"]) -> Float[Array, " N"]:
        d_in = len(self.input_scales)
        input_scales = jnp.asarray(self.input_scales, dtype=jnp.float32)

        # 1. Input rescaling: x / (input_scales * exp(log_scale_adjustment)).
        log_scale_adjustment = self.pyrox_sample(
            "log_scale_adjustment",
            self._logistic_prior((d_in,)),
        )
        scaled_x = x / (input_scales * jnp.exp(log_scale_adjustment))

        # 2. Build the four feature blocks.
        feature_blocks: list[Float[Array, "N F_block"]] = [scaled_x]

        # Fourier per input dim (only for degrees > 0).
        for col_idx, d in enumerate(self.fourier_degrees):
            if d > 0:
                feature_blocks.append(
                    fourier_features(scaled_x[:, col_idx], d, rescale=True)
                )

        # Seasonal on the time column.
        if self.seasonality_periods and any(self.num_seasonal_harmonics):
            feature_blocks.append(
                seasonal_features(
                    x[:, self.time_col],
                    self.seasonality_periods,
                    self.num_seasonal_harmonics,
                    rescale=True,
                )
            )

        # Interaction products.
        if self.interactions:
            feature_blocks.append(
                interaction_features(
                    scaled_x, jnp.asarray(self.interactions, dtype=jnp.int32)
                )
            )

        # 3. Per-block softplus(feature_gain) modulation.
        gated_blocks: list[Float[Array, "N F_block"]] = []
        for b_idx, block in enumerate(feature_blocks):
            if block.shape[-1] == 0:
                continue
            gain = self.pyrox_sample(
                f"feature_gain_{b_idx}",
                self._logistic_prior(()),
            )
            gated_blocks.append(block * jax.nn.softplus(gain))
        h = jnp.concatenate(gated_blocks, axis=-1)

        # 4. Mixed elu/tanh activation, learned mix weight.
        logit_activation_weight = self.pyrox_sample(
            "logit_activation_weight",
            self._logistic_prior(()),
        )
        alpha = jax.nn.sigmoid(logit_activation_weight)

        def activation(z: Float[Array, ...]) -> Float[Array, ...]:
            return alpha * jax.nn.elu(z) + (1.0 - alpha) * jnp.tanh(z)

        # 5. Hidden MLP layers.
        for layer_idx in range(self.depth):
            fan_in = h.shape[-1]
            W = self.pyrox_sample(
                f"layer_{layer_idx}_W",
                self._logistic_prior((fan_in, self.width)),
            )
            b = self.pyrox_sample(
                f"layer_{layer_idx}_b",
                self._logistic_prior((self.width,)),
            )
            layer_gain = self.pyrox_sample(
                f"layer_{layer_idx}_gain",
                self._logistic_prior(()),
            )
            h = h / jnp.sqrt(fan_in)
            h = activation(jax.nn.softplus(layer_gain) * (h @ W + b))

        # 6. Output linear, scaled by softplus(output_gain).
        fan_in = h.shape[-1]
        W_out = self.pyrox_sample(
            "output_W",
            self._logistic_prior((fan_in, 1)),
        )
        b_out = self.pyrox_sample(
            "output_b",
            self._logistic_prior((1,)),
        )
        output_gain = self.pyrox_sample(
            "output_gain",
            self._logistic_prior(()),
        )
        h = h / jnp.sqrt(fan_in)
        return (jax.nn.softplus(output_gain) * (h @ W_out + b_out)).squeeze(-1)

Pure-JAX feature helpers

pyrox.nn.fourier_features(x, max_degree, *, rescale=False)

Cos/sin Fourier basis at dyadic frequencies.

For each input element and each degree :math:d \in \{0, \dots, D-1\}, evaluates

.. math::

\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).

Source code in src/pyrox/nn/_features.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 :math:`d \in \{0, \dots,
    D-1\}`, evaluates

    .. math::

        \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)``.
    """
    degrees = jnp.arange(max_degree)
    # `repeat` builds (N, D) frequencies without an explicit reshape.
    z = repeat(x, "n -> n d", 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

pyrox.nn.seasonal_features(x, periods, harmonics, *, rescale=False)

Cos/sin features at multiples of :math:2\pi / \tau_p.

For each period :math:\tau_p with :math:H_p harmonics, evaluates

.. math::

\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 :math:h = 1, \dots, H_p. Returns the cos columns concatenated with the sin columns, length :math: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).

Source code in src/pyrox/nn/_features.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 :math:`2\pi / \tau_p`.

    For each period :math:`\tau_p` with :math:`H_p` harmonics, evaluates

    .. math::

        \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 :math:`h = 1, \dots, H_p`. Returns the cos columns concatenated
    with the sin columns, length :math:`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)``.
    """
    _, freq_list = seasonal_frequencies(periods, harmonics)
    if not freq_list:
        return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
    freqs = jnp.asarray(freq_list, dtype=jnp.float32)
    z = repeat(x, "n -> n f", 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=jnp.float32)
        denom = jnp.concatenate([h_within, h_within])
        feats = feats / denom
    return feats

pyrox.nn.seasonal_frequencies(periods, harmonics)

Flatten (period, harmonic_count) pairs into Python lists.

For each period :math:\tau_p with :math:H_p harmonics, emits frequencies :math:f_{p, h} = h / \tau_p for :math:h = 1, \dots, H_p. The total length is :math: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 :func: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]	math:F = \sum_p H_p.

Source code in src/pyrox/nn/_features.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 :math:`\tau_p` with :math:`H_p` harmonics, emits
    frequencies :math:`f_{p, h} = h / \tau_p` for :math:`h = 1, \dots,
    H_p`. The total length is :math:`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 :func:`seasonal_features`.

    Args:
        periods: Period values.
        harmonics: Number of harmonics per period.

    Returns:
        ``(period_index, frequency)``: two Python lists of length
        :math:`F = \sum_p H_p`.
    """
    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

pyrox.nn.interaction_features(x, pairs)

Element-wise products on selected pairs of input columns.

For each pair :math:(i, j) and each row :math:n, computes :math: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.

Source code in src/pyrox/nn/_features.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 :math:`(i, j)` and each row :math:`n`, computes
    :math:`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.
    """
    if pairs.shape[0] == 0:
        return jnp.zeros((x.shape[0], 0), dtype=x.dtype)
    # x[:, pairs] has shape (N, K, 2); reduce the last axis with prod.
    selected = x[:, pairs]
    return jnp.prod(selected, axis=-1)

pyrox.nn.standardize(x, mu, std)

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.

Source code in src/pyrox/nn/_features.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.
    """
    return (x - mu) / std

pyrox.nn.unstandardize(z, mu, std)

Inverse of :func:standardize: z * std + mu.

Source code in src/pyrox/nn/_features.py

def unstandardize(
    z: Float[Array, "*shape"],
    mu: Float[Array, "*shape"],
    std: Float[Array, "*shape"],
) -> Float[Array, "*shape"]:
    """Inverse of :func:`standardize`: ``z * std + mu``."""
    return z * std + mu