Use the JAX backend

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

Introduction to the JAX backend

By default, Meridian uses TensorFlow for its core numerical operations and probabilistic Markov Chain Monte Carlo (MCMC) sampling, providing a robust and thoroughly tested foundation for all modeling tasks.

For projects that benefit from enhanced performance and memory efficiency, Meridian provides the JAX backend. JAX encourages a functional programming style and utilizes XLA (Accelerated Linear Algebra) compilation to offer advanced performance optimizations.

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

How to enable JAX

Because the core mathematical libraries are loaded at initialization, you must instruct Meridian to use JAX before importing any Meridian modules.

To enable JAX, set the MERIDIAN_BACKEND environment variable to 'jax'.

You must set this environment variable in your script before executing any import meridian statements:

import os

# Enable JAX before importing Meridian
os.environ['MERIDIAN_BACKEND'] = 'jax'

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

Enable 32-bit precision

By default, the JAX backend in Meridian uses 64-bit precision (float64) for improved numerical stability and MCMC sampling convergence. Empirical testing demonstrates that 64-bit precision achieves higher convergence rates and more consistent posterior sampling across random seeds, without the high memory footprint of TensorFlow.

However, users may choose to disable 64-bit precision and enable 32-bit precision (float32) for workloads where:

• Faster training runtime is required: 32-bit JAX runs faster on hardware accelerators (GPUs/TPUs) compared to 64-bit JAX.
• Memory reduction is needed: 32-bit JAX provides lower 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'

# Enable JAX backend
os.environ['MERIDIAN_BACKEND'] = 'jax'

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

If an invalid string is provided to the MERIDIAN_BACKEND environment variable, Meridian will issue a RuntimeWarning and default back to standard TensorFlow execution. If the MERIDIAN_ENABLE_JAX_X64 environment variable is unset or set to 'True' or '1', Meridian defaults to 64-bit precision.

API differences when using JAX versus Tensorflow

When transitioning from TensorFlow to the JAX backend, there are key API differences you need to accommodate in your code:

Prior distributions

If you are setting custom prior distributions using PriorDistribution, you must use tfp.substrates.jax.distributions instead of tfp.distributions. For example:

TensorFlow

import tensorflow_probability as tfp
from meridian.model import prior_distribution

prior = prior_distribution.PriorDistribution(
    roi_m=tfp.distributions.LogNormal(0.2, 0.9)
)

JAX

import tensorflow_probability.substrates.jax as tfp_jax
from meridian.model import prior_distribution

prior = prior_distribution.PriorDistribution(
    roi_m=tfp_jax.distributions.LogNormal(0.2, 0.9)
)

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.

Note: These results were produced from internal analyses done using Getting Started Colab on GPU, measuring the initial and post sample_posterior state of the runtime. It is expected that users may experience variation in results depending on exact configurations. However, testing consistently demonstrates efficiency improvements with JAX compared to TensorFlow.

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.