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

[afairley]

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

[afairley]

Looking at https://github.com/google/flax/blob/main/docs/guides/flax_fundamentals/flax_basics.md way down the bottom, it looks like tx is usually(occasionally?) a parameter for update_step declarations. Anyways this is a bit of a nitpick, but I'm definitely not the only person strongly sensitized to this nit.