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