Chronax: A Jax Library for Univariate Statistical Forecasting and Conformal Inference

2026-04-17

AI Summary

This work proposes the first purely functional time series forecasting framework built on JAX, addressing the limitations of traditional libraries that rely on interpreted execution and object-oriented design, which hinder efficient parallelization and accelerator integration for large-scale heterogeneous data. By unifying preprocessing, modeling, and multi-step prediction into composable, end-to-end differentiable pure functions, the framework systematically integrates modern computational paradigms—including functional programming, automatic vectorization, and just-in-time compilation. This approach enables model-agnostic conformal uncertainty quantification, scalable multi-series forecasting, and seamless interoperability with scientific computing and machine learning ecosystems. The resulting system achieves substantial gains in predictive efficiency and scalability while natively supporting GPU/TPU acceleration.

Technology Category

• Machine Learning: Time-Series/Data Streams
• Reasoning under Uncertainty: Probabilistic Programming
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Time-series forecasting is central to many scientific and industrial domains, such as energy systems, climate modeling, finance, and retail. While forecasting methods have evolved from classical statistical models to automated, and neural approaches, the surrounding software ecosystem remains anchored to the traditional Python numerical stack. Existing libraries rely on interpreter-driven execution and object-oriented abstractions, limiting composability, large-scale parallelism, and integration with modern differentiable and accelerator-oriented workflows. Meanwhile, today's forecasting increasingly involves large collections of heterogeneous time series data, irregular covariates, and frequent retraining, placing new demands on scalability and execution efficiency. JAX offers an alternative paradigm to traditional stateful numerical computation frameworks based on pure functions and program transformations such as just-in-time compilation and automatic vectorization, enabling end-to-end optimization across CPUs, GPUs, and TPUs. However, this modern paradigm has not yet been fully incorporated into the design of forecasting systems. We introduce Chronax, a JAX-native time-series forecasting library that rethinks forecasting abstractions around functional purity, composable transformations, and accelerator-ready execution. By representing preprocessing, modeling, and multi-horizon prediction as pure JAX functions, Chronax enables scalable multi-series forecasting, model-agnostic conformal uncertainty quantification, and seamless integration with modern machine learning and scientific computing pipelines.

Problem

• time-series forecasting
• software ecosystem
• scalability
• execution efficiency
• accelerator integration

Innovation

• JAX
• time-series forecasting
• functional programming
• conformal inference
• accelerator-ready execution