blayers 0.1.0a1__tar.gz

blayers-0.1.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 George Berry
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

blayers-0.1.0a1/PKG-INFO

@@ -0,0 +1,141 @@
+Metadata-Version: 2.4
+Name: blayers
+Version: 0.1.0a1
+Summary: Bayesian layers for NumPyro and Jax
+Author-email: George Berry <george.e.berry@gmail.com>
+Project-URL: Homepage, https://github.com/georgeberry/blayers
+Project-URL: Documentation, https://georgeberry.github.io/blayers/
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: numpyro
+Provides-Extra: dev
+Requires-Dist: optax; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-check; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: autoflake; extra == "dev"
+Requires-Dist: sphinx; extra == "dev"
+Requires-Dist: furo; extra == "dev"
+Requires-Dist: coverage; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Dynamic: license-file
+
+[![codecov](https://codecov.io/gh/georgeberry/blayers/graph/badge.svg?token=ZDGT0C39QM)](https://codecov.io/gh/georgeberry/blayers) [![License](https://img.shields.io/github/license/georgeberry/blayers)](LICENSE) ![version](https://img.shields.io/badge/version-0.1.0a1-blue)
+
+# BLayers
+
+**NOTE: BLayers is in alpha. Expect changes. Feedback welcome.**
+
+The missing layers package for Bayesian inference. Inspiration from Keras and
+Tensorflow Probability, but made specifically for Numpyro + Jax.
+
+Easily build Bayesian models from parts, abstract away the boilerplate, and
+tweak priors as you wish.
+
+Fit models either using Variational Inference (VI) or your sampling method of
+choice. Use BLayer's ELBO implementation to do either batched VI or sampling
+without having to rewrite models.
+
+BLayers helps you write pure Numpyro, so you can integrate it with any Numpyro
+code to build models of arbitrary complexity. It also gives you a recipe to
+build more complex layers as you wish.
+
+## the starting point
+
+The simplest non-trivial (and most important!) Bayesian regression model form is
+the adaptive prior,
+
+```
+lmbda ~ HalfNormal(1)
+beta  ~ Normal(0, lmbda)
+y     ~ Normal(beta * x, 1)
+```
+
+BLayers takes this as its starting point and most fundamental building block,
+providing the flexible `AdaptiveLayer`.
+
+```python
+from blayers import AdaptiveLayer, gaussian_link_exp
+def model(x, y):
+    mu = AdaptiveLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+For you purists out there, we also provide a `FixedPriorLayer` for standard
+L1/L2 regression.
+
+```python
+from blayers import FixedPriorLayer, gaussian_link_exp
+def model(x, y):
+    mu = FixedPriorLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+## additional layers
+
+### factorization machines
+
+Developed in [Rendle 2010](https://jame-zhang.github.io/assets/algo/Factorization-Machines-Rendle2010.pdf) and [Rendle 2011](https://www.ismll.uni-hildesheim.de/pub/pdfs/FreudenthalerRendle_BayesianFactorizationMachines.pdf), FMs provide a low-rank approximation to the `x`-by-`x` interaction matrix. For those familiar with R syntax, it is an approximation to `y ~ x:x`, excluding the x^2 terms.
+
+To fit the equivalent of an r model like `y ~ x*x` (all main effects, x^2 terms, and one-way interaction effects), you'd do
+
+```python
+from blayers import FMLayer, gaussian_link_exp
+def model(x, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('x2', x**2) +
+        FMLayer(low_rank_dim=3)('xx', x)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+### uv decomp
+
+We also provide a standard UV deccomp for low rank interaction terms
+
+```python
+from blayers import LowRankInteractionLayer, gaussian_link_exp
+def model(x, z, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('z', z) +
+        LowRankInteractionLayer(low_rank_dim=3)('xz', x, z)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+## links
+
+We provide link functions as a convenience to abstract away a bit more Numpyro
+boilerplate.
+
+We currently provide
+
+* `gaussian_link_exp`
+
+## batched loss
+
+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.
+
+```python
+from blayers.infer import Batched_Trace_ELBO, svi_run_batched
+
+svi = SVI(model_fn, guide, optax.adam(schedule), loss=loss_instance)
+
+svi_result = svi_run_batched(
+    svi,
+    rng_key,
+    num_steps,
+    batch_size=1000,
+    **model_data,
+)
+```

blayers-0.1.0a1/README.md

@@ -0,0 +1,113 @@
+[![codecov](https://codecov.io/gh/georgeberry/blayers/graph/badge.svg?token=ZDGT0C39QM)](https://codecov.io/gh/georgeberry/blayers) [![License](https://img.shields.io/github/license/georgeberry/blayers)](LICENSE) ![version](https://img.shields.io/badge/version-0.1.0a1-blue)
+
+# BLayers
+
+**NOTE: BLayers is in alpha. Expect changes. Feedback welcome.**
+
+The missing layers package for Bayesian inference. Inspiration from Keras and
+Tensorflow Probability, but made specifically for Numpyro + Jax.
+
+Easily build Bayesian models from parts, abstract away the boilerplate, and
+tweak priors as you wish.
+
+Fit models either using Variational Inference (VI) or your sampling method of
+choice. Use BLayer's ELBO implementation to do either batched VI or sampling
+without having to rewrite models.
+
+BLayers helps you write pure Numpyro, so you can integrate it with any Numpyro
+code to build models of arbitrary complexity. It also gives you a recipe to
+build more complex layers as you wish.
+
+## the starting point
+
+The simplest non-trivial (and most important!) Bayesian regression model form is
+the adaptive prior,
+
+```
+lmbda ~ HalfNormal(1)
+beta  ~ Normal(0, lmbda)
+y     ~ Normal(beta * x, 1)
+```
+
+BLayers takes this as its starting point and most fundamental building block,
+providing the flexible `AdaptiveLayer`.
+
+```python
+from blayers import AdaptiveLayer, gaussian_link_exp
+def model(x, y):
+    mu = AdaptiveLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+For you purists out there, we also provide a `FixedPriorLayer` for standard
+L1/L2 regression.
+
+```python
+from blayers import FixedPriorLayer, gaussian_link_exp
+def model(x, y):
+    mu = FixedPriorLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+## additional layers
+
+### factorization machines
+
+Developed in [Rendle 2010](https://jame-zhang.github.io/assets/algo/Factorization-Machines-Rendle2010.pdf) and [Rendle 2011](https://www.ismll.uni-hildesheim.de/pub/pdfs/FreudenthalerRendle_BayesianFactorizationMachines.pdf), FMs provide a low-rank approximation to the `x`-by-`x` interaction matrix. For those familiar with R syntax, it is an approximation to `y ~ x:x`, excluding the x^2 terms.
+
+To fit the equivalent of an r model like `y ~ x*x` (all main effects, x^2 terms, and one-way interaction effects), you'd do
+
+```python
+from blayers import FMLayer, gaussian_link_exp
+def model(x, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('x2', x**2) +
+        FMLayer(low_rank_dim=3)('xx', x)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+### uv decomp
+
+We also provide a standard UV deccomp for low rank interaction terms
+
+```python
+from blayers import LowRankInteractionLayer, gaussian_link_exp
+def model(x, z, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('z', z) +
+        LowRankInteractionLayer(low_rank_dim=3)('xz', x, z)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+## links
+
+We provide link functions as a convenience to abstract away a bit more Numpyro
+boilerplate.
+
+We currently provide
+
+* `gaussian_link_exp`
+
+## batched loss
+
+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.
+
+```python
+from blayers.infer import Batched_Trace_ELBO, svi_run_batched
+
+svi = SVI(model_fn, guide, optax.adam(schedule), loss=loss_instance)
+
+svi_result = svi_run_batched(
+    svi,
+    rng_key,
+    num_steps,
+    batch_size=1000,
+    **model_data,
+)
+```

blayers-0.1.0a1/blayers/blayers.egg-info/PKG-INFO

@@ -0,0 +1,141 @@
+Metadata-Version: 2.4
+Name: blayers
+Version: 0.1.0a1
+Summary: Bayesian layers for NumPyro and Jax
+Author-email: George Berry <george.e.berry@gmail.com>
+Project-URL: Homepage, https://github.com/georgeberry/blayers
+Project-URL: Documentation, https://georgeberry.github.io/blayers/
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: jax
+Requires-Dist: numpyro
+Provides-Extra: dev
+Requires-Dist: optax; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-check; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: autoflake; extra == "dev"
+Requires-Dist: sphinx; extra == "dev"
+Requires-Dist: furo; extra == "dev"
+Requires-Dist: coverage; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Dynamic: license-file
+
+[![codecov](https://codecov.io/gh/georgeberry/blayers/graph/badge.svg?token=ZDGT0C39QM)](https://codecov.io/gh/georgeberry/blayers) [![License](https://img.shields.io/github/license/georgeberry/blayers)](LICENSE) ![version](https://img.shields.io/badge/version-0.1.0a1-blue)
+
+# BLayers
+
+**NOTE: BLayers is in alpha. Expect changes. Feedback welcome.**
+
+The missing layers package for Bayesian inference. Inspiration from Keras and
+Tensorflow Probability, but made specifically for Numpyro + Jax.
+
+Easily build Bayesian models from parts, abstract away the boilerplate, and
+tweak priors as you wish.
+
+Fit models either using Variational Inference (VI) or your sampling method of
+choice. Use BLayer's ELBO implementation to do either batched VI or sampling
+without having to rewrite models.
+
+BLayers helps you write pure Numpyro, so you can integrate it with any Numpyro
+code to build models of arbitrary complexity. It also gives you a recipe to
+build more complex layers as you wish.
+
+## the starting point
+
+The simplest non-trivial (and most important!) Bayesian regression model form is
+the adaptive prior,
+
+```
+lmbda ~ HalfNormal(1)
+beta  ~ Normal(0, lmbda)
+y     ~ Normal(beta * x, 1)
+```
+
+BLayers takes this as its starting point and most fundamental building block,
+providing the flexible `AdaptiveLayer`.
+
+```python
+from blayers import AdaptiveLayer, gaussian_link_exp
+def model(x, y):
+    mu = AdaptiveLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+For you purists out there, we also provide a `FixedPriorLayer` for standard
+L1/L2 regression.
+
+```python
+from blayers import FixedPriorLayer, gaussian_link_exp
+def model(x, y):
+    mu = FixedPriorLayer()('mu', x)
+    return gaussian_link_exp(mu, y)
+```
+
+## additional layers
+
+### factorization machines
+
+Developed in [Rendle 2010](https://jame-zhang.github.io/assets/algo/Factorization-Machines-Rendle2010.pdf) and [Rendle 2011](https://www.ismll.uni-hildesheim.de/pub/pdfs/FreudenthalerRendle_BayesianFactorizationMachines.pdf), FMs provide a low-rank approximation to the `x`-by-`x` interaction matrix. For those familiar with R syntax, it is an approximation to `y ~ x:x`, excluding the x^2 terms.
+
+To fit the equivalent of an r model like `y ~ x*x` (all main effects, x^2 terms, and one-way interaction effects), you'd do
+
+```python
+from blayers import FMLayer, gaussian_link_exp
+def model(x, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('x2', x**2) +
+        FMLayer(low_rank_dim=3)('xx', x)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+### uv decomp
+
+We also provide a standard UV deccomp for low rank interaction terms
+
+```python
+from blayers import LowRankInteractionLayer, gaussian_link_exp
+def model(x, z, y):
+    mu = (
+        AdaptiveLayer('x', x) +
+        AdaptiveLayer('z', z) +
+        LowRankInteractionLayer(low_rank_dim=3)('xz', x, z)
+    )
+    return gaussian_link_exp(mu, y)
+```
+
+## links
+
+We provide link functions as a convenience to abstract away a bit more Numpyro
+boilerplate.
+
+We currently provide
+
+* `gaussian_link_exp`
+
+## batched loss
+
+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.
+
+```python
+from blayers.infer import Batched_Trace_ELBO, svi_run_batched
+
+svi = SVI(model_fn, guide, optax.adam(schedule), loss=loss_instance)
+
+svi_result = svi_run_batched(
+    svi,
+    rng_key,
+    num_steps,
+    batch_size=1000,
+    **model_data,
+)
+```

blayers-0.1.0a1/blayers/blayers.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+LICENSE
+README.md
+pyproject.toml
+blayers/blayers.egg-info/PKG-INFO
+blayers/blayers.egg-info/SOURCES.txt
+blayers/blayers.egg-info/dependency_links.txt
+blayers/blayers.egg-info/requires.txt
+blayers/blayers.egg-info/top_level.txt
+tests/test_layers.py

blayers-0.1.0a1/blayers/blayers.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

blayers-0.1.0a1/blayers/blayers.egg-info/requires.txt

@@ -0,0 +1,17 @@
+numpy
+jax
+numpyro
+
+[dev]
+optax
+pytest
+pytest-check
+mypy
+black
+isort
+autoflake
+sphinx
+furo
+coverage
+pytest-cov
+pre-commit

blayers-0.1.0a1/blayers/blayers.egg-info/top_level.txt

@@ -0,0 +1 @@
+

blayers-0.1.0a1/pyproject.toml

@@ -0,0 +1,57 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "blayers"
+version = "0.1.0a1"
+description = "Bayesian layers for NumPyro and Jax"
+authors = [{ name = "George Berry", email = "george.e.berry@gmail.com" }]
+readme = "README.md"
+license = { file = "MIT" }
+requires-python = ">=3.9"
+dependencies = [
+    "numpy",
+    "jax",
+    "numpyro",
+]
+
+[project.urls]
+Homepage = "https://github.com/georgeberry/blayers"
+Documentation = "https://georgeberry.github.io/blayers/"
+
+[project.optional-dependencies]
+dev = [
+    "optax",
+    "pytest",
+    "pytest-check",
+    "mypy",
+    "black",
+    "isort",
+    "autoflake",
+    "sphinx",
+    "furo",
+    "coverage",
+    "pytest-cov",
+    "pre-commit",
+]
+
+[tool.setuptools.packages.find]
+where = ["blayers"]
+
+[tool.mypy]
+strict = true
+ignore_missing_imports = true
+explicit_package_bases = true
+disable_error_code = ["misc"]
+
+[tool.isort]
+profile = "black"
+line_length = 80
+multi_line_output = 3
+include_trailing_comma = true
+force_grid_wrap = 0
+
+[tool.black]
+line-length = 80
+target-version = ["py311"]

blayers-0.1.0a1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

blayers-0.1.0a1/tests/test_layers.py

@@ -0,0 +1,388 @@
+from typing import Any, Callable
+
+import jax
+import jax.numpy as jnp
+import jax.random as random
+import numpyro.distributions as dist
+import optax
+import pytest
+import pytest_check
+from _pytest.fixtures import SubRequest
+from numpyro import sample
+from numpyro.infer import SVI, Predictive, Trace_ELBO
+from numpyro.infer.autoguide import AutoDiagonalNormal
+
+from blayers import (
+    AdaptiveLayer,
+    EmbeddingLayer,
+    FixedPriorLayer,
+    FMLayer,
+    LowRankInteractionLayer,
+)
+from blayers.fit_tools import (
+    identity,
+    outer_product,
+    outer_product_upper_tril_no_diag,
+    rmse,
+)
+from blayers.infer import Batched_Trace_ELBO, svi_run_batched
+from blayers.links import gaussian_link_exp
+
+N_OBS = 10000
+LOW_RANK_DIM = 3
+EMB_DIM = 1
+N_EMB_CATEGORIES = 10
+
+
+# ---- Data Generating Processes --------------------------------------------- #
+
+
+def dgp_simple(n_obs: int, k: int) -> dict[str, jax.Array]:
+    lambda1 = sample("lambda1", dist.HalfNormal(1.0))
+    beta = sample("beta", dist.Normal(0, lambda1).expand([k]))
+
+    x1 = sample("x1", dist.Normal(0, 1).expand([n_obs, k]))
+
+    sigma = sample("sigma", dist.HalfNormal(1.0))
+    mu = jnp.dot(x1, beta)
+    y = sample("y", dist.Normal(mu, sigma))
+    return {
+        "x1": x1,
+        "y": y,
+        "beta": beta,
+        "lambda1": lambda1,
+        "sigma": sigma,
+    }
+
+
+def dgp_fm(n_obs: int, k: int) -> dict[str, jax.Array]:
+    x1 = sample("x1", dist.Normal(0, 1).expand([n_obs, k]))
+    lmbda = sample("lambda", dist.HalfNormal(1.0))
+    theta = sample("theta", dist.Normal(0.0, lmbda).expand([k, LOW_RANK_DIM]))
+
+    sigma = sample("sigma", dist.HalfNormal(1.0))
+    mu = FMLayer.matmul(theta, x1)
+    y = sample("y", dist.Normal(mu, sigma))
+    return {
+        "x1": x1,
+        "y": y,
+        "theta": theta,
+        "lambda": lmbda,
+        "sigma": sigma,
+    }
+
+
+def dgp_emb(n_obs: int, k: int, n_categories: int) -> dict[str, jax.Array]:
+    lmbda = sample("lambda", dist.HalfNormal(1.0))
+    beta = sample("beta", dist.Normal(0, lmbda).expand([n_categories, k]))
+    x1 = sample(
+        "x1",
+        dist.Categorical(probs=jnp.ones(n_categories) / n_categories).expand(
+            [n_obs]
+        ),
+    )
+
+    sigma = sample("sigma", dist.HalfNormal(1.0))
+
+    mu = jnp.sum(beta[x1], axis=1)
+    y = sample("y", dist.Normal(mu, sigma))
+    return {
+        "x1": x1,
+        "y": y,
+        "beta": beta,
+        "lambda": lmbda,
+        "sigma": sigma,
+    }
+
+
+def dgp_lowrank(n_obs: int, k: int) -> dict[str, jax.Array]:
+    offset = 1
+
+    x1 = sample("x1", dist.Normal(0, 1).expand([n_obs, k]))
+    x2 = sample("x2", dist.Normal(0, 1).expand([n_obs, k + offset]))
+
+    lambda1 = sample("lambda1", dist.HalfNormal(1.0))
+    theta1_lowrank = sample(
+        "theta1", dist.Normal(0.0, lambda1).expand([k, LOW_RANK_DIM])
+    )
+
+    lambda2 = sample("lambda2", dist.HalfNormal(1.0))
+    theta2_lowrank = sample(
+        "theta2", dist.Normal(0.0, lambda1).expand([k + offset, LOW_RANK_DIM])
+    )
+
+    sigma = sample("sigma", dist.HalfNormal(1.0))
+    mu = LowRankInteractionLayer.matmul(
+        theta1=theta1_lowrank,
+        theta2=theta2_lowrank,
+        x=x1,
+        z=x2,
+    )
+
+    y = sample("y", dist.Normal(mu, sigma))
+    return {
+        "x1": x1,
+        "y": y,
+        "theta1": theta1_lowrank,
+        "theta2": theta2_lowrank,
+        "lambda1": lambda1,
+        "lambda2": lambda2,
+        "sigma": sigma,
+    }
+
+
+def simulated_data(
+    dgp: Callable[..., Any],
+    **kwargs: Any,
+) -> dict[str, jax.Array]:
+    rng_key = random.PRNGKey(0)
+    predictive = Predictive(dgp, num_samples=1)
+    samples = predictive(
+        rng_key,
+        **kwargs,
+    )
+    res = {k: jnp.squeeze(v, axis=0) for k, v in samples.items()}
+    return res
+
+
+@pytest.fixture
+def simulated_data_simple() -> dict[str, jax.Array]:
+    return simulated_data(dgp_simple, n_obs=N_OBS, k=2)
+
+
+@pytest.fixture
+def simulated_data_fm() -> dict[str, jax.Array]:
+    return simulated_data(dgp_fm, n_obs=N_OBS, k=10)
+
+
+@pytest.fixture
+def simulated_data_emb() -> dict[str, jax.Array]:
+    return simulated_data(
+        dgp_emb,
+        n_obs=N_OBS,
+        k=EMB_DIM,
+        n_categories=N_EMB_CATEGORIES,
+    )
+
+
+@pytest.fixture
+def simulated_data_lowrank() -> dict[str, jax.Array]:
+    return simulated_data(
+        dgp_lowrank,
+        n_obs=N_OBS,
+        k=10,
+    )
+
+
+# ---- Models ---------------------------------------------------------------- #
+
+
+@pytest.fixture
+def linear_regression_adaptive_model() -> (
+    tuple[Callable[..., Any], list[tuple[list[str], Callable[..., jax.Array]]]]
+):
+    def model(x1: jax.Array, y: jax.Array | None = None) -> Any:
+        beta = AdaptiveLayer()("beta", x1)
+        return gaussian_link_exp(beta, y)
+
+    return model, [(["AdaptiveLayer_beta_beta"], identity)]
+
+
+@pytest.fixture
+def linear_regression_fixed_model() -> (
+    tuple[Callable[..., Any], list[tuple[list[str], Callable[..., jax.Array]]]]
+):
+    def model(x1: jax.Array, y: jax.Array | None = None) -> Any:
+        beta = FixedPriorLayer()("beta", x1)
+        return gaussian_link_exp(beta, y)
+
+    return model, [(["FixedPriorLayer_beta_beta"], identity)]
+
+
+@pytest.fixture
+def fm_regression_model() -> (
+    tuple[Callable[..., Any], list[tuple[list[str], Callable[..., jax.Array]]]]
+):
+    def model(x1: jax.Array, y: jax.Array | None = None) -> Any:
+        theta = FMLayer(low_rank_dim=LOW_RANK_DIM)("theta", x1)
+        return gaussian_link_exp(theta, y)
+
+    return (
+        model,
+        [
+            (["FMLayer_theta_theta"], outer_product_upper_tril_no_diag),
+        ],
+    )
+
+
+@pytest.fixture
+def emb_model() -> (
+    tuple[Callable[..., Any], list[tuple[list[str], Callable[..., jax.Array]]]]
+):
+    def model(x1: jax.Array, y: jax.Array | None = None) -> Any:
+        beta = EmbeddingLayer()(
+            "beta",
+            x1,
+            n_categories=N_EMB_CATEGORIES,
+            embedding_dim=EMB_DIM,
+        )
+        return gaussian_link_exp(beta, y)
+
+    return (
+        model,
+        [
+            (["EmbeddingLayer_beta_beta"], identity),
+        ],
+    )
+
+
+@pytest.fixture
+def lowrank_model() -> (
+    tuple[Callable[..., Any], list[tuple[list[str], Callable[..., jax.Array]]]]
+):
+    def model(x1: jax.Array, x2: jax.Array, y: jax.Array | None = None) -> Any:
+        beta1 = LowRankInteractionLayer(low_rank_dim=LOW_RANK_DIM)(
+            "lowrank",
+            x1,
+            x2,
+        )
+        return gaussian_link_exp(beta1, y)
+
+    return (
+        model,
+        [
+            (
+                [
+                    "LowRankInteractionLayer_lowrank_theta1",
+                    "LowRankInteractionLayer_lowrank_theta2",
+                ],
+                outer_product,
+            ),
+        ],
+    )
+
+
+# ---- Loss classes ---------------------------------------------------------- #
+
+
+@pytest.fixture
+def trace_elbo() -> Trace_ELBO:
+    return Trace_ELBO()
+
+
+@pytest.fixture
+def trace_elbo_batched() -> Batched_Trace_ELBO:
+    return Batched_Trace_ELBO(n_obs=N_OBS)
+
+
+# ---- Dispatchers ----------------------------------------------------------- #
+
+"""
+These are pytest helpers that let us cycle through fixtures. This setup is a
+little wonky and I'm sure we could come up with something better in the long
+run, but it works for now. Just make one with the name for the thing you want
+to pass to the ultimate test function.
+"""
+
+
+@pytest.fixture
+def model(request: SubRequest) -> Any:
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture
+def data(request: SubRequest) -> Any:
+    return request.getfixturevalue(request.param)
+
+
+@pytest.fixture
+def loss_instance(request: SubRequest) -> Any:
+    return request.getfixturevalue(request.param)
+
+
+# ---- Test functions -------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+    "loss_instance",
+    [
+        "trace_elbo",
+        "trace_elbo_batched",
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    ("model", "data"),
+    [
+        ("linear_regression_adaptive_model", "simulated_data_simple"),
+        ("linear_regression_fixed_model", "simulated_data_simple"),
+        ("fm_regression_model", "simulated_data_fm"),
+        ("emb_model", "simulated_data_emb"),
+        ("lowrank_model", "simulated_data_lowrank"),
+    ],
+    indirect=True,
+)
+def test_models(
+    data: Any,
+    model: Any,
+    loss_instance: Any,
+) -> Any:
+    model_fn, coef_groups = model
+    model_data = {k: v for k, v in data.items() if k in ("y", "x1", "x2")}
+
+    guide = AutoDiagonalNormal(model_fn)
+
+    num_steps = 30000
+
+    schedule = optax.cosine_onecycle_schedule(
+        transition_steps=num_steps,
+        peak_value=5e-2,
+        pct_start=0.1,
+        div_factor=25,
+    )
+
+    svi = SVI(model_fn, guide, optax.adam(schedule), loss=loss_instance)
+
+    rng_key = random.PRNGKey(2)
+
+    if isinstance(loss_instance, Trace_ELBO):
+        svi_result = svi.run(
+            rng_key,
+            num_steps=num_steps,
+            **model_data,
+        )
+    if isinstance(loss_instance, Batched_Trace_ELBO):
+        svi_result = svi_run_batched(
+            svi,
+            rng_key,
+            num_steps,
+            batch_size=1000,
+            **model_data,
+        )
+    guide_predicitive = Predictive(
+        guide,
+        params=svi_result.params,
+        num_samples=1000,
+    )
+    guide_samples = guide_predicitive(
+        random.PRNGKey(1),
+        **{k: v for k, v in model_data.items() if k != "y"},
+    )
+    guide_means = {k: jnp.mean(v, axis=0) for k, v in guide_samples.items()}
+
+    for coef_list, coef_fn in coef_groups:
+        with pytest_check.check:
+            val = rmse(
+                coef_fn(*[guide_means[x] for x in coef_list]),
+                coef_fn(*[data[x.split("_")[2]] for x in coef_list]),
+            )
+            assert val < 0.1
+
+    with pytest_check.check:
+        assert (
+            rmse(
+                guide_means["sigma"],
+                data["sigma"],
+            )
+            < 0.03
+        )