Layers

Layers#

Implements Bayesian Layers using Jax and Numpyro.

Design:

There are three levels of complexity here: class-level, instance-level, and call-level
The class-level handles things like choosing generic model form and how to multiply coefficents with data. Defined by the class Layer(BLayer) def itself.
The instance-level handles specific distributions that fit into a generic model and the initial parameters for those distributions. Defined by creating an instance of the class: Layer(*args, **kwargs).
The call-level handles seeing a batch of data, sampling from the distributions defined on the class and multiplying coefficients and data to produce an output, works like result = Layer(*args, **kwargs)(data)

Notation:

n: observations in a batch
c: number of categories of things for time, random effects, etc
d: number of coefficients
l: low rank dimension of low rank models
m: embedding dimension
u: units aka output dimension

blayers.layers.pairwise_interactions(x, z)[source]#

Compute all pairwise interactions between features in x and z.

Parameters:

x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).

Returns:

jax.Array of shape (n, d1 * d2) containing the flattened outer product x[:, i] * z[:, j] for each pair (i, j).

Return type:

Array

class blayers.layers.BLayer(*args)[source]#

Bases: ABC

Abstract base class for Bayesian layers. Lays out an interface.

Parameters:: args (Any)

abstractmethod __init__(*args)[source]#

Initialize layer parameters. This is the Bayesian model.

Parameters:: args (Any)
Return type:: None

abstractmethod __call__(*args)[source]#

Run the layer’s forward pass.

Parameters:: *args (Any) – Inputs to the layer.
Returns:: The result of the forward computation.
Return type:: jax.Array

class blayers.layers.AdaptiveLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian layer with adaptive prior using hierarchical modeling.

Generates coefficients from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\beta \sim Normal(0., \lambda)\]

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

__init__(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Parameters:

scale_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
scale_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with adaptive prior on coefficients.

Parameters:

name (str) – Variable name.
x (Array) – Input data array of shape (n, d).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.FixedPriorLayer(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Bayesian layer with a fixed prior distribution over coefficients.

Generates coefficients from the model

\[\beta \sim Normal(0., 1.)\]

Parameters:

coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Parameters:

coef_dist (Distribution) – NumPyro distribution class for the coefficients.
coef_kwargs (dict[str, float]) – Parameters to initialize the prior distribution.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with fixed prior.

Parameters:

name (str) – Variable name.
x (Array) – Input data array of shape (n, d).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.InterceptLayer(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Bayesian intercept (bias) term with a fixed prior.

Samples a scalar bias from

\[\beta \sim Normal(0., 1.)\]

and broadcasts it to every observation. No input x is needed.

Parameters:

coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Parameters:

coef_dist (Distribution) – NumPyro distribution class for the coefficients.
coef_kwargs (dict[str, float]) – Parameters to initialize the prior distribution.

__call__(name, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with fixed prior.

Parameters:

name (str) – Variable name.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (1, u).

Return type:

jax.Array

class blayers.layers.FMLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian factorization machine layer with adaptive priors.

Generates coefficients from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\beta \sim Normal(0., \lambda)\]

The shape of beta is (j, l), where j is the number if input covariates and l is the low rank dim.

Then performs matrix multiplication using the formula in Rendle (2010).

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – Distribution for scaling factor λ.
coef_dist (Distribution) – Prior for beta parameters.
coef_kwargs (dict[str, float]) – Arguments for prior distribution.
scale_kwargs (dict[str, float]) – Arguments for λ distribution.

__call__(name, x, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass through the factorization machine layer.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.FM3Layer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian order-3 factorization machine layer with adaptive prior.

Samples low-rank factors from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\theta \sim Normal(0., \lambda), \quad \theta \in \mathbb{R}^{d \times l}\]

Then computes the 3rd-order ANOVA kernel via Newton’s identity (Blondel et al. 2016). Defining power sums \(p_k = \sum_i x_i^k \theta_i^k\):

\[\text{output} = \frac{p_1^3 - 3\, p_2\, p_1 + 2\, p_3}{6}\]

This efficiently computes all 3rd-order interaction terms without enumerating all \(\binom{d}{3}\) triples.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – Distribution for scaling factor λ.
coef_dist (Distribution) – Prior for beta parameters.
coef_kwargs (dict[str, float]) – Arguments for prior distribution.
scale_kwargs (dict[str, float]) – Arguments for λ distribution.

__call__(name, x, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass through the factorization machine layer.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.LowRankInteractionLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian low-rank bilinear interaction between two feature sets (UV decomposition).

Samples separate low-rank projections for x and z from the hierarchical model

\[\lambda_1 \sim HalfNormal(1.), \quad \theta_1 \sim Normal(0., \lambda_1), \quad \theta_1 \in \mathbb{R}^{d_1 \times l}\]

\[\lambda_2 \sim HalfNormal(1.), \quad \theta_2 \sim Normal(0., \lambda_2), \quad \theta_2 \in \mathbb{R}^{d_2 \times l}\]

and computes the element-wise product of the projections, summed over the low-rank dimension:

\[\text{output} = \sum_{r=1}^{l} (x \theta_1)_r \cdot (z \theta_2)_r = x^\top (\theta_1 \theta_2^\top) z\]

This is equivalent to a rank-\(l\) approximation of the full bilinear form \(x^\top W z\).

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior. Each input gets its own scale.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
scale_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, z, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Low-rank bilinear interaction x^T (theta1 theta2^T) z between X and Z.

Projects x and z into a shared low_rank_dim-dimensional space via independent factors theta1 and theta2, then contracts.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.InteractionLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian full pairwise interaction layer with adaptive prior.

Samples one coefficient per pair of features from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\beta \sim Normal(0., \lambda), \quad \beta \in \mathbb{R}^{d_1 d_2}\]

and computes the weighted sum of all outer-product interactions:

\[\text{output} = (x \otimes z)\, \beta\]

where \(x \otimes z\) is the flattened outer product of shape \((n, d_1 d_2)\). For large inputs this scales as \(O(d_1 d_2)\) parameters; prefer LowRankInteractionLayer when \(d_1\) or \(d_2\) is large.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Initialize layer parameters. This is the Bayesian model.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

__call__(name, x, z, units=1, activation=<PjitFunction of <function identity>>)[source]#

Full pairwise interaction between feature matrices X and Z.

Samples one coefficient per (x_i, z_j) pair (d1 * d2 total) and returns the weighted sum of all outer-product interactions.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.BilinearLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian full bilinear layer with adaptive prior.

Samples a full interaction matrix from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[W \sim Normal(0., \lambda), \quad W \in \mathbb{R}^{d_1 \times d_2}\]

and computes the bilinear form:

\[\text{output} = x^\top W z\]

This learns a distinct weight for every pair \((x_i, z_j)\), making it the densest two-input layer. Has \(O(d_1 d_2)\) parameters; prefer LowRankBilinearLayer when dimensions are large.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – prior on scale of coefficients
coef_dist (Distribution) – distribution for coefficients
coef_kwargs (dict[str, float]) – kwargs for coef distribution
scale_kwargs (dict[str, float]) – kwargs for scale prior

__call__(name, x, z, units=1, activation=<PjitFunction of <function identity>>)[source]#

Full bilinear form x^T W z between feature matrices X and Z.

Samples a dense weight tensor W of shape (d1, d2, units) and contracts it against x and z.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.LowRankBilinearLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian low-rank bilinear layer with adaptive prior.

Samples shared-scale low-rank factors for both inputs from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[A \sim Normal(0., \lambda), \quad A \in \mathbb{R}^{d_1 \times l}\]

\[B \sim Normal(0., \lambda), \quad B \in \mathbb{R}^{d_2 \times l}\]

and computes the bilinear form with a rank-\(l\) weight matrix \(W = AB^\top\):

\[\text{output} = x^\top W z = (xA) \cdot (zB)\]

Compared to LowRankInteractionLayer, A and B share a single scale \(\lambda\), tying the regularisation across both inputs.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – prior on scale of coefficients
coef_dist (Distribution) – distribution for coefficients
coef_kwargs (dict[str, float]) – kwargs for coef distribution
scale_kwargs (dict[str, float]) – kwargs for scale prior

__call__(name, x, z, low_rank_dim, units=1, activation=<PjitFunction of <function identity>>)[source]#

Low-rank bilinear form x^T (A B^T) z.

Projects x and z into a shared low_rank_dim-dimensional space via shared-scale factors A and B, then contracts.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array) – Input matrix of shape (n, d2).
low_rank_dim (int) – Dimensionality of low-rank approximation.
units (int) – Number of outputs.
activation (Callable[[Array], Array]) – Activation function to apply to output.

Returns:

Output array of shape (n, u).

Return type:

jax.Array

class blayers.layers.EmbeddingLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian embedding layer for sparse categorical features.

Samples an embedding table from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\theta \sim Normal(0., \lambda), \quad \theta \in \mathbb{R}^{c \times m}\]

and performs a lookup for each observation:

\[\text{output}_i = \theta[x_i]\]

where \(c\) is the number of categories and \(m\) is the embedding dimension. For \(m = 1\) prefer RandomEffectsLayer.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
scale_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, num_categories, embedding_dim)[source]#

Forward pass through embedding lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating embeddings to use.
num_categories (int) – The number of distinct things getting an embedding
embedding_dim (int) – The size of each embedding, e.g. 2, 4, 8, etc.

Returns:

Embedding vectors of shape (n, m).

Return type:

jax.Array

class blayers.layers.RandomEffectsLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian random-effects layer — a scalar embedding per category.

Special case of EmbeddingLayer with embedding_dim=1. Samples one scalar random effect per category from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\theta \sim Normal(0., \lambda), \quad \theta \in \mathbb{R}^{c}\]

and returns the scalar for each observation’s category:

\[\text{output}_i = \theta[x_i]\]

Equivalent to a classical mixed-effects intercept with a learned variance \(\lambda^2\).

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Parameters:

scale_dist (Distribution) – NumPyro distribution class for the scale (λ) of the prior.
coef_dist (Distribution) – NumPyro distribution class for the coefficient prior.
coef_kwargs (dict[str, float]) – Parameters for the prior distribution.
scale_kwargs (dict[str, float]) – Parameters for the scale distribution.

__call__(name, x, num_categories)[source]#

Forward pass through scalar random-effect lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating which random effect to use.
num_categories (int) – The number of distinct random-effect groups.

Returns:

Random-effect values of shape (n, 1).

Return type:

jax.Array

class blayers.layers.FixedEffectsLayer(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Bayesian fixed-effects layer — per-category coefficients, fixed prior.

The no-pooling counterpart of RandomEffectsLayer: each category gets its own scalar coefficient drawn from a fixed, user-specified prior

\[\theta_c \sim Normal(0., 1.), \quad \theta \in \mathbb{R}^{c}\]

with no learned variance component. The learned scale is exactly what makes a random effect “random” (it drives the partial pooling); fixing the prior instead gives the Bayesian analogue of classical fixed effects — ridge-regularised per-category dummies, with the prior scale controlling the regularisation strength.

Each observation gets the coefficient of its category:

\[\text{output}_i = \theta[x_i]\]

Prefer RandomEffectsLayer when you have many small groups that should share information; use this layer when groups are few and well-observed, or when you explicitly do not want pooling.

Parameters:

coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Parameters:

coef_dist (Distribution) – NumPyro distribution class for the coefficients.
coef_kwargs (dict[str, float]) – Parameters to initialize the prior distribution.

__call__(name, x, num_categories)[source]#

Forward pass through scalar fixed-effect lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating which effect to use.
num_categories (int) – The number of distinct groups.

Returns:

Effect values of shape (n, 1).

Return type:

jax.Array

class blayers.layers.RandomWalkLayer(scale_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0}, scale_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Bayesian Gaussian random walk over ordered categories.

Samples i.i.d. increments from the hierarchical model

\[\lambda \sim HalfNormal(1.)\]

\[\delta_t \sim Normal(0., \lambda), \quad t = 1, \ldots, c\]

and accumulates them into positions via a cumulative sum:

\[\theta_t = \sum_{s=1}^{t} \delta_s\]

Each observation is then assigned the position of its category:

\[\text{output}_i = \theta[x_i]\]

The embedding_dim m runs m independent walks in parallel, producing output of shape (n, m). Typical use: a time index where adjacent periods share information through the walk prior.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

Initialize layer parameters. This is the Bayesian model.

Parameters:

scale_dist (Distribution)
coef_dist (Distribution)
coef_kwargs (dict[str, float])
scale_kwargs (dict[str, float])

__call__(name, x, num_categories, embedding_dim)[source]#

Forward pass through embedding lookup.

Parameters:

name (str) – Variable name scope.
x (Array) – Integer indices indicating embeddings to use.
num_categories (int) – The number of distinct things getting an embedding
embedding_dim (int) – The size of each embedding, e.g. 2, 4, 8, etc.

Returns:

Embedding vectors of shape (n, m).

Return type:

jax.Array

class blayers.layers.HorseshoeLayer(slab_scale=None, slab_df=4.0, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0})[source]#

Bases: BLayer

Bayesian layer with horseshoe prior for sparse regression.

Implements the (regularized) horseshoe prior of Piironen & Vehtari (2017).

Basic horseshoe:

\[\tau \sim HalfCauchy(1), \quad \lambda_j \sim HalfCauchy(1), \quad \beta_j \sim Normal(0,\; \tau \lambda_j)\]

Regularized horseshoe (slab_scale set) — prevents large coefficients from escaping the slab:

\[\tilde{\lambda}_j^2 = \frac{c^2 \lambda_j^2}{c^2 + \tau^2 \lambda_j^2}, \quad c^2 \sim InverseGamma(s/2,\; s/2 \cdot scale_{slab}^2)\]

Parameters:

slab_scale (float | None)
slab_df (float)
coef_dist (Distribution)
coef_kwargs (dict[str, float])

__init__(slab_scale=None, slab_df=4.0, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0})[source]#

Parameters:

slab_scale (float | None) – If set, uses the regularized horseshoe with this slab scale. None gives the plain horseshoe.
slab_df (float) – Degrees of freedom for the slab variance prior (only used when slab_scale is set).
coef_dist (Distribution) – Distribution for the coefficients. Must accept a scale keyword (derived from the horseshoe shrinkage). Defaults to Normal.
coef_kwargs (dict[str, float]) – Extra kwargs for coef_dist (beyond scale). Default {"loc": 0.0}.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Forward pass with horseshoe prior on coefficients.

Parameters:

name (str) – Variable name scope.
x (Array) – Input array of shape (n, d).
units (int) – Number of output dimensions.
activation (Callable[[Array], Array]) – Activation function.

Returns:

jax.Array of shape (n, units).

Return type:

Array

class blayers.layers.SpikeAndSlabLayer(alpha=0.5, beta=0.5, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Bases: BLayer

Sparse regression via a spike-and-slab prior.

Each coefficient has a Beta-distributed inclusion weight z_j in (0, 1). Included features (z_j ≈ 1) take the full slab coefficient; excluded features (z_j ≈ 0) are gated toward zero (the spike).

Generative model:

z_j ~ Beta(alpha, beta)          # inclusion weight (hardcoded Beta)
β_j ~ coef_dist(**coef_kwargs)   # slab coefficient
y   ~ link(z · β · x, ...)       # z gates each coefficient

The default Beta(0.5, 0.5) (Jeffreys prior) places mass near 0 and 1, encouraging features to be clearly included or excluded. The posterior mean of z_j approximates P(feature j included | data).

The slab distribution defaults to Normal(0, 1) but can be swapped for e.g. StudentT for heavier-tailed slab behaviour.

Parameters:

alpha (float) – First concentration parameter of the Beta prior on z.
beta (float) – Second concentration parameter of the Beta prior on z.
coef_dist (Distribution) – Distribution for the slab coefficients.
coef_kwargs (dict[str, float]) – Kwargs for coef_dist.

__init__(alpha=0.5, beta=0.5, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0, 'scale': 1.0})[source]#

Initialize layer parameters. This is the Bayesian model.

Parameters:

alpha (float)
beta (float)
coef_dist (Distribution)
coef_kwargs (dict[str, float])

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Parameters:

name (str) – Variable name scope.
x (Array) – Input of shape (n, d).
units (int) – Number of output dimensions.
activation (Callable[[Array], Array]) – Activation function.

Returns:

jax.Array of shape (n, units).

Return type:

Array

class blayers.layers.MixtureLayer(component_dists=(<class 'numpyro.distributions.continuous.Normal'>, <class 'numpyro.distributions.continuous.Laplace'>), component_kwargs=({'loc': 0.0, 'scale': 1.0}, {'loc': 0.0, 'scale': 1.0}), weights=None, dirichlet_concentration=1.0)[source]#

Bases: BLayer

Coefficients from a finite mixture-of-priors (e.g. Normal + Laplace).

Each coefficient is drawn from a K-component mixture

\[\beta_j \sim \sum_{k=1}^{K} w_k \, p_k(\cdot)\]

where the mixing weights w are either fixed or given a Dirichlet prior (shared across coefficients). The component indicator is marginalised analytically by numpyro.distributions.MixtureGeneral, so the log-density is smooth and works under VI and MCMC — unlike a discrete spike-and-slab indicator.

Useful for robustness (a heavy-tailed component absorbs a few outlier coefficients while the rest stay Gaussian) and elastic-net-flavoured priors (Normal + Laplace). For pure sparsity prefer HorseshoeLayer; for explicit variable selection prefer SpikeAndSlabLayer.

Parameters:

component_dists (tuple[type[Distribution], ...])
component_kwargs (tuple[dict[str, float], ...])
weights (list[float] | None)
dirichlet_concentration (float)

__init__(component_dists=(<class 'numpyro.distributions.continuous.Normal'>, <class 'numpyro.distributions.continuous.Laplace'>), component_kwargs=({'loc': 0.0, 'scale': 1.0}, {'loc': 0.0, 'scale': 1.0}), weights=None, dirichlet_concentration=1.0)[source]#

Parameters:

component_dists (tuple[type[Distribution], ...]) – NumPyro distribution classes, one per mixture component (>= 2). All must share the same (real) support.
component_kwargs (tuple[dict[str, float], ...]) – Kwargs for each component distribution.
weights (list[float] | None) – Fixed mixing weights (one per component, summing to 1). If None, a Dirichlet prior is placed on the weights.
dirichlet_concentration (float) – Symmetric Dirichlet concentration used when weights is None.

__call__(name, x, units=1, activation=<PjitFunction of <function identity>>)[source]#

Parameters:

name (str) – Variable name scope.
x (Array) – Input of shape (n, d).
units (int) – Number of output dimensions.
activation (Callable[[Array], Array]) – Activation function.

Returns:

jax.Array of shape (n, units).

Return type:

Array

blayers.layers.hsgp_L(x, c=1.5)[source]#

Boundary L = c * max(|x|) for the Hilbert-space GP basis.

Compute this once on the training inputs and pass the same L to HSGPLayer at both fit and predict time — the eigenfunction basis is only valid on a fixed domain [-L, L]. c in ~[1.2, 2.0]; larger is safer near the data edges. Center x first so it straddles 0.

Parameters:

x (Any)
c (float)

Return type:

float

class blayers.layers.HSGPLayer(lengthscale_dist=<class 'numpyro.distributions.continuous.InverseGamma'>, lengthscale_kwargs={'concentration': 5.0, 'rate': 5.0}, sigma_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, sigma_kwargs={'scale': 1.0})[source]#

Bases: BLayer

Hilbert-space approximate Gaussian process (1-D, squared-exponential).

Low-rank GP of Riutort-Mayol et al. (2020): a stationary GP on [-L, L] is approximated with m Laplacian eigenfunctions, turning the GP into a basis-function layer

\[f(x) \approx \sum_{j=1}^{m} \phi_j(x)\, \sqrt{S(\sqrt{\lambda_j})}\, \beta_j, \quad \beta_j \sim \mathrm{Normal}(0, 1)\]

where \phi_j / \lambda_j are the eigenfunctions / eigenvalues on [-L, L] and S is the squared-exponential spectral density (a function of the sampled lengthscale ell and marginal std alpha). Sits alongside blayers.splines.bspline_basis() and RandomWalkLayer as a smoother, but learns its own lengthscale and carries a proper GP interpretation.

Center / scale x so it lies within [-L, L]; pick L with hsgp_L() on the training data and reuse it at predict time. m trades accuracy for cost (~20–50 is typical); the approximation degrades for lengthscales that are very short relative to the domain.

Parameters:

lengthscale_dist (Distribution)
lengthscale_kwargs (dict[str, float])
sigma_dist (Distribution)
sigma_kwargs (dict[str, float])

__init__(lengthscale_dist=<class 'numpyro.distributions.continuous.InverseGamma'>, lengthscale_kwargs={'concentration': 5.0, 'rate': 5.0}, sigma_dist=<class 'numpyro.distributions.continuous.HalfNormal'>, sigma_kwargs={'scale': 1.0})[source]#

Parameters:

lengthscale_dist (Distribution) – Prior distribution class for the GP lengthscale.
lengthscale_kwargs (dict[str, float]) – Kwargs for the lengthscale prior.
sigma_dist (Distribution) – Prior distribution class for the GP marginal std.
sigma_kwargs (dict[str, float]) – Kwargs for the marginal-std prior.

__call__(name, x, L, m, units=1, activation=<PjitFunction of <function identity>>)[source]#

Parameters:

name (str) – Variable name scope.
x (Array) – 1-D input of shape (n,) or (n, 1) within [-L, L].
L (float) – Domain boundary (see hsgp_L()). Fixed across fit/predict.
m (int) – Number of basis functions.
units (int) – Number of output dimensions.
activation (Callable[[Array], Array]) – Activation function.

Returns:

jax.Array of shape (n, units).

Return type:

Array

Layers

Contents

Layers#