Fix JAX 0.11.0 sharding compatibility and SFT trainer signature for Tunix integration - #4708

Fix JAX 0.11.0 sharding compatibility and SFT trainer signature for Tunix integration#4708

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

fix_tunix_sft_notebook_jax11AI-Hypercomputer/maxtext:fix_tunix_sft_notebook_jax11Copy head branch name to clipboard