GitHub repository: romainligneul/dynamax
Description: State Space Models library in JAX