Function _shard_arg

imperative/python/megengine/xla/sharding.py:374–376

(arg, devices, arg_indices, sharding=None)

def shard_args(devices, indices, args, shardings=None):
    def _shard_arg(arg, devices, arg_indices, sharding=None):
        arg = canonicalize_arg(arg)
        return shard_arg_handlers[type(arg)](arg, devices, arg_indices, sharding)

    if shardings is None:
        return [_shard_arg(arg, devices, indices[i]) for i, arg in enumerate(args)]