flax

9bf7283c - Correct state variable name

Correct state variable name This pr corrects the ``update_step()`` function in the Managing Parameters and State notebook, to pass the ``updated_state``, rather than passing the ``state``, as it is.