We have officially published linax, a library for state space models written in Jax. You can find it here.