MCPcopy Create free account
hub / github.com/apple/axlearn / _mesh

Function _mesh

axlearn/common/checkpointer_orbax_test.py:25–27  ·  view source on GitHub ↗
(mesh_shape: Sequence[int])

Source from the content-addressed store, hash-verified

23
24
25def _mesh(mesh_shape: Sequence[int]):
26 devices = mesh_utils.create_device_mesh(mesh_shape)
27 return jax.sharding.Mesh(devices, ("data", "model"))
28
29
30class OrbaxCheckpointerTest(test_utils.TestCase):

Callers 1

test_indexMethod · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected