Using variable declared at a broader scope in a function is bad form
afairley opened this issue · comments
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.
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.