search

probml / dynamax

State Space Models library in JAX
https://probml.github.io/dynamax/
MIT License
675 stars 78 forks source link

Implement optax-like optimizer for CMGF update step #230

Closed petergchang closed 2 years ago

petergchang commented 2 years ago

Implement optax-like optimizer for CMGF update step

murphyk commented 2 years ago

Isn't this now resolved by https://github.com/probml/ssm-jax/blob/main/ssm_jax/cond_moments_gaussian_filter/optimizer.py