OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

Jan Williams; Dima Tretiak; Steven L. Brunton; J. Nathan Kutz; Krithika Manohar

cs.LGMar 16, 2026

OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

Authors: Jan Williams, Dima Tretiak, Steven L. Brunton, J. Nathan Kutz, Krithika Manohar

Organizations: Department of Mechanical Engineering, University of Washington, USA · Department of Applied Mathematics, University of Washington, USA · Department of Electrical and Computer Engineering, University of Washington, USA · Autodesk Research, London, UK

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.

Figures & tables

	Lang.	GPU	Auto. Diff.	Parallel./Vect.	Forecast.	Classif.	Control	Cont. Time
ORC	Python	✓	✓	✓/✓	✓	✓	✓	✓
ReservoirPy	Python	✓*	$\times$	$\times$ / $\times$	✓	✓	$\times$	$\times$
ReservoirComputing.jl	Julia	✓	✓	✓/ $\times$	✓	✓	$\times$	$\times$

Table 1: Comparison of reservoir computing libraries across key features. ✓ indicates full support;

\times

indicates no support. *ReservoirPy’s GPU support is available via its JAX backend (v0.4.0+) but does not fully exploit JAX’s functional programming model. Parallelizable denotes native support for parallel RC architectures as in (Pathak et al. 2018) and vectorizable denotes native support for vectorization (e.g. vmap ).

Figure 1: ORC three-layer pipeline architecture. Each reservoir computer (RC) is decomposed into (i) an “embedding” function that lifts a low-dimensional signal

u_{t}

to the reservoir dimension, (ii) a “driver” function that propagates the reservoir state, and (iii) a “readout” that maps back to a target

y

. For control and forecasting RCs, the target

y

is typically

u_{t+1}

(either in the presence of a forcing term or not) and for classification the target

y

is a label.

Figure 2: ORC with GPU acceleration enables significantly faster performance than ReservoirPy, even when using ReservoirPy’s JAX backend. Panel (a) shows the time per forecast step of RC models with varying reservoir dimension trained to forecast the Lorenz system, while panel (b) shows training tine for RC models with fixed reservoir dimension of 2000 but varying number of training samples. Performance of the two libraries with and without GPU acceleration are shown. GPU results were obtained running on an NVIDIA A40 GPU and CPU results were obtained with an Apple M2 chip.

Reservoir Computing

Reservoir computing uses recurrent dynamical systems to transform temporal inputs into features that support efficient learning with simple readouts.

See all 107 papers

Kolmogorov-Arnold Reservoir Computing
Jun 18, 2026Juntian Huang, Jürgen Kurths, Ying TangDynamical Systems Time Series Forecasting
Institute of Fundamental and Frontier Sciences, University of Electronic Science and Technology of China, Chengdu 611731, China · School of Physics, University of Electronic Science and Technology of China, Chengdu 611731, China · Key Laboratory of Quantum Physics and Photonic Quantum Information, Ministry of Education, University of Electronic Science and Technology of China, Chengdu 611731, China +4
Frequency-Based Reservoir computing
Jul 27, 2026Arthur S PowanweFrequency-Domain Feature Learning Reservoir Computing
Department of Mathematics, Western University, London, ON, Canada
Chronax: A Jax Library for Forecasting and Conformal Inference
Apr 17, 2026Xan Carey, Arsh Singh, Yash Deshmukh +8Time Series Forecasting
Simulacrum New York City, NY, USA

Jun 18, 2026J/K move · Enter open · S save

Kolmogorov-Arnold Reservoir Computing

Juntian Huang, Jürgen Kurths, Ying Tang

Institute of Fundamental and Frontier Sciences, University of Electronic Science and Technology of China, Chengdu 611731, China · School of Physics, University of Electronic Science and Technology of China, Chengdu 611731, China · Key Laboratory of Quantum Physics and Photonic Quantum Information, Ministry of Education, University of Electronic Science and Technology of China, Chengdu 611731, China+4

Dynamical Systems Time Series Forecasting Reservoir Computing Kolmogorov-Arnold Networks

Open paper AskPDF arXiv

Reservoir computing offers a lightweight framework for forecasting dynamical systems but may struggle to capture long-range dependencies due to limited representational capacity. Conventional reservoir computing recurrently uses fixed reservoirs with hyperparameter sensitivity, while the next generation reservoir computing removes recurrence at the cost of rapidly growing feature dimensions. Here, we develop Kolmogorov-Arnold Reservoir Computing (KARC), which replaces reservoirs with explicit basis-function expansions inspired by the Kolmogorov-Arnold representation theorem. We rigorously show that KARC is a lightweight design of Kolmogorov-Arnold networks (KANs), preserving the potential expressive capacity of KANs while admitting efficient closed-form training of reservoir computing. At comparable cost, KARC outperforms existing reservoir computing methods on challenging benchmarks including partial differential equations. It can also be integrated with generative diffusion models for facilitating text-to-image generation. This work thus establishes a principled bridge between reservoir computing and KANs, yielding a unified framework for efficient dynamical forecasting and generative modeling.

Abstract

Reservoir Computing

Explore similar work

Kolmogorov-Arnold Reservoir Computing

Frequency-Based Reservoir computing

Chronax: A Jax Library for Forecasting and Conformal Inference