Hotfix for Maxtext regression with JAX 0.9 changes#629

ipanfilo wants to merge 2 commits into

devROCm/TransformerEngine:devfrom

ipanfilo/jax09_maxtext_hotfixROCm/TransformerEngine:ipanfilo/jax09_maxtext_hotfixCopy head branch name to clipboard