Fix typo in Part1 tutorial notebook

[rootkiller6788]

Summary

Fixes a typo in the MomentumOptimizer example in docs/notebooks/Part1_Introduction.ipynb (issue #296).

In the update method, the MomentumOptState is constructed with:

momentums=jax.tree_util.tree_unflatten(struct, output_params),

but it should use output_momentums. As written, the returned optimizer state's momentums field is populated with the updated parameters rather than the updated momentum accumulators, so the momentum update shown in the tutorial is silently incorrect (the notebook's own output confirms this: momentums equals params).

Change

One-line fix:

-        momentums=jax.tree_util.tree_unflatten(struct, output_params),
+        momentums=jax.tree_util.tree_unflatten(struct, output_momentums),

Testing

Verified the notebook still parses as valid JSON (nbformat 4, 40 cells). No runtime environment changes; the fix only corrects which accumulator list is passed to tree_unflatten for the momentums field.