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 pairwise interaction layer with adaptive prior.

Samples one coefficient per pair of features from the hierarchical model

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

and computes the weighted sum of the interaction design.

Two modes:

Within a single feature set (z omitted): the unique pairs \(x_i x_j\) for \(i < j\) — no squares, no duplicates — \(\binom{d}{2}\) coefficients, in lexicographic i < j order.
Between two feature sets (z given): the full flattened outer product \(x \otimes z\) of shape \((n, d_1 d_2)\).

Scales as \(O(d^2)\) parameters; prefer LowRankInteractionLayer when \(d\) is large, or HorseshoeInteractionLayer for a sparse (variable-selecting) prior over the pairs.

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=None, units=1, activation=<PjitFunction of <function identity>>)[source]#

Pairwise interaction design times a per-pair coefficient.

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array | None) – Optional second feature set of shape (n, d2). If omitted, the interactions are the unique within-x pairs i < j.
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(tau0=1.0, 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(\tau_0), \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:

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

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

Parameters:

tau0 (float) – Scale of the HalfCauchy prior on the GLOBAL shrinkage tau. This is the knob that decides whether the layer selects. At p >> n the default of 1.0 is very loose: tau drifts to O(10) and, under a regularized horseshoe, the per-coefficient scale tau * sqrt(c^2 l^2 / (c^2 + tau^2 l^2)) collapses to just c for any l >> c/tau — so the prior degenerates to Normal(0, slab_scale) and shrinks nothing. Piironen & Vehtari suggest tau0 ~ (p0 / (p - p0)) * sigma / sqrt(n) for an expected p0 non-zero coefficients; in practice something like 0.05 is a reasonable starting point for sparse selection.
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.HorseshoeInteractionLayer(tau0=1.0, slab_scale=None, slab_df=4.0, coef_dist=<class 'numpyro.distributions.continuous.Normal'>, coef_kwargs={'loc': 0.0})[source]#

Bases: HorseshoeLayer

Sparse pairwise interactions under a horseshoe prior.

Builds an explicit interaction design and places a (regularized) horseshoe prior on the per-pair coefficients. Local shrinkage pulls most interactions to zero and leaves the few real ones standing, so this is the layer to reach for to identify sparse interactions (as opposed to InteractionLayer’s single global scale, which cannot localize).

Two modes:

Within a single feature set (z omitted): the unique pairs x_i x_j for i < j — no squares, no duplicates — C(d, 2) columns. Column k is the k-th pair in lexicographic i < j order (row-major upper triangle), so posterior coefficients map back to feature pairs.
Between two feature sets (z given): the full d1 * d2 outer product x_i z_j (like InteractionLayer); column k is (i, j) = divmod(k, d2).

Costs O(d^2) coefficients — for large inputs where you only need prediction, prefer LowRankInteractionLayer or FMLayer. Inherits its prior configuration (slab_scale, slab_df, coef_dist, coef_kwargs) from HorseshoeLayer.

Parameters:

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

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

Parameters:

name (str) – Variable name scope.
x (Array) – Input matrix of shape (n, d1).
z (Array | None) – Optional second feature set of shape (n, d2). If omitted, the interactions are the unique within-x pairs i < j.
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, weight_scale=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)
weight_scale (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, weight_scale=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 logistic-normal prior is placed on the weights: softmax of Normal(0, weight_scale) logits. This keeps the weight latent in unconstrained space so the layer fits under VI, MCMC, and SVGD — a raw Dirichlet simplex site breaks SVGD’s particle flattening (its unconstrained dimension differs from its constrained one).
weight_scale (float) – Prior standard deviation of the Normal logits used when weights is None. Larger spreads the weights more.

__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#