flax

Update flax/docs_nnx/guides/performance.md and .ipynb

#4919

Merged