MCPcopy Create free account
hub / github.com/NVIDIA/DALI / test_positional_sharding_workflow

Function test_positional_sharding_workflow

dali/test/python/jax_plugin/jax_server.py:92–98  ·  view source on GitHub ↗
(process_id)

Source from the content-addressed store, hash-verified

90
91
92def test_positional_sharding_workflow(process_id):
93 mesh = Mesh(jax.devices(), axis_names=("device",))
94 sharding = NamedSharding(mesh, PartitionSpec("device"))
95
96 run_distributed_sharing_test(sharding=sharding, process_id=process_id)
97
98 log.info("Passed positional sharding workflow test")
99
100
101def test_named_sharding_workflow(process_id):

Callers 1

Calls 2

infoMethod · 0.45

Tested by

no test coverage detected