implement deep kalman filter

[murphyk]

Convert the pyro demo, which uses pytorch under the hood, to numpyro, which uses jax under the hood. Use flax for the neural network part. Then reimplement in pure jax/flax, without relying on numpyro black magic.