timestral-mlx

Timestral time-series forecasting on Apple silicon, ported to MLX.

This repo holds no weights. It is the home page for the code, which loads Mistral AI's released checkpoints straight from their own repos. The network runs on the Apple GPU in float32, and the float64 normalization statistics stay on MLX's CPU stream.

• Code: https://github.com/vikiival/timestral-mlx
• PyPI: https://pypi.org/project/timestral-mlx/

Use

pip install timestral-mlx

import numpy as np
from timestral_mlx import TimestralForecaster

forecaster = TimestralForecaster.from_pretrained("mistralai/Timestral-209M")

context = np.random.randn(2, 1, 2000).cumsum(-1)   # (series, variates, time); NaN = missing
forecast = forecaster.forecast(context, prediction_length=96)

forecast.median            # (2, 1, 96)
forecast.quantile(0.9)     # (2, 1, 96)

Parity with the PyTorch reference

Largest disagreement divided by the largest forecast value, on the released checkpoints:

checkpoint	2000 → 96	1500 → 500	2000 → 2500 (3 blocks)	3 variates + covariate, 800 → 64
Timestral-5M	5.4e-6	1.1e-6	9.7e-5	3.4e-6
Timestral-209M	9.4e-7	1.1e-6	4.6e-6	6.4e-7

On an M3 Pro the port is about 2.3 to 4.7 times faster than the reference on the CPU. Details are in the README.

Weights and license

The code is Apache 2.0. The weights are Mistral AI's, under their Modified MIT license, which grants no rights to anyone whose company or employer has more than $20 million in monthly revenue.