lindermanlab / S5

MIT License
248 stars 43 forks source link

JAX implementation #3

Closed EelcoHoogendoorn closed 1 year ago

EelcoHoogendoorn commented 1 year ago

Hi, congrats on such an elegant and no doubt very high impact piece of work. And the JAX code in your paper is much appreciated! However, I was curious if you have plans for a more end-to-end JAX implementation on github. Right now im looking at porting the normal-hippo init to JAX for instance; seems doable, but the easier to reproduce the better I suppose.

EelcoHoogendoorn commented 1 year ago

Never mind; got too many tabs open I suppose; I see its already there!