{"id":34097224,"url":"https://github.com/georgeberry/blayers","last_synced_at":"2026-04-02T14:55:33.681Z","repository":{"id":300801348,"uuid":"1003817980","full_name":"georgeberry/blayers","owner":"georgeberry","description":"The missing layers package for Bayesian inference","archived":false,"fork":false,"pushed_at":"2026-03-17T02:52:12.000Z","size":4018,"stargazers_count":21,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-03-17T09:08:54.841Z","etag":null,"topics":["bayesian","bayesian-statistics","jax","ml","numpyro","python","statistics"],"latest_commit_sha":null,"homepage":"https://georgeberry.github.io/blayers/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/georgeberry.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":".github/CODEOWNERS","security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-06-17T17:59:46.000Z","updated_at":"2026-03-16T22:21:11.000Z","dependencies_parsed_at":"2025-06-23T17:50:19.735Z","dependency_job_id":"4c624f4b-2303-4402-87d6-d44630524d4a","html_url":"https://github.com/georgeberry/blayers","commit_stats":null,"previous_names":["georgeberry/blayers"],"tags_count":15,"template":false,"template_full_name":null,"purl":"pkg:github/georgeberry/blayers","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/georgeberry%2Fblayers","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/georgeberry%2Fblayers/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/georgeberry%2Fblayers/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/georgeberry%2Fblayers/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/georgeberry","download_url":"https://codeload.github.com/georgeberry/blayers/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/georgeberry%2Fblayers/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":31308452,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-04-02T12:59:32.332Z","status":"ssl_error","status_checked_at":"2026-04-02T12:54:48.875Z","response_time":89,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.5:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["bayesian","bayesian-statistics","jax","ml","numpyro","python","statistics"],"created_at":"2025-12-14T15:51:29.424Z","updated_at":"2026-04-02T14:55:33.672Z","avatar_url":"https://github.com/georgeberry.png","language":"Python","readme":"[![Coverage Status](https://coveralls.io/repos/github/georgeberry/blayers/badge.svg?branch=main)](https://coveralls.io/github/georgeberry/blayers?branch=main) [![License](https://img.shields.io/github/license/georgeberry/blayers)](LICENSE) [![PyPI](https://img.shields.io/pypi/v/blayers)](https://pypi.org/project/blayers/) [![Read - Docs](https://img.shields.io/badge/Read-Docs-2ea44f)](https://georgeberry.github.io/blayers/) [![View - GitHub](https://img.shields.io/badge/View-GitHub-89CFF0)](https://github.com/georgeberry/blayers) [![PyPI Downloads](https://static.pepy.tech/badge/blayers)](https://pepy.tech/projects/blayers)\n\n\n\n# BLayers\n\nThe missing layers package for Bayesian inference.\n\n**BLayers is in beta, errors are possible! We invite you to contribute on [GitHub](https://github.com/georgeberry/blayers).**\n\n## Write code immediately\n\n```\npip install blayers\n```\n\ndeps are: `numpyro`, `jax`, and `optax`.\n\n## Concept\n\n\u003cimg width=\"646\" height=\"258\" alt=\"image\" src=\"https://github.com/user-attachments/assets/21608d4a-fe83-4ebd-a8eb-a67774ea115f\" /\u003e\n\n\nEasily build Bayesian models from parts, abstract away the boilerplate, and\ntweak priors as you wish.\n\nInspiration from Keras and Tensorflow Probability, but made specifically for Numpyro + Jax.\n\nBLayers provides tools to\n\n- Quickly build Bayesian models from layers which encapsulate useful model parts\n- Fit models either using Variational Inference (VI) or your sampling method of\nchoice without having to rewrite models\n- Write pure Numpyro to integrate with all of Numpyro's super powerful tools\n- Add more complex layers (model parts) as you wish\n- Fit models in a greater variety of ways with less code\n\n## The starting point\n\nThe simplest non-trivial (and most important!) Bayesian regression model form is\nthe adaptive prior,\n\n```\nscale ~ HalfNormal(1)\nbeta  ~ Normal(0, scale)\ny     ~ Normal(beta * x, 1)\n```\n\nBLayers encapsulates a generative model structure like this in a `BLayer`. The\nfundamental building block is the `AdaptiveLayer`.\n\n```python\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\ndef model(x, y):\n    mu = AdaptiveLayer()('mu', x)\n    return gaussian_link(mu, y)\n```\n\nAll `AdaptiveLayer` is doing is writing Numpyro for you under the hood. This\nmodel is exactly equivalent to writing the following, just using way less code.\n\n```python\nimport jax.numpy as jnp\nfrom numpyro import distributions, sample\n\ndef model(x, y):\n    # Adaptive layer does all of this\n    input_shape = x.shape[1]\n    # adaptive prior\n    scale = sample(\n        name=\"scale\",\n        fn=distributions.HalfNormal(1.),\n    )\n    # beta coefficients for regression\n    beta = sample(\n        name=\"beta\",\n        fn=distributions.Normal(loc=0., scale=scale),\n        sample_shape=(input_shape,),\n    )\n    mu = jnp.einsum('ij,j-\u003ei', x, beta)\n\n    # the link function does this\n    sigma = sample(name='sigma', fn=distributions.Exponential(1.))\n    return sample('obs', distributions.Normal(mu, sigma), obs=y)\n```\n\n### Mixing it up\n\nThe `AdaptiveLayer` is also fully parameterizable via arguments to the class, so let's say you wanted to change the model from\n\n```\nscale ~ HalfNormal(1)\nbeta  ~ Normal(0, scale)\ny     ~ Normal(beta * x, 1)\n```\n\nto\n\n```\nscale ~ Exponential(1.)\nbeta  ~ LogNormal(0, scale)\ny     ~ Normal(beta * x, 1)\n```\n\nyou can just do this directly via arguments\n\n```python\nfrom numpyro import distributions\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\ndef model(x, y):\n    mu = AdaptiveLayer(\n        scale_dist=distributions.Exponential,\n        coef_dist=distributions.LogNormal,\n        scale_kwargs={'rate': 1.},\n        coef_kwargs={'loc': 0.}\n    )('mu', x)\n    return gaussian_link(mu, y)\n```\n\n### \"Factories\"\n\nSince Numpyro traces `sample` sites and doesn't record any parameters on the class, you can re-use with a particular generative model structure freely.\n\n```python\nfrom numpyro import distributions\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\nmy_lognormal_layer = AdaptiveLayer(\n    scale_dist=distributions.Exponential,\n    coef_dist=distributions.LogNormal,\n    scale_kwargs={'rate': 1.},\n    coef_kwargs={'loc': 0.}\n)\n\ndef model(x, y):\n    mu = my_lognormal_layer('mu1', x) + my_lognormal_layer('mu2', x**2)\n    return gaussian_link(mu, y)\n```\n\n## Layers\n\nThe full set of layers included with BLayers:\n\n- `AdaptiveLayer` — Adaptive prior layer: `scale ~ HalfNormal(1)`, `beta ~ Normal(0, scale)`.\n- `FixedPriorLayer` — Fixed prior over coefficients (e.g., Normal or Laplace), no hierarchical scale.\n- `InterceptLayer` — Intercept-only layer (bias term).\n- `EmbeddingLayer` — Bayesian embeddings for sparse categorical features.\n- `RandomEffectsLayer` — Classical random-effects (embedding with output dim 1).\n- `FMLayer` — Factorization Machine (order 2) for pairwise interaction terms.\n- `FM3Layer` — Factorization Machine (order 3).\n- `LowRankInteractionLayer` — Low-rank interaction between two feature sets.\n- `InteractionLayer` — All pairwise interactions between two feature sets.\n- `BilinearLayer` — Bilinear interaction: `x^T W z`.\n- `LowRankBilinearLayer` — Low-rank bilinear interaction.\n- `RandomWalkLayer` — Gaussian random walk prior over an ordered index (e.g., time).\n- `HorseshoeLayer` — Horseshoe prior for sparse regression; global-local shrinkage via HalfCauchy.\n- `SpikeAndSlabLayer` — Spike-and-slab prior; `z ~ Beta(0.5, 0.5)` inclusion weights times a configurable slab.\n- `AttentionLayer` — Multi-head self-attention over the feature dimension with FT-Transformer tokenisation ([Gorishniy et al. 2021](https://arxiv.org/abs/2106.11959)). `head_dim` is per-head so total embedding dim is `head_dim * num_heads` — adding heads increases capacity.\n\nAll layer prior kwargs are validated at construction time — bad kwargs raise `TypeError` immediately.\n\n## Links\n\nWe provide link helpers in `links.py` to reduce Numpyro boilerplate. Available links:\n\n- `gaussian_link` — Gaussian likelihood with configurable sigma prior (see below).\n- `lognormal_link` — LogNormal likelihood with configurable sigma prior.\n- `logit_link` — Bernoulli link for logistic regression.\n- `poisson_link` — Poisson link with log-rate input.\n- `negative_binomial_link` — NegativeBinomial2 for overdispersed counts; learned concentration via `Exponential`.\n- `ordinal_link` — Cumulative logit / proportional odds for ordinal outcomes.\n- `zip_link` — Zero-inflated Poisson for count data with excess zeros.\n- `beta_link` — Beta regression for proportions strictly in (0, 1).\n\n### `gaussian_link` and `lognormal_link`\n\nBoth links are built on a common base and support three scale modes:\n\n```python\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\n# Default: sigma ~ Exp(1) learned from data\ngaussian_link(mu, y)\n\n# Fixed known scale (e.g. from XGBoost quantile regression)\ngaussian_link(mu, y, scale=pred_std)\n\n# Learned scale from a layer — softplus applied internally for stable gradients\nraw = AdaptiveLayer()(\"log_sigma\", x)\ngaussian_link(mu, y, untransformed_scale=raw)\n```\n\nSwap the sigma prior via `functools.partial`:\n\n```python\nfrom functools import partial\nimport numpyro.distributions as dists\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\n# HalfNormal prior instead of Exponential\nhn_gaussian = partial(gaussian_link, sigma_dist=dists.HalfNormal, sigma_kwargs={\"scale\": 1.0})\n\ndef model(x, y=None):\n    mu = AdaptiveLayer()(\"mu\", x)\n    return hn_gaussian(mu, y)\n```\n\n## Splines\n\nNon-linear transformations via B-splines. Compute the basis matrix once with `make_knots` + `bspline_basis`, then pass it to any layer.\n\n```python\nfrom blayers.splines import make_knots, bspline_basis\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\n\nknots = make_knots(x_train, num_knots=10)   # clamped knot vector from data quantiles\n\ndef model(x, y=None):\n    B = bspline_basis(x, knots)             # (n, num_basis) design matrix\n    f = AdaptiveLayer()(\"f\", B)\n    return gaussian_link(f, y)\n```\n\nAdditive models are straightforward:\n\n```python\nknots1 = make_knots(x1_train, num_knots=10)\nknots2 = make_knots(x2_train, num_knots=10)\n\ndef model(x1, x2, y=None):\n    f1 = AdaptiveLayer()(\"f1\", bspline_basis(x1, knots1))\n    f2 = AdaptiveLayer()(\"f2\", bspline_basis(x2, knots2))\n    return gaussian_link(f1 + f2, y)\n```\n\n## fit() helpers\n\n`fit()` handles the guide, ELBO, batching, and LR schedule. The same model runs unchanged under VI, MCMC, or SVGD.\n\n```python\nfrom blayers.fit import fit\nfrom blayers.decorators import autoreshape\nfrom blayers.layers import AdaptiveLayer, InterceptLayer\nfrom blayers.links import gaussian_link\n\n@autoreshape\ndef model(x, y=None):\n    mu = AdaptiveLayer()(\"beta\", x)\n    intercept = InterceptLayer()(\"intercept\")\n    return gaussian_link(mu + intercept, y)\n\n# Variational Inference (default)\nresult = fit(model, y=y, num_steps=1000, batch_size=256, lr=0.01, x=X)\n\n# MCMC\nresult = fit(model, y=y, method=\"mcmc\", num_mcmc_samples=1000, num_warmup=500, x=X)\n\n# SVGD\nresult = fit(model, y=y, method=\"svgd\", num_steps=1000, num_particles=20, x=X)\n```\n\n`result.predict()` returns a `Predictions` object with `.mean`, `.std`, and `.samples`. `result.summary()` returns posterior stats per latent variable.\n\n```python\npreds = result.predict(x=X, num_samples=500)\nsummary = result.summary(x=X)\n```\n\nKeyword arguments that are JAX arrays are treated as **data** (batched during training). Non-array kwargs are bound as **constants**.\n\n## Batched loss\n\nThe default Numpyro way to fit batched VI models is to use `plate`, which confuses\nme a lot. Instead, BLayers provides `Batched_Trace_ELBO` which does not require\nyou to use `plate` to batch in VI. Just drop your model in.\n\n```python\nfrom numpyro.infer import SVI\nfrom numpyro.infer.autoguide import AutoDiagonalNormal\nimport optax\nfrom blayers.vi_infer import Batched_Trace_ELBO, svi_run_batched\n\nloss = Batched_Trace_ELBO(num_obs=len(y), batch_size=1000)\nguide = AutoDiagonalNormal(model_fn)\nsvi = SVI(model_fn, guide, optax.adam(0.01), loss=loss)\n\nsvi_result = svi_run_batched(\n    svi,\n    rng_key,\n    batch_size=1000,\n    num_steps=500,\n    **model_data,\n)\n```\n\n**⚠️⚠️⚠️ `numpyro.plate` + `Batched_Trace_ELBO` do not mix. ⚠️⚠️⚠️**\n\n`Batched_Trace_ELBO` is known to have issues when your model uses `numpyro.plate`. If your model needs plates, either:\n1. Batch via `plate` and use the standard `Trace_ELBO`, or\n1. Remove plates and use `Batched_Trace_ELBO` + `svi_run_batched`.\n\n`Batched_Trace_ELBO` will warn if your model has plates.\n\n\n### Reparameterizing\n\nTo fit MCMC models well it is crucial to [reparameterize](https://num.pyro.ai/en/latest/reparam.html). BLayers helps you do this via `@autoreparam`, which automatically applies `LocScaleReparam` to all `LocScale` distributions in your model (Normal, LogNormal, StudentT, Cauchy, Laplace, Gumbel).\n\n```python\nfrom numpyro.infer import MCMC, NUTS\nfrom blayers.layers import AdaptiveLayer\nfrom blayers.links import gaussian_link\nfrom blayers.decorators import autoreparam\n\ndata = {...}\n\n@autoreparam\ndef model(x, y):\n    mu = AdaptiveLayer()('mu', x)\n    return gaussian_link(mu, y)\n\nkernel = NUTS(model)\nmcmc = MCMC(\n    kernel,\n    num_warmup=500,\n    num_samples=1000,\n    num_chains=1,\n    progress_bar=True,\n)\nmcmc.run(\n    rng_key,\n    **data,\n)\n```\n","funding_links":[],"categories":[],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgeorgeberry%2Fblayers","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgeorgeberry%2Fblayers","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgeorgeberry%2Fblayers/lists"}