# How to Use TimesFM with JAX/Flax: A Complete Inference Guide

> Learn how to use TimesFM with JAX/Flax! This guide covers inference with JAX and Flax models, supporting JIT compilation and multi-device parallelism for efficient time series forecasting.

- Repository: [Google Research/timesfm](https://github.com/google-research/timesfm)
- Tags: how-to-guide
- Published: 2026-04-02

---

**You can run TimesFM inference on JAX using `TimesFmJax` in [`v1/src/timesfm/timesfm_jax.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_jax.py) or on Flax using `TimesFM_2p5_200M_flax` in [`src/timesfm/timesfm_2p5/timesfm_2p5_flax.py`](https://github.com/google-research/timesfm/blob/main/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py), both supporting JIT compilation and multi-device parallelism via `pmap`.**

The `google-research/timesfm` repository provides production-ready inference APIs for **TimesFM with JAX/Flax**, enabling scalable time-series forecasting on accelerators. Both implementations share a common patch-based transformer backbone but differ in abstraction levels—JAX offers low-level control while Flax provides high-level NNX modules with built-in quantization utilities.

## JAX Implementation: TimesFmJax

The **JAX** API centers on the `TimesFmJax` class defined in [`v1/src/timesfm/timesfm_jax.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_jax.py). This wrapper handles checkpoint restoration, JIT compilation, and device-parallel inference through a functional programming approach.

### Loading Pretrained Checkpoints

Initialize the model by specifying hyperparameters that match your checkpoint architecture, then load weights using `TimesFmCheckpoint`. The `load_from_checkpoint` method automatically downloads from the Hugging Face hub when `path=None` and restores weights via PaxML's `checkpoints.restore_checkpoint`.

```python
from timesfm.timesfm_base import TimesFmCheckpoint
from timesfm.timesfm_jax import TimesFmJax

api = TimesFmJax(
    per_core_batch_size=4,
    context_len=512,
    output_patch_len=128,
    input_patch_len=32,
    model_dims=1280,
    num_layers=20,
    num_heads=16,
)

ckpt = TimesFmCheckpoint(
    path=None,
    huggingface_repo_id="google/timesfm-2.5-200m-jax",
    type=None,
    step=None,
)
api.load_from_checkpoint(ckpt)

```

### Compiling the Decoder

During initialization, `load_from_checkpoint` invokes `jit_decode` to compile the decoder for a specific context length and forecast horizon. This freezes tensor shapes and creates a `pmap`-ed kernel that distributes computation across available devices (`jax.local_device_count`).

### Running Forecast Inference

Call the `forecast` method with a list of 1-D numpy arrays and frequency flags. The method internally iterates over `global_batch_size` and invokes the parallelized decode kernel. It returns the mean prediction and a full quantile tensor.

```python
import numpy as np

inputs = [np.random.randn(600).astype(np.float32) for _ in range(8)]
freq = [0] * len(inputs)  # 0 = high frequency

mean, full = api.forecast(inputs, freq=freq)
print(mean.shape)  # (batch, horizon)

print(full.shape)  # (batch, horizon, 1 + #quantiles)

```

## Flax Implementation: TimesFM_2p5_200M_flax

The **Flax** implementation in [`src/timesfm/timesfm_2p5/timesfm_2p5_flax.py`](https://github.com/google-research/timesfm/blob/main/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py) uses the NNX API to provide an object-oriented interface. The `TimesFM_2p5_200M_flax` class wraps `TimesFM_2p5_200M_flax_module` and adds utilities for continuous quantile heads and quantile-crossing correction.

### Model Initialization and Checkpoint Loading

Use the `from_pretrained` classmethod to download and restore checkpoints via `orbax.checkpoint.StandardCheckpointer`. This eliminates manual configuration of PaxML primitives.

```python
from timesfm.timesfm_2p5.timesfm_2p5_flax import TimesFM_2p5_200M_flax

model = TimesFM_2p5_200M_flax.from_pretrained(
    model_id="google/timesfm-2.5-200m-flax"
)

```

### Configuration and Compilation

Define a `ForecastConfig` from [`src/timesfm/configs.py`](https://github.com/google-research/timesfm/blob/main/src/timesfm/configs.py) specifying `max_context` (must be a multiple of input patch size `p`) and `max_horizon` (must be a multiple of output patch size `o`). The `compile` method validates these constraints, builds an `nnx.pmap` kernel (`compiled_decode_kernel`), and stores a partially-applied callable.

```python
from timesfm.configs import ForecastConfig

fc = ForecastConfig(
    max_context=512,
    max_horizon=256,
    per_core_batch_size=2,
    force_flip_invariance=False,
    use_continuous_quantile_head=True,
    fix_quantile_crossing=True,
)

model.compile(fc, dryrun=True)

```

### Executing the Compiled Decode Kernel

Invoke `compiled_decode` with JAX arrays for inputs and boolean masks (True indicates padding). The method returns the mean forecast (index 5 of the quantile stack) and the full quantile predictions. The Flax module handles `_force_flip_invariance_fn`, `_use_continuous_quantile_head_fn`, and `_fix_quantile_crossing_fn` automatically based on configuration flags.

```python
import jax.numpy as jnp

global_batch = model.global_batch_size
inputs = jnp.zeros((global_batch, fc.max_context), dtype=jnp.float32)
masks = jnp.zeros((global_batch, fc.max_context), dtype=jnp.bool_)

mean_forecast, full_forecast = model.compiled_decode(
    fc.max_horizon, inputs, masks
)

```

## Key Architectural Components

Both APIs rely on shared architectural primitives implemented across [`v1/src/timesfm/timesfm_base.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_base.py) and [`src/timesfm/timesfm_2p5/timesfm_2p5_base.py`](https://github.com/google-research/timesfm/blob/main/src/timesfm/timesfm_2p5/timesfm_2p5_base.py).

### Patch-Based Tokenization

TimesFM reshapes raw time-series into **input patches** of length `p` and **output patches** of length `o` before feeding them to the stacked transformer (`self.stacked_xf`). Ensure your `context_len` or `max_context` is divisible by `p` and your horizon is divisible by `o`; otherwise, the compilation step will adjust dimensions automatically and log the changes.

### Quantile Handling and Advanced Features

The Flax implementation exposes three optional post-processing utilities:
- **`_force_flip_invariance_fn`**: Enforces prediction consistency under time-series reversal.
- **`_use_continuous_quantile_head_fn`**: Activates continuous quantile output heads.
- **`_fix_quantile_crossing_fn`**: Corrects monotonicity violations in quantile predictions.

## Summary

- **Two APIs, one model**: Use `TimesFmJax` for pure JAX control or `TimesFM_2p5_200M_flax` for high-level Flax NNX abstractions.
- **Checkpoint compatibility**: JAX uses PaxML restoration while Flax uses Orbax; both support automatic Hugging Face hub downloads.
- **Compilation requirements**: Context length must be a multiple of input patch size `p`; forecast horizon must be a multiple of output patch size `o`.
- **Parallelism**: Both implementations use `pmap` to distribute inference across all available JAX devices.
- **Quantile outputs**: Both return mean forecasts and full quantile tensors; Flax adds optional continuous heads and crossing correction.

## Frequently Asked Questions

### What is the difference between TimesFmJax and TimesFM_2p5_200M_flax?

**`TimesFmJax`** provides a functional, low-level API in [`v1/src/timesfm/timesfm_jax.py`](https://github.com/google-research/timesfm/blob/main/v1/src/timesfm/timesfm_jax.py) that directly manages PaxML checkpoints and JIT compilation via `jit_decode`. **`TimesFM_2p5_200M_flax`** offers an object-oriented NNX interface with `from_pretrained` loading, encapsulated compilation via `model.compile()`, and built-in utilities for advanced quantile handling.

### How do I configure the forecast horizon and context length?

Set `context_len` (JAX) or `max_context` (Flax) and the horizon parameters as multiples of the model's patch sizes. The input patch length `p` and output patch length `o` are defined in the base classes. If you provide incompatible lengths, the Flax `compile` method will automatically round to valid multiples and issue a warning.

### Does TimesFM support multi-GPU inference with JAX?

Yes. Both implementations automatically detect `jax.local_device_count()` and use `pmap` to shard batches across devices. The `per_core_batch_size` parameter controls the batch size per accelerator, with the global batch size calculated as `per_core_batch_size * device_count`.

### What checkpoint formats are supported?

The JAX API accepts PaxML-format checkpoints via `TimesFmCheckpoint`, while the Flax API uses Orbax standards through `StandardCheckpointer`. Both support loading directly from the Hugging Face Hub using repository IDs like `google/timesfm-2.5-200m-jax` or `google/timesfm-2.5-200m-flax`.