Pith. sign in

REVIEW 26 cited by

BlackJAX: Composable Bayesian inference in JAX

Not yet reviewed by Pith; the record is open.

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2402.10797 v2 pith:AZXHQQ3L submitted 2024-02-16 cs.MS cs.LGstat.COstat.ML

BlackJAX: Composable Bayesian inference in JAX

classification cs.MS cs.LGstat.COstat.ML
keywords blackjaxbayesianinferencemethodsalgorithmscomposabledesignedease
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved
0 comments
read the original abstract

BlackJAX is a library implementing sampling and variational inference algorithms commonly used in Bayesian computation. It is designed for ease of use, speed, and modularity by taking a functional approach to the algorithms' implementation. BlackJAX is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs. The library integrates well with probabilistic programming languages by working directly with the (un-normalized) target log density function. BlackJAX is intended as a collection of low-level, composable implementations of basic statistical 'atoms' that can be combined to perform well-defined Bayesian inference, but also provides high-level routines for ease of use. It is designed for users who need cutting-edge methods, researchers who want to create complex sampling methods, and people who want to learn how these work.

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 26 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. On the triple nature of the PSR J0435+3233 system

    astro-ph.HE 2026-07 conditional novelty 7.0

    PSR J0435+3233 is a hierarchical triple: its anomalous spin-down is an acceleration artifact from a ~1.2 M⊙ outer companion on a ~70-year eccentric orbit.

  2. Torus Graphs for Large Scale Neural Phase Analysis

    cs.LG 2026-05 unverdicted novelty 7.0

    Scalable Torus Graph models with stochastic inference allow large-scale analysis of neural phase couplings and their dynamics in LFP data.

  3. gemlib.mcmc: composable kernels for Metropolis-within-Gibbs sampling schemes

    stat.CO 2026-05 unverdicted novelty 7.0

    gemlib.mcmc supplies composable kernel abstractions for Metropolis-within-Gibbs sampling via writer monads, allowing concise expression and reuse of complex MCMC algorithms for partially observed epidemic models.

  4. Rapid Hubble constant inference from GW170817 using GPU-accelerated nested sampling: prior sensitivity and the limits of post-hoc reweighting

    astro-ph.CO 2026-06 accept novelty 6.5

    Direct sampling under a uniform-in-dL prior raises P(H0>120) from 0.017 to 0.159 for GW170817; post-hoc reweighting recovers only 0.041 because a low-dL mode is undersampled.

  5. Binary neutron stars in the next-generation era: Multi-messenger detection prospects and constraints on the equation of state, mass distribution, and cosmology

    astro-ph.HE 2026-07 conditional novelty 6.0

    With ET (and ET+CE), mock multi-messenger BNS catalogues yield ~40–500 EM counterparts per year and, under ideal recovery, constrain R1.4 to ~0.2 km and H0 to ~1 km s−1 Mpc−1.

  6. Ab Initio Real-Time Gravitational-Wave Parameter Estimation

    gr-qc 2026-07 accept novelty 6.0

    Slice-within-Gibbs nested sampling on modern GPUs delivers well-calibrated BNS parameter estimation in ~12 minutes uncompressed and ~89 seconds with heterodyning, from cold priors.

  7. Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms

    gr-qc 2026-07 accept novelty 6.0

    A piecewise MLP surrogate emulates NRSur7dq4 over its full domain at NR-faithful accuracy with ~1 ms GPU latency and a fully differentiable JAX likelihood pipeline.

  8. The Universal Warmup Path: Automatic Preconditioner Selection for HMC

    stat.CO 2026-07 conditional novelty 6.0

    A dimension-scheduled multi-chain controller automatically selects diagonal vs low-rank HMC preconditioners from warmup evidence and outperforms fixed Fisher and diagonal baselines on ESS per gradient.

  9. SchwarMAX: a GPU-friendly Schwarzschild orbit-superposition modelling framework

    astro-ph.GA 2026-06 unverdicted novelty 6.0

    SchwarMAX delivers a fast GPU-native Schwarzschild modeling code that recovers density profiles and bar pattern speed from mock IFU data of a simulated barred galaxy.

  10. Numerical polology: towards next-generation model-building for cosmology

    astro-ph.CO 2026-06 unverdicted novelty 6.0

    Numerical polology framework samples coupling space to discover ghost-free tensor field theories up to rank three for cosmology, then applies resulting priors to black hole superradiance, dynamical dark energy, and GW data.

  11. Rapid Hubble constant inference from GW170817 using GPU-accelerated nested sampling: prior sensitivity and the limits of post-hoc reweighting

    astro-ph.CO 2026-06 accept novelty 6.0

    GPU-accelerated nested sampling on GW170817 demonstrates that switching to a uniform-in-dL prior shifts the H0 tail and median far more than post-hoc reweighting captures, due to an under-sampled (dL, iota) bimodality.

  12. KiDS-Legacy: Joint analysis of second- and third-order cosmic shear

    astro-ph.CO 2026-06 unverdicted novelty 6.0

    Joint 2nd- and 3rd-order cosmic shear analysis on KiDS-Legacy data produces Ω_m = 0.297^{+0.056}_{-0.040} and S_8 = 0.806^{+0.025}_{-0.023}, consistent with Planck and prior KiDS results while improving Ω_m precision.

  13. Flow-based generative models for amortized Bayesian inference in regression and inverse PDE problems

    physics.comp-ph 2026-06 unverdicted novelty 6.0

    Flow-ABI trains flow-matching models on historical data to produce a set-conditioned functional posterior sampler that delivers near-real-time Bayesian inference for regression and inverse PDE tasks without per-observ...

  14. GenSBI: Generative Methods for Simulation-Based Inference in JAX

    cs.LG 2026-05 unverdicted novelty 6.0

    GenSBI delivers JAX-native implementations of generative SBI methods with transformer backbones and reports near-ideal calibration scores on standard benchmarks.

  15. Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

    cs.LG 2026-04 unverdicted novelty 6.0

    Tempered sequential Monte Carlo samples from a Boltzmann-tilted distribution over controllers to optimize trajectories and policies under differentiable dynamics.

  16. Tempered Sequential Monte Carlo for Trajectory and Policy Optimization with Differentiable Dynamics

    cs.LG 2026-04 unverdicted novelty 6.0

    Tempered sequential Monte Carlo samples efficiently from a temperature-annealed distribution over controller parameters to solve trajectory and policy optimization under differentiable dynamics.

  17. Towards E-Value Based Stopping Rules for Bayesian Deep Ensembles

    cs.LG 2026-04 unverdicted novelty 6.0

    E-value sequential tests enable early stopping of MCMC sampling in Bayesian deep ensembles, often needing only a fraction of the full budget while improving over standard deep ensembles.

  18. Global structure of the time delay likelihood

    stat.ME 2026-02 unverdicted novelty 6.0

    Time delay likelihoods modeled with Gaussian processes develop a boundary-driven W-shape with a global maximum at the true delay and rises at observation window edges, misleading nested sampling and biasing H0 high.

  19. Image reconstruction with the JWST Interferometer

    astro-ph.IM 2025-10 unverdicted novelty 6.0

    Dorito enables diffraction-limited image reconstruction from JWST AMI observations by deconvolving images or Fourier observables using maximum entropy and total variation regularization.

  20. Dynamic or Systematic? Bayesian model selection between dark energy and supernova biases

    astro-ph.CO 2025-09 unverdicted novelty 6.0

    A magnitude offset between low- and high-redshift supernovae beats the Bayesian evidence for flexknot dark energy and reduces DES-5Y/DESI tension.

  21. Empirical-Bayes Unfolding of $\gamma$-ray Spectra

    astro-ph.IM 2026-06 unverdicted novelty 5.0

    Empirical-Bayes hierarchical unfolding for gamma-ray spectra with Poisson ON/OFF likelihood, adaptive Richardson-Lucy prior, and NUTS posterior sampling, yielding spectra consistent with frequentist regularized ML.

  22. Star-formation variability on the star-forming main sequence during the Epoch of Reionization

    astro-ph.GA 2026-06 unverdicted novelty 5.0

    PSD modeling of SFR scatter at six timescales shows dominant variability on 10-30 Myr scales, stronger in lower-mass galaxies at z=3-8.

  23. Corrected Integrated Laplace Approximation for Bayesian Inference in Latent Gaussian Models

    stat.ML 2026-05 unverdicted novelty 5.0

    An importance sampling correction is added to integrated Laplace approximation so that the approximate posterior for latent Gaussian models converges to the true posterior as the number of samples grows.

  24. bde: A Python Package for Bayesian Deep Ensembles via MILE

    cs.LG 2026-05 unverdicted novelty 5.0

    bde is a new Python package that implements Bayesian deep ensembles via efficient JAX-based Microcanonical Langevin Ensembles for tabular regression and classification with uncertainty estimates.

  25. Confronting eikonal and post-Kerr methods with numerical evolution of scalar field perturbations in spacetimes beyond Kerr

    gr-qc 2026-01 unverdicted novelty 5.0

    Numerical simulations benchmark the eikonal and post-Kerr approximations for quasinormal modes in deformed Kerr spacetimes, quantifying their errors relative to expected observational precision.

  26. The mass of TOI-1883 b: A low density super-Neptune in the ridge regime transiting an early-M dwarf

    astro-ph.EP 2026-06 unverdicted novelty 4.0

    Mass of 13.7 Earth masses and density 0.4 g cm^{-3} measured for TOI-1883 b, a super-Neptune in the ridge regime around an early-M dwarf, with implications for disk migration and photoevaporation.