The missing layers package for Bayesian regression.
BLayers is in beta, errors are possible! We invite you to contribute on GitHub.
pip install blayers
deps are: numpyro, jax, and optax.
Easily build Bayesian models from parts, abstract away the boilerplate, and tweak priors as you wish.
Inspiration from Keras and Tensorflow Probability, but made specifically for Numpyro + Jax.
Scope. BLayers works best for structured Bayesian regression — GLMs, hierarchical /
mixed-effects models, factorization machines, splines, and sparse priors. Layers
are meant to be added together into a linear predictor (mu = layer1(...) + layer2(...) + ...). You can stack them into a deep net, but better tools exist for this:
numpyro.contrib.module's
random_flax_module / random_haiku_module instead — they drop a full Flax or
Haiku net into a NumPyro model with priors on the weights.
BLayers provides tools to
- Quickly build Bayesian models from layers which encapsulate useful model parts
- Fit models either using Variational Inference (VI) or your sampling method of choice without having to rewrite models
- Write pure Numpyro to integrate with all of Numpyro's super powerful tools
- Add more complex layers (model parts) as you wish
- Fit models in a greater variety of ways with less code
The simplest non-trivial (and most important!) Bayesian regression model form is
the adaptive prior (note I'm using a . to denote the regression sigma and any
prior on it just to save space),
scale ~ HalfNormal(1)
beta ~ Normal(0, scale)
y ~ Normal(beta * x, .)
BLayers encapsulates a generative model structure like this in a BLayer. The
fundamental building block is the AdaptiveLayer.
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
def model(x, y):
mu = AdaptiveLayer()('mu', x)
return gaussian_link(mu, y)All AdaptiveLayer is doing is writing Numpyro for you under the hood. This
model is exactly equivalent to writing the following, just using way less code.
import jax.numpy as jnp
from numpyro import distributions, sample
def model(x, y):
# Adaptive layer does all of this
input_shape = x.shape[1]
# adaptive prior
scale = sample(
name="scale",
fn=distributions.HalfNormal(1.),
)
# beta coefficients for regression
beta = sample(
name="beta",
fn=distributions.Normal(loc=0., scale=scale),
sample_shape=(input_shape,),
)
mu = jnp.einsum('ij,j->i', x, beta)
# the link function does this
sigma = sample(name='sigma', fn=distributions.Exponential(1.))
return sample('obs', distributions.Normal(mu, sigma), obs=y)The AdaptiveLayer is also fully parameterizable via arguments to the class, so let's say you wanted to change the model from
scale ~ HalfNormal(1)
beta ~ Normal(0, scale)
y ~ Normal(beta * x, .)
to
scale ~ Exponential(1.)
beta ~ LogNormal(0, scale)
y ~ Normal(beta * x, .)
you can just do this directly via arguments
from numpyro import distributions
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
def model(x, y):
mu = AdaptiveLayer(
scale_dist=distributions.Exponential,
coef_dist=distributions.LogNormal,
scale_kwargs={'rate': 1.},
coef_kwargs={'loc': 0.}
)('mu', x)
return gaussian_link(mu, y)Since Numpyro traces sample sites and doesn't record any parameters on the class, you can re-use with a particular generative model structure freely.
from numpyro import distributions
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
my_lognormal_layer = AdaptiveLayer(
scale_dist=distributions.Exponential,
coef_dist=distributions.LogNormal,
scale_kwargs={'rate': 1.},
coef_kwargs={'loc': 0.}
)
def model(x, y):
mu = my_lognormal_layer('mu1', x) + my_lognormal_layer('mu2', x**2)
return gaussian_link(mu, y)The full set of layers included with BLayers:
AdaptiveLayer— Adaptive prior layer:scale ~ HalfNormal(1),beta ~ Normal(0, scale).FixedPriorLayer— Fixed prior over coefficients (e.g., Normal or Laplace), no hierarchical scale.InterceptLayer— Intercept-only layer (bias term).EmbeddingLayer— Bayesian embeddings for sparse categorical features.RandomEffectsLayer— Classical random-effects (embedding with output dim 1); learned variance component → partial pooling.FixedEffectsLayer— Per-category coefficients with a fixed, user-specified prior; the no-pooling counterpart ofRandomEffectsLayer.FMLayer— Factorization Machine (order 2) for pairwise interaction terms.FM3Layer— Factorization Machine (order 3).LowRankInteractionLayer— Low-rank interaction between two feature sets.InteractionLayer— All pairwise interactions between two feature sets.BilinearLayer— Bilinear interaction:x^T W z.LowRankBilinearLayer— Low-rank bilinear interaction.RandomWalkLayer— Gaussian random walk prior over an ordered index (e.g., time).HorseshoeLayer— Horseshoe prior for sparse regression; global-local shrinkage via HalfCauchy.SpikeAndSlabLayer— Spike-and-slab prior;z ~ Beta(0.5, 0.5)inclusion weights times a configurable slab.MixtureLayer— Finite mixture-of-priors on coefficients (default Normal + Laplace) with aDirichlet(or fixed) weight; the component indicator is marginalised so it works under VI and MCMC. Good for robustness / elastic-net-style priors.HSGPLayer— Hilbert-space approximate Gaussian process (1-D, squared-exponential; Riutort-Mayol et al. 2021). A GP smoother that learns its own lengthscale; usehsgp_L(x_train)to pick the domain boundary.
All layer prior kwargs are validated at construction time — bad kwargs raise TypeError immediately.
We provide link helpers in links.py to reduce Numpyro boilerplate. Available links:
gaussian_link— Gaussian likelihood with configurable sigma prior (see below).lognormal_link— LogNormal likelihood with configurable sigma prior.student_t_link— StudentT likelihood for robust regression (defaultdf=4).gamma_link— Gamma likelihood (log link) for positive continuous data; learned shape.exponential_link— Exponential likelihood (log link) for positive / survival data.logit_link— Bernoulli link for binary logistic regression.categorical_link— Categorical / softmax link for multiclass classification (units = num_classes).poisson_link— Poisson link with log-rate input.negative_binomial_link— NegativeBinomial2 for overdispersed counts; learned concentration viaExponential.ordinal_link— Cumulative logit / proportional odds for ordinal outcomes.zip_link— Zero-inflated Poisson for count data with excess zeros.zinb_link— Zero-inflated NegativeBinomial2 for overdispersed, zero-heavy counts.beta_link— Beta regression for proportions strictly in (0, 1).
All three share a common location-scale base and support three scale modes:
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
# Default: sigma ~ Exp(1) learned from data
gaussian_link(mu, y)
# Fixed known scale (e.g. from XGBoost quantile regression)
gaussian_link(mu, y, scale=pred_std)
# Learned scale from a layer — softplus applied internally for stable gradients
raw = AdaptiveLayer()("log_sigma", x)
gaussian_link(mu, y, untransformed_scale=raw)Swap the sigma prior via functools.partial:
from functools import partial
import numpyro.distributions as dists
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
# HalfNormal prior instead of Exponential
hn_gaussian = partial(gaussian_link, sigma_dist=dists.HalfNormal, sigma_kwargs={"scale": 1.0})
def model(x, y=None):
mu = AdaptiveLayer()("mu", x)
return hn_gaussian(mu, y)Non-linear transformations via B-splines. Compute the basis matrix once with make_knots + bspline_basis, then pass it to any layer.
from blayers.splines import make_knots, bspline_basis
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
knots = make_knots(x_train, num_knots=10) # clamped knot vector from data quantiles
def model(x, y=None):
B = bspline_basis(x, knots) # (n, num_basis) design matrix
f = AdaptiveLayer()("f", B)
return gaussian_link(f, y)Additive models are straightforward:
knots1 = make_knots(x1_train, num_knots=10)
knots2 = make_knots(x2_train, num_knots=10)
def model(x1, x2, y=None):
f1 = AdaptiveLayer()("f1", bspline_basis(x1, knots1))
f2 = AdaptiveLayer()("f2", bspline_basis(x2, knots2))
return gaussian_link(f1 + f2, y)HSGPLayer is a Hilbert-space approximate GP (Riutort-Mayol et al. 2021) — a smoother like splines, but it learns its own lengthscale and carries a proper GP interpretation. Pick the domain boundary L once on the training inputs with hsgp_L and reuse it at predict time; m is the number of basis functions (~20–50).
from blayers.layers import HSGPLayer, hsgp_L
from blayers.links import gaussian_link
from blayers.decorators import autoreshape
L = hsgp_L(x_train) # domain boundary = 1.5 * max(|x|); fixed across fit/predict
@autoreshape
def model(x, y=None, L=L, m=30):
f = HSGPLayer()("f", x, L=L, m=m)
return gaussian_link(f, y)Like splines, HSGP terms add for a GAM-style additive model (f1(x1) + f2(x2) + ...). Center/scale each input so it lies within [-L, L].
MixtureLayer draws each coefficient from a finite mixture of priors (default Normal + Laplace) — useful for robustness (a heavy-tailed component absorbs outlier coefficients) or elastic-net-flavoured priors. The mixing weights get a Dirichlet prior by default, or pass fixed weights=. The component indicator is marginalised internally, so it works under both VI and MCMC.
from blayers.layers import MixtureLayer
from blayers.links import gaussian_link
def model(x, y=None):
mu = MixtureLayer()("beta", x) # Normal + Laplace, Dirichlet weights
return gaussian_link(mu, y)For pure sparsity prefer HorseshoeLayer; for explicit variable selection prefer SpikeAndSlabLayer.
fit() handles the guide, ELBO, batching, and LR schedule. The same model runs unchanged under VI, MCMC, or SVGD.
from blayers.fit import fit
from blayers.decorators import autoreshape
from blayers.layers import AdaptiveLayer, InterceptLayer
from blayers.links import gaussian_link
@autoreshape
def model(x, y=None):
mu = AdaptiveLayer()("beta", x)
intercept = InterceptLayer()("intercept")
return gaussian_link(mu + intercept, y)
# Variational Inference (default)
result = fit(model, y=y, num_steps=1000, batch_size=256, lr=0.01, x=X)
# MCMC
result = fit(model, y=y, method="mcmc", num_mcmc_samples=1000, num_warmup=500, x=X)
# SVGD
result = fit(model, y=y, method="svgd", num_steps=1000, num_particles=20, x=X)result.predict() returns a Predictions object with .mean, .std, and .samples. result.summary() returns posterior stats per latent variable.
preds = result.predict(x=X, num_samples=500)
summary = result.summary(x=X)Keyword arguments that are JAX arrays are treated as data (batched during training). Non-array kwargs are bound as constants.
result.to_arviz() hands the fit to ArviZ for R-hat,
ESS, divergences, PSIS-LOO, and the full plotting suite — reusing NumPyro's own
ArviZ bridge rather than reinventing diagnostics. Install with pip install blayers[arviz] (arviz ≥ 1.0, Python ≥ 3.12).
import arviz as az
# MCMC: divergences, R-hat, ESS, and log-likelihood come through automatically
idata = fit(model, y=y, method="mcmc", num_chains=2, x=X).to_arviz()
az.summary(idata) # R-hat / ESS per latent
az.loo(idata) # PSIS-LOO
# VI: pass the observed y (and inputs) so the log_likelihood group can be built
idata_vi = fit(model, y=y, num_steps=2000, x=X).to_arviz(y=y, x=X)
# Compare models on out-of-sample predictive fit
az.compare({"mcmc": idata, "vi": idata_vi})SVGD is not supported by to_arviz() (too few particles to be a meaningful
sample for LOO); fit with method="mcmc" or method="vi" for comparison.
The default Numpyro way to fit batched VI models is to use plate, which confuses
me a lot. Instead, BLayers provides Batched_Trace_ELBO which does not require
you to use plate to batch in VI. Just drop your model in.
from numpyro.infer import SVI
from numpyro.infer.autoguide import AutoDiagonalNormal
import optax
from blayers.vi_infer import Batched_Trace_ELBO, svi_run_batched
loss = Batched_Trace_ELBO(num_obs=len(y), batch_size=1000)
guide = AutoDiagonalNormal(model_fn)
svi = SVI(model_fn, guide, optax.adam(0.01), loss=loss)
svi_result = svi_run_batched(
svi,
rng_key,
batch_size=1000,
num_steps=500,
**model_data,
)numpyro.plate + Batched_Trace_ELBO do not mix.
Batched_Trace_ELBO does not support numpyro.plate: its N / batch_size log-likelihood rescaling double-counts plate-subsampled sites and yields an incorrect ELBO. If your model needs plates, either:
- Batch via
plateand use the standardTrace_ELBO, or - Remove plates and use
Batched_Trace_ELBO+svi_run_batched.
Batched_Trace_ELBO raises ValueError if your model contains a plate.
To fit MCMC models well it is crucial to reparameterize. BLayers helps you do this via @autoreparam, which automatically applies LocScaleReparam to all LocScale distributions in your model (Normal, LogNormal, StudentT, Cauchy, Laplace, Gumbel).
Note:
fit(method="mcmc")already applies@autoreparamfor you (controlled byautoreparam_model=True, on by default). You only need to apply the decorator yourself when driving NUTS / HMC manually, as shown below.
from numpyro.infer import MCMC, NUTS
from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
from blayers.decorators import autoreparam
data = {...}
@autoreparam
def model(x, y):
mu = AdaptiveLayer()('mu', x)
return gaussian_link(mu, y)
kernel = NUTS(model)
mcmc = MCMC(
kernel,
num_warmup=500,
num_samples=1000,
num_chains=1,
progress_bar=True,
)
mcmc.run(
rng_key,
**data,
)