Use the JAX backend

This guide explains how to use the JAX backend in Meridian.

Introduction to the JAX backend

Meridian uses JAX as its default numerical backend for core numerical operations and probabilistic Markov Chain Monte Carlo (MCMC) sampling (starting in Meridian 2.0). JAX encourages a functional programming style and utilizes XLA (Accelerated Linear Algebra) compilation to offer advanced performance optimizations and memory efficiency.

The legacy TensorFlow backend is deprecated and will be removed in a future release.

Tutorial: To see JAX in action, see the Getting started with JAX notebook.

Backend configuration

By default, Meridian runs on JAX. You don't need to configure any environment variables to use JAX.

Legacy TensorFlow backend (deprecated)

If you need to temporarily run with the legacy TensorFlow backend, set the MERIDIAN_BACKEND environment variable to 'tensorflow' before importing Meridian:

import os

# Select legacy TensorFlow backend (deprecated)
os.environ['MERIDIAN_BACKEND'] = 'tensorflow'

# Now it is safe to import Meridian modules
from meridian.model import model
from meridian.data import load

Precision configuration

By default, Meridian runs with 64-bit precision (float64) on JAX.

Users can choose to run with 32-bit precision (float32) instead, for example:

  • Faster training runtime: 32-bit floating-point operations can run faster on hardware accelerators (such as GPUs or TPUs).
  • Lower memory usage: 32-bit precision reduces memory consumption during MCMC sampling.

To use 32-bit precision, set the MERIDIAN_ENABLE_JAX_X64 environment variable to 'False' (or '0') before importing Meridian:

import os

# Disable 64-bit precision (enable 32-bit precision)
os.environ['MERIDIAN_ENABLE_JAX_X64'] = 'False'

# Now it is safe to import Meridian modules
from meridian.model import model

If the MERIDIAN_ENABLE_JAX_X64 environment variable is unset or set to 'True' or '1', Meridian defaults to 64-bit precision.

Type consistency

Because Meridian operates in 64-bit precision by default, ensure that all user-provided values, custom arrays, and distribution parameters maintain type consistency:

  • Float literals and arrays: Standard Python float literals (such as 0.2, 0.9) default to 64-bit floats. When creating NumPy arrays for priors or inputs, use np.float64 or dtype=np.float64 to match the default precision.
  • Construct custom prior distributions with matching precision: When defining custom prior distributions in PriorDistribution, ensure all distribution parameters (such as loc, scale, concentration0, and concentration1) match the active precision (64-bit float by default).

API differences when using JAX versus TensorFlow

When using the JAX backend, there are key API differences to keep in mind:

Prior distributions

Meridian models use TensorFlow Probability on JAX (tensorflow_probability.substrates.jax). When configuring custom prior distributions under JAX, import tensorflow_probability.substrates.jax as tfp_jax and construct distributions using tfp_jax.distributions.

Ensure all custom distribution parameters use 64-bit precision (such as np.float64 or 64-bit float arrays) to maintain type consistency with Meridian's default precision settings.

JAX

import numpy as np
import tensorflow_probability.substrates.jax as tfp_jax
from meridian.model import constants
from meridian.model import prior_distribution

# Parameters use 64-bit precision
roi_mu = np.float64(0.2)
roi_sigma = np.float64(0.9)
prior = prior_distribution.PriorDistribution(
    roi_m=tfp_jax.distributions.LogNormal(
        roi_mu, roi_sigma, name=constants.ROI_M
    )
)

TensorFlow (Deprecated)

import tensorflow_probability as tfp
from meridian.model import constants
from meridian.model import prior_distribution

roi_mu = 0.2
roi_sigma = 0.9
prior = prior_distribution.PriorDistribution(
    roi_m=tfp.distributions.LogNormal(
        roi_mu, roi_sigma, name=constants.ROI_M
    )
)

Explicit seed requirement

When using the JAX backend, an explicit seed is required for stochastic functions (for example, in sample_posterior()). While TensorFlow uses a global random number generator that automatically picks a random seed, JAX makes this seed explicit. We found no statistically significant differences in ROI estimates or budget shifts across different seeds.

# Explicitly set a seed for MCMC sampling when using the JAX backend
mmm.sample_posterior(
    n_chains=2,
    n_adapt=1000,
    n_burnin=500,
    n_keep=1000,
    seed=0,
)

For more information on JAX random numbers and seeds, refer to the JAX pseudorandom numbers documentation.

Numerical differences and reproducibility

Because TensorFlow and JAX compile their computational graphs differently, you may observe minor numerical differences in your posterior estimates when switching to JAX using the same data and random seeds.

While posterior distributions might not be identical across backends, the differences are generally small and not statistically significant for business metrics such as ROI and budget allocation. This ensures that switching to the JAX backend maintains the integrity of your model's insights.

Performance considerations

Internal testing found JAX supercharged initial model runs, cutting average runtime by ~40% and memory usage by ~70%, compared to TensorFlow when using GPUs. JAX also streamlined model iterations, enabling 2x faster runtimes, 4x less memory usage, and uninterrupted workflows by eliminating the need for kernel restarts.

Because of the increased memory efficiency, you have more headroom to adjust computationally intensive parameters. For example, in Meridian.sample_posterior(), you might increase the unrolled_leapfrog_steps argument (e.g., from 1 to 5). This can accelerate convergence by increasing the trajectory length of the No-U-Turn-Sampler (NUTS) without exceeding hardware memory limits. You can also increase the n_adapt parameter to further aid convergence during the adaptation phase.

Previous
Set the adstock_decay_spec parameter
Next
National models