improve scaling up guide#5354

cgarciae wants to merge 2 commits intonnx-basics-treegoogle/flax:nnx-basics-treefrom

improve-scaling-up-guidegoogle/flax:improve-scaling-up-guideCopy head branch name to clipboard