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

## Introduction to the JAX backend

Meridian uses [JAX](https://docs.jax.dev/en/latest/) 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](https://developers.google.com/meridian/notebook/meridian-getting-started)
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

> [!WARNING]
> **Important:** JAX enforces strict type checking and does not automatically promote or cast mismatched floating-point types in many operations. Mixing `float32` and `float64` values will result in a `TypeError` (for example, `TypeError: Cannot mix
> float32 and float64`).

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.

> [!CAUTION]
> **Caution:** Standard `tfp.distributions` from TensorFlow Probability is incompatible with the JAX backend and will raise runtime errors. Always use `tfp_jax.distributions`.

### 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](https://docs.jax.dev/en/latest/random-numbers.html).

## 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]
> **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.