Function _shard_dim_0

tests/python/disco/test_loader.py:52–54

(src, num_shards, tgt)

@register_global_func("tests.disco.shard_dim_0", override=True)
def _shard_dim_0(src, num_shards, tgt):
    s_0, s_1 = src.shape
    tgt.copyfrom(src.numpy().reshape(num_shards, s_0 // num_shards, s_1))


@register_global_func("tests.disco.shard_dim_1", override=True)