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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Carey, Xan, Deshmukh, Yash, Huang, Aileen, Jadhav, Sunit, Tekawade, Omkar, Yang, Lorraine, Tiwary, Anvesha, Riano, Gerardo, Greenwald, Amy, Goktas, Denizalp
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915943401127936
author Carey, Xan
Deshmukh, Yash
Huang, Aileen
Jadhav, Sunit
Tekawade, Omkar
Yang, Lorraine
Tiwary, Anvesha
Riano, Gerardo
Greenwald, Amy
Goktas, Denizalp
author_facet Carey, Xan
Deshmukh, Yash
Huang, Aileen
Jadhav, Sunit
Tekawade, Omkar
Yang, Lorraine
Tiwary, Anvesha
Riano, Gerardo
Greenwald, Amy
Goktas, Denizalp
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2604_16719
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Chronax: A Jax Library for Univariate Statistical Forecasting and Conformal Inference
Carey, Xan
Deshmukh, Yash
Huang, Aileen
Jadhav, Sunit
Tekawade, Omkar
Yang, Lorraine
Tiwary, Anvesha
Riano, Gerardo
Greenwald, Amy
Goktas, Denizalp
Machine Learning
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.
title Chronax: A Jax Library for Univariate Statistical Forecasting and Conformal Inference
topic Machine Learning
url https://arxiv.org/abs/2604.16719