flax

6e63497f - Merge pull request #3942 from google:fix-nnx-basics-notebook