Getting the statespace models running with numba is high on my priority list. I don’t like being shackled to jax either.