camail-official/linax

linax is a light weight collection of state space models implemented in JAX ⚡️