Change nnx ModelAndOptimizer to Optimizer

[Gerixmus]

Flax optimizer is currently under nnx.Optimizer while all the docs and notebooks reference nnx.ModelAndOptimizer. This leads to this error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[/tmp/ipython-input-3566937807.py](https://localhost:8080/#) in <cell line: 0>()
      2 import optax
      3 
----> 4 optimizer = nnx.ModelAndOptimizer(model, optax.sgd(learning_rate=0.05))
      5 
      6 def loss_fun(model: nnx.Module, data: jax.Array, labels: jax.Array):

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

This PR replaces all occurrences of nnx.ModelAndOptimizer with nnx.Optimizer in docs and notebooks.

[Gerixmus]

Import flax in Google Colab will install flax 0.10.6. This version doesn't require wrt but from version 0.11+ the parameter is required. Now because I tried changing the version in requirements I got this in Read the Docs:

The conflict is caused by:
--
  | The user requested flax==0.10.6
  | jax-ai-stack 2025.9.3 depends on flax==0.12.0

I guess changing any file in docs/source now would break Read the Docs build due to flax v0.11 Optimizer updates.