in digits_vae.ipynb, error: no attribute ModelAndOptimizer

[joshtburdick]

When I open docs/source/digits_vae.ipynb in Google Colab, I get errors in cell [15] (by the original numbering) :

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipython-input-3209941142.py in <cell line: 0>()
      8 )
      9 
---> 10 optimizer = nnx.ModelAndOptimizer(model, optax.adam(1e-3))
     11 
     12 @nnx.jit

AttributeError: module 'flax.nnx' has no attribute 'ModelAndOptimizer'

Presumably this isn't the error you wanted to demo debugging :)

Gemini suggests replacing ModelAndOptimizer with Optimizer. When I do that throughout the notebook, it runs.