jbilcke/timestral-mlx

Timestral MLX

Run all four Mistral Timestral models locally on Apple silicon with MLX. This independent community port includes the forecasting pipeline, FP32 weights, pinned provenance and numerical validation reports. No training, fine-tuning or quantization is performed.

Variant	Weights	Layers	Width	Local bundle
Timestral-5M	20.84 MB	8	192	5M/
Timestral-28M	111.30 MB	12	384	28M/
Timestral-123M	490.45 MB	20	640	123M/
Timestral-209M	837.27 MB	24	768	209M/

Exact tensor counts, byte sizes, source revisions and file checksums are recorded in models.json and each bundle's manifest.json. File size does not measure peak runtime memory.

Download and install

Use an Apple silicon Mac with a visible Metal GPU and Python 3.10 or later. The tested environment is Python 3.12, MLX 0.32.4 and NumPy 2.5.4; other combinations have not been validated. This Python MLX release targets macOS. It does not include a Core ML, Neural Engine, Swift or iOS runtime.

Install the download tool, then explicitly download the code and your chosen model. Replace 5M with 28M, 123M or 209M for another size. Pin revision to a trusted commit for reproducible downloads.

python3 -m pip install huggingface_hub

from huggingface_hub import snapshot_download
snapshot_download(
    "jbilcke/timestral-mlx",
    local_dir="./timestral-mlx",
    allow_patterns=[
        "README.md", "pyproject.toml", "LICENSE*", "NOTICE", "requirements.txt",
        "timestral_mlx/*.py", "examples/*.py", "models.json", "SHA256SUMS", "5M/*",
    ],
    # revision="<trusted commit SHA>",
)

cd timestral-mlx
python3 -m venv .venv
. .venv/bin/activate
python -m pip install .
python examples/forecast.py ./5M --horizon 96

Create virtual environments on a normal macOS filesystem. Drives with large allocation units or no symlink support can consume excessive space or prevent environment creation. Checkpoints themselves can live on an external drive.

Forecast

import numpy as np
from timestral_mlx import TimestralForecaster

forecaster = TimestralForecaster.from_pretrained("./5M")
context = np.sin(np.arange(512, dtype=np.float32) * 2 * np.pi / 24)
forecast = forecaster.forecast(context, prediction_length=96)

print(forecast.quantiles.shape)  # (1, 1, 96, 21)
print(forecast.median.shape)     # (1, 1, 96)
print(forecast.quantile(0.9))

A local directory load and subsequent inference make no network requests. After installing the runtime, an explicit Hub load is also available:

forecaster = TimestralForecaster.from_pretrained(
    "jbilcke/timestral-mlx", variant="28M",
    # revision="<trusted commit SHA>",
    # cache_dir="/path/to/a/symlink-capable/cache",
    # local_files_only=True,
)

Local loads verify the config, weights and original weight licence against the bundle manifest, then load all MLX parameters strictly. A self-described checksum detects corruption; pin a trusted repository revision or manifest hash for authenticity.

Inputs, outputs and covariates

Inputs are NumPy arrays or MLX arrays convertible to NumPy. context accepts (T,), (B,T) or (B,V,T), where B counts independent series and V counts jointly forecast variates. NaN marks missing values; observation_mask can mark additional missing steps. Empty inputs are rejected. Context exceeding 32,768 steps is truncated to its most recent steps. The pipeline left-pads incomplete 32-step patches.

The result holds NumPy float32 quantiles (B,V,H,21) at probabilities 0.01, 0.05, 0.10, 0.15, ..., 0.90, 0.95, 0.99. Up to 1,024 steps are decoded per block. Longer horizons feed each block's median back into context. Symmetric averaging and clipping for nonnegative observed context are enabled by default, matching upstream. Disable them with symmetric_averaging=False and clip_nonnegative=False.

forecast = forecaster.forecast(
    context, 96,
    past_covariates=past,     # (B,C_past,T)
    known_covariates=known,   # (B,C_known,T+96)
    categorical_past_covariates=past_codes,
    categorical_known_covariates=known_codes,
)

Categorical covariates use numeric category codes, with NaN for missing values. Codes remain float64 during target encoding so distinct large codes stay distinct. For a univariate target they use the upstream smoothed target encoding; multivariate targets receive raw codes. Numeric and categorical covariates may be combined. Covariate rows remain exogenous through the upstream time and variate attention masks. Known 0/1 indicators retain the upstream normalization rule. Infinite covariate values are rejected.

Execution and fidelity

Patch embedding, gated time/variate attention, partial RoPE, SwiGLU, RMSNorm and the monotone quantile head execute in FP32 with MLX on the GPU. CPU NumPy handles packing, masks, target encoding, normalization and refinement. Causal and advancing-frame Welford statistics retain float64, which MLX Metal does not support. PyTorch is used only by the separate reference-validation tools; runtime inference does not import or require it. This pipeline is specific to Timestral and is not an mlx_lm.generate or Transformers pipeline implementation.

Conversion preserves every original FP32 tensor value and layout, verified after serialization. Only parameter names containing ._gate. change to .gate. because MLX excludes private attributes from parameter registration. Every original source revision and original checkpoint hash is retained.

The supplied reports compare all four checkpoints against pinned Mistral PyTorch CPU FP32 inference on deterministic synthetic fixtures, including multivariate batches, missing values, numeric and categorical covariates, 1,057-step block feedback, constants and all-missing contexts. Acceptance requires finite ordered quantiles, normalized RMSE at most 1e-4, and normalized maximum error at most 2e-3. The synthetic checks establish implementation parity, not real-world forecast accuracy. No GIFT-Eval, TIME or fev-bench benchmark was rerun. Timings in the reports include first-call overhead and are not controlled performance measurements. The trained full 32,768-step context is supported by the configuration; all four sizes passed the full-context smoke checks, including exact truncation from 32,785 to 32,768 steps and inference without importing PyTorch.

Reproduce

git clone https://github.com/mistralai/timestral.git upstream
git -C upstream checkout a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff
python -m pip install -r requirements-validation.txt
python scripts/download.py 5M 28M 123M 209M
python scripts/convert.py 5M 28M 123M 209M
python -m pytest -q
python scripts/validate.py 5M 28M 123M 209M
python scripts/smoke.py
python scripts/build_wheel.py

Downloads are explicit, resumable, and verify each upstream Git/LFS blob. Conversion proceeds smallest first. scripts/publish.py assembles an explicit inventory, rejects failed/stale validation, excludes AppleDouble ._*, .DS_Store, caches and temporary files, and verifies public Hub hashes after upload. The repository includes source code and a wheel under dist/ for installation without a source checkout.

Licences and attribution

Model weights: the original Modified MIT License, attribution 2026 - Mistral AI, is preserved verbatim in LICENSE-WEIGHTS and in each model folder's LICENSE. It includes a restriction for companies with global consolidated monthly revenue above US$20 million in the preceding month; consult the original licence for its full terms. The converted weights remain subject to that licence.

Source code: Apache License 2.0, adapted from mistralai/timestral, with the upstream licence in LICENSE and modification attribution in NOTICE. The code licence does not relicense the weights. Each bundle also retains its original upstream model card as upstream-model-card.md.

This is an independent community conversion and is not an official Mistral release.