Fit

Fit#

High-level fitting API for blayers models.

Reduces boilerplate when fitting Bayesian models by providing sensible defaults for guides, optimizers, learning rate schedules, and prediction.

Example

from blayers.layers import AdaptiveLayer
from blayers.links import gaussian_link
from blayers.fit import fit

def model(x, y=None):
    mu = AdaptiveLayer()('mu', x)
    return gaussian_link(mu, y)

# Fit with batched VI
result = fit(model, y=y_train, batch_size=1024, num_epochs=100, x=x_train)

# Predict on new data
preds = result.predict(x=x_test)
print(preds.mean, preds.std)

# Fit with MCMC
result = fit(model, y=y_train, method="mcmc", x=x_train)
preds = result.predict(x=x_test)

# Fit with Stein Variational Gradient Descent
result = fit(model, y=y_train, method="svgd", num_steps=500, x=x_train)
preds = result.predict(x=x_test)

Constants (non-array kwargs) are automatically bound via functools.partial, so you never need to wrap your model manually:

def model(x, n_conditions, y=None):
    ...

# n_conditions is an int → auto-bound; x is an array → batched
result = fit(model, y=y_train, batch_size=4096, num_epochs=250,
             x=x_train, n_conditions=10)
class blayers.fit.Predictions(mean, std, samples)[source]#

Bases: object

Posterior predictive output from FittedModel.predict().

Parameters:

  • mean (Array)

  • std (Array)

  • samples (Array)

mean: Array#

Point predictions averaged over posterior samples. Shape (n,).

std: Array#

Predictive standard deviation over posterior samples. Shape (n,).

samples: Array#

Raw posterior predictive draws. Shape (num_samples, n, ...).

__init__(mean, std, samples)#
Parameters:

  • mean (Array)

  • std (Array)

  • samples (Array)

Return type:

None

class blayers.fit.FittedModel(model_fn, method, params=None, guide=None, losses=None, posterior_samples=None, mcmc=None, num_particles=None)[source]#

Bases: object

A fitted blayers model.

Created by fit(). Provides predict() for posterior predictive inference and summary() for inspecting latent variable posteriors.

Parameters:

  • model_fn (Callable[[...], Any])

  • method (str)

  • params (dict[str, Any] | None)

  • guide (Any | None)

  • losses (Array | None)

  • posterior_samples (dict[str, Any] | None)

  • mcmc (Any | None)

  • num_particles (int | None)

model_fn: Callable[[...], Any]#

The model function with any constants already bound.

method: str#

One of "vi", "mcmc", or "svgd".

params: dict[str, Any] | None#

SVI / SVGD parameters.

guide: Any | None#

Fitted variational / Stein guide.

losses: Array | None#

Per-step loss curve (VI / SVGD).

posterior_samples: dict[str, Any] | None#

MCMC posterior samples (MCMC only).

__init__(model_fn, method, params=None, guide=None, losses=None, posterior_samples=None, mcmc=None, num_particles=None)#
Parameters:

  • model_fn (Callable[[...], Any])

  • method (str)

  • params (dict[str, Any] | None)

  • guide (Any | None)

  • losses (Array | None)

  • posterior_samples (dict[str, Any] | None)

  • mcmc (Any | None)

  • num_particles (int | None)

Return type:

None

mcmc: Any | None#

The fitted numpyro.infer.MCMC object (MCMC only). Retained so to_arviz() can hand it straight to arviz.from_numpyro for divergences, R-hat, ESS, and log-likelihood.

num_particles: int | None#

Number of Stein particles (SVGD only).

predict(*, num_samples=100, seed=1, **data)[source]#

Generate posterior predictive predictions on new data.

Parameters:

  • num_samples (int) – Number of posterior samples to draw. For VI this controls the guide; for MCMC all posterior samples are used regardless.

  • seed (int) – Random seed for the predictive distribution. Fixed by default, so repeated calls return identical draws — vary it to see Monte Carlo variability (identical reruns are not method determinism).

  • **data – Model inputs excluding y. Constants that were auto-bound during fit() should not be passed again.

Return type:

Predictions

Notes

.mean / .std assume a continuous outcome. For the classification / discrete links (logit_link, categorical_link, ordinal_link, count links) the label mean is not meaningful — work from .samples instead (e.g. per-observation modal class, or class probabilities via (samples == k).mean(axis=0)).

summary(*, num_samples=1000, seed=2, **data)[source]#

Summarize the posterior of each latent variable.

Parameters:

  • num_samples (int) – Samples to draw from the guide (VI only; ignored for MCMC).

  • seed (int) – Random seed. Fixed by default, so repeated calls return identical draws — vary it to see Monte Carlo variability.

  • **data – Model inputs (excluding y) needed so the guide can determine parameter shapes. Required for VI; ignored for MCMC.

Returns:

{site_name: {"mean": ..., "std": ..., "q025": ..., "q975": ..., "shape": ...}}

Return type:

dict

to_arviz(*, y=None, num_samples=1000, seed=3, **data)[source]#

Convert the fit to an ArviZ InferenceData.

Reuses NumPyro’s own ArviZ bridge rather than reimplementing diagnostics:

  • MCMC — delegates to arviz.from_numpyro(mcmc), which carries over posterior draws, sample stats (divergences), and per-observation log-likelihood. y / **data are ignored (already baked into the completed MCMC run).

  • VI — draws from the fitted guide, computes per-observation log-likelihood with numpyro.infer.log_likelihood, and assembles an InferenceData via arviz.from_dict. The observed y (and any model inputs) are required so the log_likelihood and observed_data groups can be built.

The result plugs straight into arviz.summary (R-hat / ESS), arviz.waic, arviz.loo, and the ArviZ plotting suite — so model comparison is az.compare({"a": fit_a.to_arviz(...), ...}).

Parameters:

  • y (jax.Array, optional) – Observed target. Required for VI (to build the log_likelihood and observed_data groups); ignored for MCMC.

  • num_samples (int) – Posterior draws to take from the guide (VI only).

  • seed (int) – Random seed for guide sampling (VI only).

  • **data – Model inputs (e.g. x) needed to evaluate the model. Ignored for MCMC.

Return type:

arviz.InferenceData

Notes

SVGD is not supported: its handful of Stein particles do not form a meaningful sample for WAIC/LOO, and its params do not map cleanly onto model sites. Fit with method="mcmc" or method="vi" for ArviZ-based comparison.

blayers.fit.fit(model_fn, *, y, method='vi', batch_size=None, num_epochs=None, num_steps=None, lr=0.01, schedule='cosine', guide=None, optimizer=None, num_warmup=500, num_mcmc_samples=1000, num_chains=1, autoreparam_model=True, num_particles=10, kernel_fn=None, seed=0, **kwargs)[source]#

Fit a blayers model via variational inference, MCMC, or SVGD.

Keyword arguments that are JAX/numpy arrays are treated as data and batched during training. Non-array keyword arguments (ints, floats, strings, etc.) are treated as constants and bound to the model via functools.partial so they don’t need to be passed again at predict time.

Parameters:

  • model_fn (Callable) – A NumPyro model function that accepts y as a keyword argument.

  • y (jax.Array) – Target / observed values.

  • method ("vi", "mcmc", or "svgd") – Inference method. Default "vi".

  • batch_size (int, optional) – Mini-batch size for VI. If None the full dataset is used each step (appropriate for small datasets).

  • num_epochs (int, optional) – Number of full passes through the data. Exactly one of num_epochs or num_steps is required for VI and SVGD.

  • num_steps (int, optional) – Total number of gradient updates. Exactly one of num_epochs or num_steps is required for VI and SVGD.

  • lr (float) – Peak learning rate (default 0.01). For SVGD this is the Adagrad step size. Ignored when optimizer is given.

  • schedule (str) – LR schedule name: "cosine" (default), "warmup_cosine", or "constant". Only used for VI.

  • guide (type or AutoGuide instance, optional) – Variational family. Pass a class (instantiated on model_fn) or a ready-to-use instance. Default: AutoDiagonalNormal. Not used for SVGD (which auto-generates an AutoDelta guide).

  • optimizer (optax.GradientTransformation, optional) – A fully-constructed optax optimizer. When provided, lr and schedule are ignored. Not used for SVGD.

  • num_warmup (int) – MCMC warmup iterations (default 500).

  • num_mcmc_samples (int) – MCMC posterior samples to draw (default 1000).

  • num_chains (int) – Number of MCMC chains (default 1).

  • autoreparam_model (bool) – Automatically reparameterize LocScale distributions for MCMC (default True).

  • num_particles (int) – Number of Stein particles (default 10). Only used for SVGD.

  • kernel_fn (SteinKernel, optional) – Kernel for SVGD. Default: RBFKernel().

  • seed (int) – Random seed (default 0).

  • **kwargs – Model inputs. Arrays → batched data. Non-arrays → constants bound via partial.

Returns:

Object with .predict(**data) and .summary(**data) methods.

Return type:

FittedModel

Examples

Batched VI (the common case for large datasets):

>>> result = fit(model, y=y_train, batch_size=4096, num_epochs=250,
...              x=x_train, n_conditions=10)
>>> preds = result.predict(x=x_test)

Full-dataset VI (small datasets):

>>> result = fit(model, y=y_train, num_steps=20000, x=x_train)

MCMC:

>>> result = fit(model, y=y_train, method="mcmc", x=x_train)

SVGD:

>>> result = fit(model, y=y_train, method="svgd", num_steps=500,
...              num_particles=20, x=x_train)
blayers.fit.sample_prior(model_fn, *, num_samples=500, seed=0, **data)[source]#

Draw from a model’s prior (and prior predictive) before fitting.

Runs the model with no observed y, so every latent site and the "obs" site are sampled straight from the prior. Use it to sanity-check that your priors imply sensible outcomes — the core of the “tweak priors as you wish” workflow — before committing to inference.

Parameters:

  • model_fn (Callable) – A blayers / NumPyro model. Pass inputs the same way you would to fit(), but without y — supplying y would condition the "obs" site and defeat the purpose.

  • num_samples (int) – Number of prior draws (default 500).

  • seed (int) – Random seed (default 0).

  • **data – Model inputs (e.g. x) and any constants, used to fix the shapes of the sampled sites.

Returns:

{site_name: array of shape (num_samples, *site_shape)}, including the prior-predictive "obs" site.

Return type:

dict

Examples

>>> prior = sample_prior(model, x=x_train, num_samples=1000)
>>> prior["obs"].shape        # (1000, n) prior-predictive outcomes
>>> prior["AdaptiveLayer_mu_beta"].mean(axis=0)   # prior mean of a latent