Fix compatibility with flax-basic_example.ipynb after JAX update#6247
Merged
JanuszL merged 3 commits intoNVIDIA:mainfrom Mar 10, 2026
Merged
Fix compatibility with flax-basic_example.ipynb after JAX update#6247JanuszL merged 3 commits intoNVIDIA:mainfrom
JanuszL merged 3 commits intoNVIDIA:mainfrom