Lessons learned from building an agent to convert the R Forecast Package into a JAX api
Abstract
Lessons learned from building an agent to convert the R Forecast Package into a JAX api
Skander Hannachi, Jasmeet Bhatia, Dennis Kashkin, Anna Novakovska - Applied AI Engineering, Google Cloud
{shannachi@ jasmeetbhatia@ kashkin@ anovakovska@}google.com
Accepted at ISF 2026
Abstract:
The Forecast package in R, is one of the most popular and established frameworks for learning and working with statistical (local) time series models. However, over the last decade or so, the language of choice for data analysis and mathematical modeling has become Python, especially, since, unlike R, the latter provides multiple options for running models on dedicated hardware accelerators (TPUs/GPUs), and on distributed compute infrastructure.
In this study, we implement an LLM agent based code conversion pipeline that automatically converts the code from the Forecast package (R, C++), into JAX/PAX, a more recent modeling framework dedicated specifically for running compute heavy modeling tasks on TPUs and GPUs. We run our specialized code conversion agent against three core models of the Forecast package: TBATs, auto.arima(), and ETS(), with an aim to provide the exact same user experience as the original R API, but with the underlying model selection, model fitting, and forecast generation running all using JAX and PAX operations, running in a Python environment.
We report the results of our experiments and discuss the challenges we observed while running it. We compile these into a skills markdown file, which can be used by other agents intended to perform similar experiments. We provide the JAX based implementations, along with the skills file in an accompanying open source repo.
This is not the first attempt at converting the Forecast package into Python. Those other efforts however are significantly labor intensive, especially when it comes to ensuring parity with the source modeling APIs and quality control in general. Moreover, such projects rely heavily on the long term commitment of both community members and institutional contributors to the conversion effort.
The purpose of our effort is to show how this agent based process can be applied to automate any data science package upgrade or language conversion process with minimal contributors required outside of the core package maintainer team. Especially since the concept of agent skills files makes the process inherently self-improving, both within the scope of a single conversion effort, as well as across multiple long term conversion efforts. For example, the same approach can be applied to a future effort for upgrading the Forecast package to work with Julia, an even more recent and promising modeling language, while benefitting from the R-2JAX lessons learned.
Skander Hannachi, Jasmeet Bhatia, Dennis Kashkin, Anna Novakovska - Applied AI Engineering, Google Cloud
{shannachi@ jasmeetbhatia@ kashkin@ anovakovska@}google.com
Accepted at ISF 2026
Abstract:
The Forecast package in R, is one of the most popular and established frameworks for learning and working with statistical (local) time series models. However, over the last decade or so, the language of choice for data analysis and mathematical modeling has become Python, especially, since, unlike R, the latter provides multiple options for running models on dedicated hardware accelerators (TPUs/GPUs), and on distributed compute infrastructure.
In this study, we implement an LLM agent based code conversion pipeline that automatically converts the code from the Forecast package (R, C++), into JAX/PAX, a more recent modeling framework dedicated specifically for running compute heavy modeling tasks on TPUs and GPUs. We run our specialized code conversion agent against three core models of the Forecast package: TBATs, auto.arima(), and ETS(), with an aim to provide the exact same user experience as the original R API, but with the underlying model selection, model fitting, and forecast generation running all using JAX and PAX operations, running in a Python environment.
We report the results of our experiments and discuss the challenges we observed while running it. We compile these into a skills markdown file, which can be used by other agents intended to perform similar experiments. We provide the JAX based implementations, along with the skills file in an accompanying open source repo.
This is not the first attempt at converting the Forecast package into Python. Those other efforts however are significantly labor intensive, especially when it comes to ensuring parity with the source modeling APIs and quality control in general. Moreover, such projects rely heavily on the long term commitment of both community members and institutional contributors to the conversion effort.
The purpose of our effort is to show how this agent based process can be applied to automate any data science package upgrade or language conversion process with minimal contributors required outside of the core package maintainer team. Especially since the concept of agent skills files makes the process inherently self-improving, both within the scope of a single conversion effort, as well as across multiple long term conversion efforts. For example, the same approach can be applied to a future effort for upgrading the Forecast package to work with Julia, an even more recent and promising modeling language, while benefitting from the R-2JAX lessons learned.
