0.3.1

.github/workflows/claude-code-review.yml

@@ -41,4 +41,3 @@ jobs:
           prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}'
           # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
           # or https://code.claude.com/docs/en/cli-reference for available options
-

.github/workflows/claude.yml

@@ -47,4 +47,3 @@ jobs:
           # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
           # or https://code.claude.com/docs/en/cli-reference for available options
           # claude_args: '--allowed-tools Bash(gh pr:*)'
-

.github/workflows/docs.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11
+          python-version: "3.12"
 
       - name: Install project with dev dependencies
         run: pip install .[dev]

.github/workflows/publish.yml

@@ -18,7 +18,7 @@ jobs:
 
       - uses: actions/setup-python@v5
         with:
-          python-version: "3.11"
+          python-version: "3.12"
 
       - name: Build
         run: |

.github/workflows/test.yml

@@ -9,7 +9,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
-          python-version: 3.11
+          python-version: "3.12"
 
       - name: Install deps
         run: |
@@ -22,4 +22,4 @@ jobs:
       - name: Upload coverage to Coveralls
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: coveralls --service=github
+        run: coveralls --service=github

.pre-commit-config.yaml

@@ -7,28 +7,25 @@ repos:
       - id: end-of-file-fixer
       - id: check-yaml
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 23.12.1
     hooks:
       - id: black
+        language_version: python3.11
   - repo: https://github.com/PyCQA/isort
     rev: 5.12.0  # Known working version (not Poetry-based)
     hooks:
       - id: isort
         language_version: python3.11
 
-  - repo: https://github.com/psf/black
-    rev: 23.12.1
-    hooks:
-      - id: black
-        language_version: python3.11
-
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.8.0
     hooks:
       - id: mypy
         additional_dependencies: [types-requests]  # Add any stubs your project needs
         language_version: python3.11
-        exclude: ^docs/
+        # Type-check the library strictly; tests are exercised by pytest, not
+        # type-checked (the suite is only partially annotated).
+        exclude: ^(docs|tests)/
 
   - repo: https://github.com/myint/autoflake
     rev: v2.2.1

CHANGELOG.md

@@ -0,0 +1,58 @@
+# Changelog
+
+All notable changes to BLayers are documented here. The format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the project aims
+to follow semantic versioning (with the usual 0.x caveat that minor releases
+may carry breaking changes).
+
+## [0.3.1]
+
+### Added
+- `FittedModel.to_arviz()` — convert a fit to an ArviZ `InferenceData` for
+  diagnostics (R-hat, ESS, divergences) and model comparison (PSIS-LOO via
+  `az.loo`, `az.compare`). MCMC uses `arviz.from_numpyro`; VI builds the
+  `log_likelihood` group via `numpyro.infer.log_likelihood`. SVGD is
+  unsupported. Requires the optional `blayers[arviz]` extra (arviz >= 1.0).
+- `sample_prior(model, **inputs)` — draw from the prior / prior predictive
+  before fitting, for prior checks.
+- `categorical_link` — Categorical (softmax) likelihood for multiclass
+  classification (`units = num_classes`).
+- `__version__` on the top-level package.
+
+### Changed
+- **`Batched_Trace_ELBO` now raises `ValueError` on models that use
+  `numpyro.plate`** (previously it emitted a `UserWarning` and continued).
+  The `num_obs / batch_size` rescaling double-counts plate-subsampled sites,
+  so the ELBO was silently wrong — better to fail closed. Use the standard
+  `Trace_ELBO` with plates instead.
+- **Requires Python >= 3.12** (was declared `>=3.9`, but the codebase's
+  `X | None` annotations never actually supported 3.9; arviz 1.x also needs
+  3.12). CI, docs, and publish workflows now run on 3.12.
+- Documented BLayers' scope in the README: it is a **structured Bayesian
+  regression** toolkit (GLMs, hierarchical models, factorization machines,
+  splines, sparse priors) whose layers are *added* into a linear predictor —
+  not stacked into a deep network. For true Bayesian neural nets, use
+  NumPyro's `random_flax_module` / `random_haiku_module`.
+
+### Removed
+- **`AttentionLayer`** — removed. It was the one primitive at odds with the
+  library's additive/interpretable focus, and mean-field VI serves its
+  weight space poorly. If you need it, pin `blayers==0.3.0`, or use
+  NumPyro's neural-network module integration.
+
+### Fixed
+- Minibatch VI now **shuffles** the data each epoch (`svi_run_batched` /
+  `yield_batches`), instead of iterating the same fixed batches in the same
+  order every pass. Removes a bias in the ELBO gradient estimate, especially
+  on sorted data.
+- `EmbeddingLayer` / `RandomEffectsLayer` / `RandomWalkLayer` index lookups no
+  longer collapse a single-row batch to a scalar and now accept float-typed
+  indices (`reshape(-1).astype(int)` instead of `squeeze()`).
+- Fixed a duplicated (and mutually conflicting) `black` hook in the
+  pre-commit config.
+
+### Documentation
+- `Batched_Trace_ELBO` documents that it assumes **all latents are global**
+  (per-observation latents are unsupported in batched mode).
+- `FittedModel.predict` notes that `.mean` / `.std` are not meaningful for
+  classification / discrete links — work from `.samples` instead.

README.md

@@ -1,4 +1,4 @@
-[![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)
+[![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)](https://github.com/georgeberry/blayers/blob/main/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)
 
 
 
@@ -26,6 +26,18 @@ tweak priors as you wish.
 
 Inspiration from Keras and Tensorflow Probability, but made specifically for Numpyro + Jax.
 
+**Scope.** BLayers is 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(...) + ...`), the way you'd build a GLM or GAM — not stacked into a deep
+network. Each term stays interpretable, and the priors and inference (NUTS / VI /
+SVGD) are chosen for honest posteriors over a modest number of meaningful
+parameters. If you want a true Bayesian *neural network* (composed nonlinear
+layers, weight-space inference), reach for
+[`numpyro.contrib.module`](https://num.pyro.ai/en/stable/primitives.html#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
@@ -161,7 +173,6 @@ The full set of layers included with BLayers:
 - `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.
-- `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.
 
 All layer prior kwargs are validated at construction time — bad kwargs raise `TypeError` immediately.
 
@@ -171,16 +182,18 @@ We provide link helpers in `links.py` to reduce Numpyro boilerplate. Available l
 
 - `gaussian_link` — Gaussian likelihood with configurable sigma prior (see below).
 - `lognormal_link` — LogNormal likelihood with configurable sigma prior.
-- `logit_link` — Bernoulli link for logistic regression.
+- `student_t_link` — StudentT likelihood for robust regression (default `df=4`).
+- `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 via `Exponential`.
 - `ordinal_link` — Cumulative logit / proportional odds for ordinal outcomes.
 - `zip_link` — Zero-inflated Poisson for count data with excess zeros.
 - `beta_link` — Beta regression for proportions strictly in (0, 1).
 
-### `gaussian_link` and `lognormal_link`
+### `gaussian_link`, `lognormal_link`, and `student_t_link`
 
-Both links are built on a common base and support three scale modes:
+All three share a common location-scale base and support three scale modes:
 
 ```python
 from blayers.layers import AdaptiveLayer
@@ -277,6 +290,31 @@ 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**.
 
+### Diagnostics & model comparison (ArviZ)
+
+`result.to_arviz()` hands the fit to [ArviZ](https://python.arviz.org) 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).
+
+```python
+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.
+
 ## Batched loss
 
 The default Numpyro way to fit batched VI models is to use `plate`, which confuses
@@ -304,17 +342,19 @@ svi_result = svi_run_batched(
 
 **⚠️⚠️⚠️ `numpyro.plate` + `Batched_Trace_ELBO` do not mix. ⚠️⚠️⚠️**
 
-`Batched_Trace_ELBO` is known to have issues when your model uses `numpyro.plate`. If your model needs plates, either:
+`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:
 1. Batch via `plate` and use the standard `Trace_ELBO`, or
 1. Remove plates and use `Batched_Trace_ELBO` + `svi_run_batched`.
 
-`Batched_Trace_ELBO` will warn if your model has plates.
+`Batched_Trace_ELBO` **raises `ValueError`** if your model contains a plate.
 
 
 ### Reparameterizing
 
 To 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).
 
+> **Note:** `fit(method="mcmc")` already applies `@autoreparam` for you (controlled by `autoreparam_model=True`, on by default). You only need to apply the decorator yourself when driving NUTS / HMC manually, as shown below.
+
 ```python
 from numpyro.infer import MCMC, NUTS
 from blayers.layers import AdaptiveLayer

TODO.md

@@ -0,0 +1,62 @@
+# BLayers roadmap / to-dos
+
+Positioning: BLayers is a **structured Bayesian regression** toolkit (GLMs,
+hierarchical models, factorization machines, splines, sparse priors) — layers are
+*added* into a linear predictor, not stacked into a deep net. The list below is
+ordered by value/effort for that niche.
+
+## Done
+- [x] Cut `AttentionLayer` from core (off-brand; worst-served by mean-field VI;
+      interpretability oversold). Removed from `layers.py`, `__init__.py`, README,
+      and tests.
+- [x] README scope note clarifying GLM/GAM focus + pointer to
+      `random_flax_module` for true Bayesian neural nets.
+
+## Tier 1 — Bayesian workflow tooling (highest leverage)
+Inference exists; evaluation barely does. Mostly plumbing over NumPyro/ArviZ.
+- [x] `FittedModel.to_arviz()` — MCMC via `az.from_numpyro` (posterior +
+      sample_stats + log_likelihood, coerced to NumPy for arviz-stats); VI via
+      `az.from_dict` + `numpyro.infer.log_likelihood`. Unlocks `az.summary`
+      (R-hat/ESS), `az.loo` (PSIS-LOO), and `az.compare`. SVGD unsupported.
+      Optional `blayers[arviz]` extra (arviz >= 1.0, Python >= 3.12).
+      NOTE: arviz 1.x dropped WAIC — LOO is the comparison metric.
+- [ ] MCMC diagnostics surfaced in `summary()` too (R-hat, ESS, divergence
+      count) for users who don't reach for ArviZ.
+- [x] `sample_prior(model, num_samples=...)` prior-predictive helper. Returns
+      the raw draws dict (latents + prior-predictive `obs`); rejects `y`.
+      Exported from `blayers`.
+
+## Tier 2 — Close the GLM likelihood gaps
+- [x] `categorical_link` (multiclass softmax) — takes `(n, num_classes)` logits
+      from a layer's `units=K`; reads K from the trailing dim. Exported.
+- [ ] `gamma_link` / `exponential_link` (positive continuous, survival).
+- [ ] Censored / Tobit likelihood.
+- [ ] `zinb_link` (zero-inflated negative binomial; have ZIP, not ZINB).
+
+## Tier 3 — Production readiness
+- [ ] `FittedModel.save()` / `load()` (params are pytrees — pickle / orbax /
+      safetensors).
+- [ ] Guide shortcuts in `fit()`: `guide="mvn" | "lowrank" | "flow" | "laplace"`,
+      plus `init_loc_fn` passthrough. (Diagonal-normal VI underestimates the
+      posterior correlations that hierarchical models produce.)
+
+## Tier 4 — New marquee layer
+- [ ] Hilbert-Space approximate GP layer (HSGP, Riutort-Mayol et al.) — reduces to
+      a basis-function layer, fast/batchable, sits naturally next to splines and
+      `RandomWalkLayer`.
+
+## Correctness / robustness fixes (small, do alongside)
+- [ ] `_utils.yield_batches` never shuffles — same fixed batches, same order every
+      epoch. Add per-epoch permutation (biases minibatch VI, esp. on sorted data).
+- [ ] Document that `Batched_Trace_ELBO` assumes **all latents are global** (it
+      rescales the whole observed log-lik by N/B and never subsamples local
+      latents). State as a hard constraint, not just a plate warning.
+- [ ] `EmbeddingLayer` / `RandomEffectsLayer` use `theta[x.squeeze()]` — `squeeze`
+      collapses a size-1 batch to a scalar index and misbehaves on multi-column x.
+      Prefer `x.reshape(-1).astype(int)`.
+- [ ] Note that `predict`/`summary` default seeds are constant (1, 2) so identical
+      reruns aren't mistaken for method determinism.
+
+## Docs
+- [ ] Short "how BLayers composes with `random_flax_module`" note for people who
+      want to mix structured terms with a neural component.

blayers/__init__.py

@@ -1,24 +1,27 @@
+from importlib.metadata import PackageNotFoundError, version
+
+from blayers.decorators import autoreparam, autoreshape
+from blayers.fit import FittedModel, Predictions, fit, sample_prior
 from blayers.layers import (
     AdaptiveLayer,
-    AttentionLayer,
     BilinearLayer,
     EmbeddingLayer,
     FixedPriorLayer,
-    FMLayer,
     FM3Layer,
+    FMLayer,
     HorseshoeLayer,
     InteractionLayer,
     InterceptLayer,
     LowRankBilinearLayer,
     LowRankInteractionLayer,
-    pairwise_interactions,
     RandomEffectsLayer,
     RandomWalkLayer,
     SpikeAndSlabLayer,
+    pairwise_interactions,
 )
-
 from blayers.links import (
     beta_link,
+    categorical_link,
     gaussian_link,
     logit_link,
     lognormal_link,
@@ -29,21 +32,15 @@
     zip_link,
 )
 
-from blayers.decorators import (
-    autoreparam,
-    autoreshape,
-)
-
-from blayers.fit import (
-    fit,
-    FittedModel,
-    Predictions,
-)
+try:
+    __version__ = version("blayers")
+except PackageNotFoundError:  # package not installed (e.g. running from source)
+    __version__ = "0.0.0"
 
 __all__ = [
+    "__version__",
     # Layers
     "AdaptiveLayer",
-    "AttentionLayer",
     "BilinearLayer",
     "EmbeddingLayer",
     "FixedPriorLayer",
@@ -60,6 +57,7 @@
     "SpikeAndSlabLayer",
     # Links
     "beta_link",
+    "categorical_link",
     "gaussian_link",
     "logit_link",
     "lognormal_link",
@@ -73,6 +71,7 @@
     "autoreshape",
     # Fit
     "fit",
+    "sample_prior",
     "FittedModel",
     "Predictions",
 ]