Expanding the JAX SSM Library & Exploring Particle-Based Variational Inference (JAX ML Textbook)

For this project, I plan to explore state space models(SSM) and variational inference(VI) in depth. First, I will translate legacy numpy/PyTorch codes into JAX/Flax and contribute additional codes and figures of pedagogical utility to the relevant chapters in the PRML textbook. Then, I will expand the JAX SSM Library(JSL), with a primary focus on implementing Gaussian smoothers and particle smoothers. Last, I will explore some of the state-of-the-art particle-based variational inference models and implement them to add to the pyprobml repository.