OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

2026-03-15

AI Summary

This work proposes OpenReservoirComputing, the first open-source reservoir computing framework built on JAX and Equinox, to address efficiency bottlenecks in time series forecasting, classification, and chaotic system modeling. By leveraging JAX’s automatic differentiation, just-in-time (JIT) compilation, and GPU/TPU acceleration capabilities—combined with Equinox’s composable neural network design—the framework enables end-to-end differentiable, highly modular, and scalable reservoir architectures. This integration significantly accelerates prototyping and large-scale reservoir training compared to existing approaches, demonstrating superior performance across a range of time series tasks.

Technology Category

• Machine Learning: Time-Series/Data Streams
• Computer Vision: Learning & Optimization for CV
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

OpenReservoirComputing (ORC) is a Python library for reservoir computing (RC) written in JAX (Bradbury et al. 2018) and Equinox (Kidger and Garcia 2021). JAX is a Python library for high-performance numerical computing that enables automatic differentiation, just-in-time (JIT) compilation, and GPU/TPU acceleration, while Equinox is a neural network framework for JAX. RC is a form of machine learning that functions by lifting a low-dimensional sequence or signal into a high-dimensional dynamical system and training a simple, linear readout layer from the high-dimensional dynamics back to a lower-dimensional quantity of interest. The most common application of RC is time-series forecasting, where the goal is to predict a signal's future evolution. RC has achieved state-of-the-art performance on this task, particularly when applied to chaotic dynamical systems. In addition, RC approaches can be adapted to perform classification and control tasks. ORC provides both modular components for building custom RC models and built-in models for forecasting, classification, and control. By building on JAX and Equinox, ORC offers GPU acceleration, JIT compilation, and automatic vectorization. These capabilities make prototyping new models faster and enable larger and more powerful reservoir architectures. End-to-end differentiability also enables seamless integration with other deep learning models built with Equinox.

Problem

• Reservoir Computing
• GPU acceleration
• JAX
• time-series forecasting
• differentiability

Innovation

• Reservoir Computing
• JAX
• GPU Acceleration
• End-to-End Differentiability
• Just-In-Time Compilation