Implement EKF-MLP

[petergchang]

Description

• Reimplemented JSL Demo using SSM-JAX EKF code.
• Reorganized portions of code and added comments explaining how EKF is used for readability.
• Omitted saving figures for consistency with rest of the API.

Issue

#32

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[murphyk]

This looks good. Some small requests:

• I would prefer if each plot is its own figure, so we can save them separately, and then in latex, i can decide which ones to use. (Probably 10, 20, 30, and final=?? steps).
• use same vertical scale for all plots (eg -30..30)
• also create ekf_mlp.ipynb that imports this code, runs it, and shows the plots (cf https://github.com/probml/ssm-jax/blob/main/ssm_jax/lgssm/demos/kf_linreg.ipynb), so that i can point book readers to this colab.

[petergchang]

@murphyk
Thanks for your review, I submitted new commits that address your requests.
I can't import the demo module before the PR gets merged, so the Jupyter notebook does not work yet,
but it should work after the merge.