OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

• URL: http://arxiv.org/abs/2603.14802v1
• Date: Mon, 16 Mar 2026 03:59:46 GMT
• Title: OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX
• Abstract summary: OpenReservoirComputing (ORC) is a Python library for reservoir computing (RC) written in JAX.<n>JAX is a Python library for high-performance numerical computing.<n>ORC offers GPU acceleration, JIT compilation, and automatic vectorization.
• Score: 3.6137862091320936
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• CodeJeNN: A simple C++ neural network generator for physics applications [1.1599040452563385]
  Machine learning has shown speedups for numerical methods in physics applications, but integrating Python-based libraries into high-performance C++ solvers creates performance bottlenecks.<n>We present CodeJeNN, which bridges this gap by auto-generating self-contained C++ code from trained Keras models for inference.
  arXiv  Detail & Related papers  (2026-07-02T20:30:40Z)
• A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures [58.26804959656713]
  We present EB-JEPA, an open-source library for learning representations and world models using Joint-Embedding Predictive Architectures (JEPAs)<n>JEPAs learn to predict in representation space rather than pixel space, avoiding the pitfalls of generative modeling.<n>We show how these representations can drive action-conditioned world models, achieving a 97% planning success rate on the Two Rooms navigation task.
  arXiv  Detail & Related papers  (2026-02-03T14:56:24Z)
• AutoHete: An Automatic and Efficient Heterogeneous Training System for LLMs [68.99086112477565]
  Transformer-based large language models (LLMs) have demonstrated exceptional capabilities in sequence modeling and text generation.<n>Existing heterogeneous training methods significantly expand the scale of trainable models but introduce substantial communication overheads and CPU workloads.<n>We propose AutoHete, an automatic and efficient heterogeneous training system compatible with both single- GPU and multi- GPU environments.
  arXiv  Detail & Related papers  (2025-02-27T14:46:22Z)
• Adaptable Embeddings Network (AEN) [49.1574468325115]
  We introduce Adaptable Embeddings Networks (AEN), a novel dual-encoder architecture using Kernel Density Estimation (KDE) AEN allows for runtime adaptation of classification criteria without retraining and is non-autoregressive. The architecture's ability to preprocess and cache condition embeddings makes it ideal for edge computing applications and real-time monitoring systems.
  arXiv  Detail & Related papers  (2024-11-21T02:15:52Z)
• Scaling Tractable Probabilistic Circuits: A Systems Perspective [53.76194929291088]
  PyJuice is a general implementation design for PCs that improves prior art in several regards.<n>It is 1-2 orders of magnitude faster than existing systems at training large-scale PCs.<n>PyJuice consumes 2-5x less memory, which enables us to train larger models.
  arXiv  Detail & Related papers  (2024-06-02T14:57:00Z)
• Enabling High-Sparsity Foundational Llama Models with Efficient Pretraining and Deployment [56.44025052765861]
  Large language models (LLMs) have revolutionized Natural Language Processing (NLP), but their size creates computational bottlenecks. We introduce a novel approach to create accurate, sparse foundational versions of performant LLMs. We show a total speedup on CPUs for sparse-quantized LLaMA models of up to 8.6x.
  arXiv  Detail & Related papers  (2024-05-06T16:03:32Z)
• JORA: JAX Tensor-Parallel LoRA Library for Retrieval Augmented Fine-Tuning [16.86356520836045]
  We introduce a novel framework for PEFT-compatible fine-tuning of Llama-2 models, leveraging distributed training. Our framework uniquely utilizes JAX's just-in-time (JIT) compilation and tensor-sharding for efficient resource management. Our experiments show more than 12x improvement in runtime compared to Hugging Face/DeepSpeed implementation with four GPU while consuming less than half the VRAM per GPU.
  arXiv  Detail & Related papers  (2024-03-17T23:02:04Z)
• Detach-ROCKET: Sequential feature selection for time series classification with random convolutional kernels [0.7499722271664144]
  We introduce Sequential Feature Detachment (SFD) to identify and prune non-essential features in ROCKET-based models. SFD can produce models with better test accuracy using only 10% of the original features. We also present an end-to-end procedure for determining an optimal balance between the number of features and model accuracy.
  arXiv  Detail & Related papers  (2023-09-25T20:24:36Z)
• SynJax: Structured Probability Distributions for JAX [3.4447129363520337]
  SynJax provides efficient vectorized implementation of inference algorithms for structured distributions. We can build large-scale differentiable models that explicitly model structure in the data.
  arXiv  Detail & Related papers  (2023-08-07T04:20:38Z)
• RAF: Holistic Compilation for Deep Learning Model Training [17.956035630476173]
  In this paper, we present RAF, a deep learning compiler for training. Unlike existing DLCs, RAF accepts a forward model and in-house generates a training graph. RAF is able to systematically consolidate graph optimizations for performance, memory and distributed training.
  arXiv  Detail & Related papers  (2023-03-08T17:51:13Z)
• RcTorch: a PyTorch Reservoir Computing Package with Automated Hyper-Parameter Optimization [1.2891210250935146]
  Reservoir computers (RCs) are among the fastest to train of all neural networks. RC adoption has lagged because of the model's sensitivity to its hyper- parameters. We introduce RcTorch, a PyTorch based RC neural network package with automated HP tuning.
  arXiv  Detail & Related papers  (2022-07-12T22:24:36Z)
• Convolutional Tensor-Train LSTM for Spatio-temporal Learning [116.24172387469994]
  We propose a higher-order LSTM model that can efficiently learn long-term correlations in the video sequence. This is accomplished through a novel tensor train module that performs prediction by combining convolutional features across time. Our results achieve state-of-the-art performance-art in a wide range of applications and datasets.
  arXiv  Detail & Related papers  (2020-02-21T05:00:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.