Optimizing Chronax Performance using JAX JIT Compilation
Chronax models leverage JAX's Just-In-Time (JIT) compilation via XLA (Accelerated Linear Algebra) to achieve high performance, especially on GPUs. This guide explains the Chronax execution model, focusing on the difference between the initial "Cold Start" and subsequent "Warm Start" runs.
Prerequisites
- Chronax installed with JAX backend (ideally with GPU support,
jax[cuda]). - Data must be provided as
jnp.ndarray(JAX NumPy array). - All input arrays must be of type
float32.
import jax.numpy as jnp
import time
from chronax.models import AutoARIMA # Example model
Steps
1. Understand the Cold Start (JIT Compilation)
The very first time you call a core model function, such as model.fit() or model.forecast(), JAX triggers the JIT compilation process. This compiles the entire computation graph (the model logic) into highly optimized machine code (XLA).
This initial compilation step is known as the "Cold Start" and is typically the slowest part of the execution. The time taken here measures both the compilation overhead and the initial training time.
# Create synthetic data
y = jnp.arange(100, dtype=jnp.float32)
model = AutoARIMA()
# Phase 1: Cold Start (Compilation + Training)
# This call triggers JAX JIT compilation.
print("Starting Cold Start...")
start_time = time.time()
model.fit(y[:80])
cold_start_time = time.time() - start_time
print(f"Cold Start Time (Compilation + Fit): {cold_start_time:.4f}s")
2. Measure the Warm Start (Inference)
Once the model has been compiled, JAX caches the compiled graph. Subsequent calls to the same function signature (e.g., calling model.predict(h=20) multiple times) reuse this cached graph.
This reuse eliminates the compilation overhead, resulting in extremely fast execution times, known as the "Warm Start." This measures the pure inference speed of the optimized model.
# Phase 2: Warm Start (Inference)
# Subsequent calls reuse the compiled graph, resulting in high speed.
print("Starting Warm Start...")
start_time = time.time()
out = model.predict(h=20)
warm_start_time = time.time() - start_time
print(f"Warm Start Time (Inference): {warm_start_time:.4f}s")
# Access the forecast mean
print(f"Forecast mean shape: {out['mean'].shape}")
3. Utilize GPU Acceleration
Chronax models automatically leverage GPU acceleration if JAX is configured correctly and a compatible device is available. JAX's XLA backend is designed to maximize performance on accelerators (GPUs/TPUs).
When running Chronax, ensure your environment is set up for GPU usage. The performance gains from the Warm Start phase are most pronounced when running on accelerated hardware.
4. Account for Prediction Interval Overhead
If you request prediction intervals (e.g., 95% confidence level), the model must perform additional computations (often involving sampling or bootstrapping) beyond calculating the mean forecast.
While the core model remains compiled, generating these intervals adds measurable overhead compared to a simple mean forecast. Always benchmark interval generation separately if speed is critical.
# Phase 3: Interval Overhead
# Requesting intervals adds computation time.
print("Starting Interval Run...")
start_time = time.time()
out_interval = model.predict(h=20, level=95)
interval_time = time.time() - start_time
print(f"Interval Run Time: {interval_time:.4f}s")
# Access the prediction interval bounds
print(f"Lower bound shape: {out_interval['lo-95'].shape}")
Full example
This example demonstrates the typical performance profile of a Chronax model: a slow initial call followed by rapid subsequent execution.
import jax.numpy as jnp
import time
from chronax.models import AutoARIMA
# 1. Setup Data (must be float32)
y = jnp.arange(100, dtype=jnp.float32)
model = AutoARIMA()
# --- COLD START ---
# The first call triggers JAX JIT compilation.
print("--- Cold Start (Compilation + Fit) ---")
start_time = time.time()
model.fit(y[:80])
cold_start_time = time.time() - start_time
print(f"Time: {cold_start_time:.4f}s\n")
# --- WARM START (Inference) ---
# Subsequent calls reuse the compiled graph.
print("--- Warm Start (Inference Only) ---")
start_time = time.time()
mean_forecast = model.predict(h=20)
warm_start_time = time.time() - start_time
print(f"Time: {warm_start_time:.4f}s")
print(f"Speedup factor: {cold_start_time / warm_start_time:.1f}x\n")
# --- INTERVAL OVERHEAD ---
# Requesting intervals adds computation.
print("--- Interval Run (Inference + Intervals) ---")
start_time = time.time()
interval_forecast = model.predict(h=20, level=90)
interval_time = time.time() - start_time
print(f"Time: {interval_time:.4f}s")
print(f"Interval Overhead vs Warm: {(interval_time / warm_start_time - 1) * 100:.1f}%\n")
# Verify results
print(f"Forecast mean: {mean_forecast['mean'][:5]}")
print(f"90% Lower bound: {interval_forecast['lo-90'][:5]}")
Next steps
- Consult the JAX documentation for advanced control over JIT compilation and device placement.
- Learn how to use the
chronax.models.AutoARIMAfor real-world data. - Explore the
jnp.ndarraydocumentation for efficient data handling. - See the guide on "Adding Prediction Intervals" for more detail on the
levelparameter.