Using variable declared at a broader scope in a function is bad form

In person I would argue vociferously that this is kind of a major sin . tx should either be passed in as an argument or both tx and update_step should live on some object of which tx should be a variable. Maybe the neural network instance, that way all of the code in this tutorial can be typed into a single notebook without the later declaration of update_step shadowing the first one.
https://github.com/google/flax/blob/514c11199152fab44de6002983f44e6b49aa0622/docs/guides/flax_fundamentals/state_params.rst?plain=1#L59