Reservoir Computing with Flax

[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 Flax 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 awesome library.

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

-> A helpful survey on reservoir computing
-> I have left the same idea in the JAX's discussion page

[Wotaker]

I have found the support for reserviour computing in the discussion page, but we still lack a decent implementation in JAX.
Here is a link for both threads:

• Implementation attempt by @GJBoth
• Unanswered thread relevant to the topic