Skip to content

API Reference

pyrox is probabilistic modeling with Equinox and NumPyro: Bayesian neural networks, Gaussian processes, and composable GP building blocks. The reference is organised by subpackage:

Section What's inside
Core The Equinox-to-NumPyro bridge — PyroxModule, PyroxParam, PyroxSample, Parameterized, pyrox_method
GP Kernels, sparse/variational GPs, guides, likelihoods, non-Gaussian inference strategies, pathwise samplers, Markov (Kalman) GPs, multi-output GPs
NN Bayesian and uncertainty-aware layers — dense variants, spectral / random-feature layers, SNGP, ensembles, heteroscedastic heads, the BNF stack
NN — Geo encoders Longitude/latitude, cyclic, spherical-harmonic, and Slepian input encoders
NN — Conditioning FiLM, affine conditioners, and hypernetworks for conditional neural fields
NN — MFN Multiplicative filter networks (Fourier / Gabor) and their Bayesian variants
Inference Ensemble MAP / VI runners and primitives on top of NumPyro
Preprocessing Pandas-aware spatiotemporal feature extraction for the BNF workflow
Estimator The high-level scikit-learn-style BayesianNeuralFieldEstimator

Conventions

A few patterns hold across the whole package:

  • Modules are pytrees. Everything is an equinox.Module (often a Parameterized subclass from Core) — immutable, jit / grad / vmap-safe, and rendered into NumPyro models via pyrox_sample sites.

  • Three integration patterns. Deterministic modules work with eqx.tree_at surgery (Pattern A), PyroxModule registers params/sites automatically (Pattern B), and Parameterized adds constraint-aware set_prior / autoguide / set_mode("model" | "guide") (Pattern C). See the regression masterclass notebooks for the same model written all three ways.

  • gaussx owns the linear algebra. Solver strategies (gaussx.DenseSolver, CGSolver, BBMMSolver, …), structured operators, and Gaussian distributions come from gaussx; every pyrox entry point with a solver= keyword accepts any gaussx.AbstractSolverStrategy.

  • geonnax owns the deterministic cores. Spherical encoders, SIREN / MFN backbones, random-feature maps, and basis functions are re-exported from geonnax; pyrox wraps them with priors and posteriors.

pyrox: Equinox-to-NumPyro bridge primitives and ensemble inference.

The GP building blocks live in the pyrox-gp package (pyrox_gp) and the Bayesian NN layers in pyrox-nn (pyrox_nn); both build on the primitives exported here.

__version__ = '0.1.1' module-attribute