Replace 'with mesh' with 'with jax.set_mesh(mesh)'#4034

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

with_mesh_upgradeAI-Hypercomputer/maxtext:with_mesh_upgradeCopy head branch name to clipboard