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

Function _mesh

axlearn/common/checkpointer_test.py:64–66  ·  view source on GitHub ↗
(mesh_shape: Sequence[int])

Source from the content-addressed store, hash-verified

62
63
64def _mesh(mesh_shape: Sequence[int]):
65 devices = mesh_utils.create_device_mesh(mesh_shape)
66 return jax.sharding.Mesh(devices, ("data", "model"))
67
68
69def _checkpointer_config(

Callers 15

test_save_and_restoreMethod · 0.70
test_save_and_restore_meshMethod · 0.70
test_save_restore_files_countMethod · 0.70
test_save_and_restore_latest_validMethod · 0.70
test_save_can_override_on_gcsMethod · 0.70
test_gdaMethod · 0.70
test_custom_dictMethod · 0.70
test_elastic_input_iteratorMethod · 0.70
test_input_iteratorMethod · 0.70
test_python_savableMethod · 0.70
test_grainMethod · 0.70
test_garbage_collectionMethod · 0.70

Calls

no outgoing calls

Tested by

no test coverage detected