Create a Jax version of the ESN

[eliseoe]

Issue:

I think it is worth exploring the possibility of creating a JAX version of our current implementation to leverage the benefits and optimizations of JAX to speed up our ESN training. In my preliminary tests, the jitted version of step speeds up the computation when using a reservoir of 100. I believe the benefits will increase for larger reservoirs

Tasks:

• Implement a JAX version of the standard ESN
• Adapt the validation strategy to JAX

Additional Context:

Precisely, I believe a mix of scipy ridge regression and JAX will lead us to the fastest version. While the change to JAX in the standard ESN is mostly straightforward, there is more adjusting needed in validation.py

[eliseoe]

This task was preliminarily completed. A first Jax esn version is on the main branch, with the validation adapted.

Memory issue still exists (see Issue: Resolve gpu memory issue with JAX esn).

Closing this issue.