OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

TL;DR

OpenReservoirComputing (ORC) is a GPU-accelerated Python library in JAX for reservoir computing, enabling fast prototyping and scalable models for time-series forecasting, classification, and control tasks.

Contribution

ORC introduces a modular, GPU-accelerated reservoir computing library in JAX that simplifies building, training, and integrating reservoir models with deep learning frameworks.

Findings

Enables high-performance reservoir computing with GPU acceleration.

Supports flexible, modular reservoir model construction.

Facilitates larger, more complex reservoir architectures.

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.