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.

flax/docs/guides/flax_fundamentals/state_params.rst

Line 59 in 514c111

updates, opt_state = tx.update(grads, opt_state) # Defined below.

[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.