Skip to main content

BayesForge for Python

**A unified probabilistic programming library, bringing JAX-powered BayesForge to the Python, R and Julia ecosystem.**\ *Run bespoke models on CPU, GPU, or TPU with Julia's native syntax.*

Website bioRxiv Python License: GPL (>= 3)

------------------------------------------------------------------------

One Mental Model. Three Languages.

BayesForge (BF) provides a unified experience across Julia, Python, and R. 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 across all implementations.
  • NumPyro Power: All interfaces compile down to XLA via JAX for blazing fast inference.
  • Rich Diagnostics: Seamless integration with ArviZ for posterior analysis.

Compare the Syntax

Python Syntax Julia Syntax R Syntax
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)
@BF 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
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)
}

Built for Speed

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

Benchmark: Network Size 400 (Lower is Better)

+----------------+--------------------------------+----------------------+ | Engine | Execution Time | Relative Performance | +:===============+:===============================+:=====================+ | STAN (CPU) | ████████████████████████████ | Baseline | +----------------+--------------------------------+----------------------+ | BF (CPU) | ████████████ | ~30x Faster | +----------------+--------------------------------+----------------------+ | BF (GPU) | ██ | ~200x Faster | +----------------+--------------------------------+----------------------+

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


Installation & Setup

1. Install Python

Download and install Python 3.10 or later

2. Install Package

From pip

pip install BayesForge

Development Installation

pip install git+https://github.com/BGN-for-ASNA/BF.git

Or clone the repository and activate it locally:

git clone https://github.com/BGN-for-ASNA/BF.git
cd BF

Then in Python:

pip install -e .

3. Initialize Environment

from BF import BF
m = BF()

4. Select Backend

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

# Initialize on CPU (default)
m = BF(platform="cpu")

# Or on GPU (requires JAX GPU installation)
m = BF(platform="gpu")

Quick Start

from BayesForge import bf

# Initialize BF
m = bf('cpu')

# Generate some data
x = m.dist.normal(0, 1, shape=(100,), sample=True)
y = m.dist.normal(0.2 + 0.6 * x, 1.2, sample=True)

# Define a Bayesian linear regression model
def linear_model(x, y):
    alpha = m.dist.normal(loc=0, scale=1, name="alpha")
    beta  = m.dist.normal(loc=0, scale=1, name="beta")
    sigma = m.dist.exponential(1, name="sigma")
    mu = alpha + beta * x
    m.dist.normal(mu, sigma, obs=y)


# Fit the model
m.fit(linear_model, num_warmup=1000, num_samples=1000, num_chains=1)

# Display results
m.summary()

# Plot results with @pyplot
m.plot_trace()

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

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

Discrete

  • m.dist.bernoulli, m.dist.binomial
  • m.dist.poisson, m.dist.negative_binomial
  • m.dist.geometric, m.dist.discrete_uniform
  • m.dist.beta_binomial, m.dist.zero_inflated_poisson

Multivariate

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

Time Series & Stochastic Processes

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

Mixtures & Truncated

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

(See package documentation for the full list)


Documentation

For full documentation and examples visitOfficial website](https://s-sosa.com/BF)


Platform Support

  • ✅ Linux
  • ✅ macOS
  • ✅ Windows

GPU support available on compatible systems with JAX GPU installation.


Related Packages

  • BFR - R implementation
  • BFJ - Julia implementation

::: {align="center"} BayesForge

Based on "The BayesForge library for Python, R, Julia" by Sosa, McElreath, & Ross (2026).

Official website | Issues | Quick Start

© 2026 BayesForge Team. Released under GPL-3.0. :::

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

bayesforge-0.0.57.tar.gz (499.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

bayesforge-0.0.57-py3-none-any.whl (628.2 kB view details)

Uploaded Python 3

File details

Details for the file bayesforge-0.0.57.tar.gz.

File metadata

  • Download URL: bayesforge-0.0.57.tar.gz
  • Upload date:
  • Size: 499.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/7.0.0 CPython/3.13.5

File hashes

Hashes for bayesforge-0.0.57.tar.gz
Algorithm Hash digest
SHA256 5b01a463b82d962961dab182c9f4afbde1a93c7136e07f822d92368ca2826634
MD5 aaaf19b4c5d0646681bbb79b94d875ee
BLAKE2b-256 4e5bf8b9fbf5a4a6fb5cbcf39f0fbfae955f0d1baf7714ed83002484a3443f6b

See more details on using hashes here.

File details

Details for the file bayesforge-0.0.57-py3-none-any.whl.

File metadata

  • Download URL: bayesforge-0.0.57-py3-none-any.whl
  • Upload date:
  • Size: 628.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/7.0.0 CPython/3.13.5

File hashes

Hashes for bayesforge-0.0.57-py3-none-any.whl
Algorithm Hash digest
SHA256 8a93a6ffb0c44037c80c1ff88abda3f75c063d27719cf08fa7ff5944dd8a4465
MD5 ae2403b145c36e6f7b9115829b69f441
BLAKE2b-256 f1a562c22d5999b30812b30c0ffbaebacbc8381bbaaee110946a4947877bc44b

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page