Reservoir Computing with JAX

[Wotaker]

Hi, the Resrviour Computing (RC) is an especially interesting architecture of a Recurent Neural Networks (RNN). It is a convinient tool for the analysys of dynamic systems. I am currently working on my Master thesis studying brain effective connectivity and I use the RC during the studies. Unfortunately none of the autograd libraries (TF, PyTorch nor JAX ecosysem) officially supports this awesome computing approach.

My idea was to implement the Echo State Networks (ESN), the most basic of the architectures in RC, using JAX and to benefit both from autograd and JIT acceleration. But I do not have enough time nor skillset to do it on my own, so I plant this idea here, so you can maybe integrate it in your ecosystem.

For now, I am using the PyRCN library with a scikit-learn like interface, but it can not be accelerated using GPU.

-> An awesome survey on reservoir computing
-> Maybe it is a relevant issue for the Haiku library, so they could add a relevant module.