For some reason these imports are elided on read the docs

I mean, they are inferable by the reader, but it looks like they are intended to be visible in the finished product and getting eaten somewhere along the publication tool chain?

I would also say that this tutorial reads a little bit bouncier between levels of technical explanation than one might hope. I think the call out to the reader to supply their own loss function is more of a math textbook kind of exercise left to the reader than something that belongs in explanatory documentation. The bit about BatchNormalization being dependent on batch size and therefore not properly vectorizable also seems like a theoretical concern that should be left out or handled differently.
https://github.com/google/flax/blob/514c11199152fab44de6002983f44e6b49aa0622/docs/guides/flax_fundamentals/state_params.rst?plain=1#L12