Update JAX tutorial to reflect changes in device mesh creation and sharding strategies. - #288

Update JAX tutorial to reflect changes in device mesh creation and sharding strategies.#288

Davidnet wants to merge 1 commit into

jax-ml:mainjax-ml/jax-ai-stack:mainfrom

Davidnet:davidnet-update-jax-llm-pretrainingDavidnet/jax-ai-stack:davidnet-update-jax-llm-pretrainingCopy head branch name to clipboard