flax

afbc5021 - Merge pull request #4149 from google:nnx-fix-nnx-basics-prints