[State-space model; dependent on #100] Replace autograd with Jax

[adam2392]

Autograd is not actively developed anymore and is superseded by jax. In #100, @jadrew43 will implement things first with autograd since he's comfortable w/ that already. This GH issue will track when #100 is done, to refactor the code to replace all autograd functions w/ jax.

Hey @adam2392 one of the CI errors is due to autograd: ModuleNotFoundError: No module named 'autograd'. I'd really like to move to the next step of the proposal and work to integrate jax (to replace autograd) later on in order to keep pace with the summer milestones. Is it alright if I install autograd in the dependencies for mnedev (assuming that will fix the error) for now and work to integrate jax later on?

Originally posted by @jadrew43 in #100 (comment)