Replace jax.sharding.use_mesh with jax.set_mesh. jax.set_mesh can act as a global setter or a context manager.#65

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/RecML:mainfrom

test_789513590AI-Hypercomputer/RecML:test_789513590Copy head branch name to clipboard