Bayesian Inference using 'numpyro' and 'XLA'

A high-performance probabilistic programming library that aims to unify the modeling experience by providing an intuitive model-building syntax together with the flexibility of low-level abstraction coding. It also includes pre-built functions for high-level abstraction and supports hardware-accelerated computation for improved scalability, including parallelization, vectorization, and execution on CPU (Central Processing Unit), GPU (Graphics Processing Unit), or TPU (Tensor Processing Unit) using 'JAX' (Just-In-Time compiled Accelerated linear algebra) as the computational backend: Sosa (2026) .

Bayesian Inference Reimagined for R, Python and Julia

One Mental Model. Three Languages.

BFR (BI) provides a unified experience across R, Python, and Julia. Whether you work in R's formula syntax, Python's object-oriented approach, or Julia's mathematical elegance, the model logic remains consistent.

• ✅ Zero Context Switching: Variable names, distribution signatures, and model logic remain consistent.
• ✅ NumPyro Power: Both interfaces compile down to XLA via JAX for blazing fast inference.
• ✅ Rich Diagnostics: Seamless integration with ArviZ for posterior analysis.

Compare the Syntax

model <- function(height, weight){
  # Priors
  sigma = bf.dist.uniform(0, 50, name='sigma', shape=c(1))
  alpha = bf.dist.normal(178, 20, name='alpha', shape=c(1))
  beta  = bf.dist.normal(0, 1, name='beta', shape=c(1))

  # Likelihood
  mu = alpha + beta * weight
  bf.dist.normal(mu, sigma, obs=height)
}

def model(height, weight):
    # Priors
    sigma = bf.dist.uniform(0, 50, name='sigma', shape=(1,))
    alpha = bf.dist.normal(178, 20, name='alpha', shape=(1,))
    beta  = bf.dist.normal(0, 1, name='beta', shape=(1,))

    # Likelihood
    mu = alpha + beta * weight
    bf.dist.normal(mu, sigma, obs=height)

@BI function model(weight, height)
    # Priors
    sigma = bf.dist.uniform(0, 50, name='sigma', shape=(1,))
    alpha = bf.dist.normal(178, 20, name='alpha', shape=(1,))
    beta  = bf.dist.normal(0, 1, name='beta', shape=(1,))

    # Likelihood
    mu = alpha + beta * weight
    bf.dist.normal(mu, sigma, obs=height)
end

Built for Speed

Leveraging Just-In-Time (JIT) compilation, BI outperforms traditional engines on standard hardware and unlocks massive scalability on GPU clusters for large datasets.

Benchmark: Network Size 100 (Lower is Better)

> Comparison of execution time for a Social Relations Model. Source: Sosa et al. (2026).

Installation & Setup

1. Install Package

Use devtools to pull the latest development version from GitHub.

if (!requireNamespace("devtools", quietly = TRUE)) install.packages("devtools")
devtools::install_github("https://github.com/BGN-for-ASNA/BIR")

2. Initialize Environment

Run the starting test to create the Python virtual environment managed by reticulate.

library(BayesForge)
# Run the starting test to install Python dependencies
BF_starting_test()

3. Select Backend

Choose 'cpu', 'gpu', or 'tpu' when importing the library.

# Initialize on CPU (default) or GPU/TPU
m <- importBF(platform = 'cpu')

Features

Data Manipulation

• One-hot encoding
• Index variable conversion
• Scaling and normalization

Modeling (via NumPyro)

• Linear & Generalized Linear Models: Regression, Binomial, Poisson, Negative Binomial, etc.
• Hierarchical/Multilevel Models: Varying intercepts and slopes.
• Time Series & Processes: Gaussian Processes, Gaussian Random Walks, State Space Models.
• Mixture Models: GMM, Dirichlet Process Mixtures.
• Network Models: Network-based diffusion, Block models.
• Bayesian Neural Networks (BNN).

Diagnostics (via ArviZ)

• Posterior summary statistics and plots.
• Trace plots, Density plots, Autocorrelation.
• WAIC and LOO (ELPD) model comparison.
• R-hat and Effective Sample Size (ESS).

Available Distributions

The package provides wrappers for a comprehensive set of distributions from NumPyro.

Continuous

• bf.dist.normal, bf.dist.uniform, bf.dist.student_t
• bf.dist.cauchy, bf.dist.halfcauchy, bf.dist.halfnormal
• bf.dist.gamma, bf.dist.inverse_gamma, bf.dist.exponential
• bf.dist.beta, bf.dist.beta_proportion
• bf.dist.laplace, bf.dist.asymmetric_laplace
• bf.dist.log_normal, bf.dist.log_uniform
• bf.dist.pareto, bf.dist.weibull, bf.dist.gumbel
• bf.dist.chi2, bf.dist.gompertz

Discrete

• bf.dist.bernoulli, bf.dist.bfnomial
• bf.dist.poisson, bf.dist.negative_binomial
• bf.dist.geometric, bf.dist.discrete_uniform
• bf.dist.beta_binomial, bf.dist.zero_inflated_poisson

Multivariate

• bf.dist.multivariate_normal, bf.dist.multivariate_student_t
• bf.dist.dirichlet, bf.dist.dirichlet_multinomial
• bf.dist.multinomial
• bf.dist.lkj, bf.dist.lkj_cholesky
• bf.dist.wishart, bf.dist.wishart_cholesky

Time Series & Stochastic Processes

• bf.dist.gaussian_random_walk
• bf.dist.gaussian_state_space
• bf.dist.euler_maruyama
• bf.dist.car (Conditional AutoRegressive)

Mixtures & Truncated

• bf.dist.mixture, bf.dist.mixture_same_family
• bf.dist.truncated_normal, bf.dist.truncated_cauchy
• bf.dist.lower_truncated_power_law

(See package documentation for the full list)

Documentation

For full documentation of functions and parameters, you can use the built-in R help or the package helper:

# Open package documentation
bf.doc()

# Help for a specific function
?bf.dist.normal

Platform Support

• ✅ Linux
• ✅ macOS
• ✅ Windows

GPU support available on compatible systems with JAX GPU installation.

Related Packages

• BF - Python implementation
• BFJ - J implementation