Function _shard_dim_1

tests/python/disco/test_loader.py:58–60

(src, num_shards, tgt)

@register_global_func("tests.disco.shard_dim_1", override=True)
def _shard_dim_1(src, num_shards, tgt):
    s_0, s_1 = src.shape
    tgt.copyfrom(src.numpy().reshape(s_0, num_shards, s_1 // num_shards).transpose(1, 0, 2))


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