Commit 8a2ba76

jax authors

committed

Enable Creating Device Mesh with Physical Axes Splits.

PiperOrigin-RevId: 617994892

1 parent 5532e55 commit 8a2ba76Copy full SHA for 8a2ba76

2 files changed

+623

-58

lines changed

jax/experimental
- mesh_utils.py
tests
- mesh_utils_test.py

2 files changed

+623

-58

lines changed

Comments

(0)