def _init_chunked(self, dims, tensor_init_params: TensorInitParams):
    current_rank = dist.get_rank(self._process_group)
    sharding_dim = self._sharding_spec.dim
    if (not isinstance(sharding_dim, int)):
        raise ValueError(f'Sharding dim needs to be an integer, found: {sharding_dim}')
    if ((sharding_dim >= len(dims)) or (sharding_dim < (- len(dims)))):
        raise ValueError(f'Invalid sharding dim: {sharding_dim}')
    dim_size = dims[sharding_dim]
    remote_devices = self._sharding_spec.placements
    chunks = len(remote_devices)
    split_size = get_split_size(dim_size, chunks)
    shards_metadata = []
    for (idx, remote_device) in enumerate(remote_devices):
        (rank, local_device) = _parse_and_validate_remote_device(self._process_group, remote_device)
        sharded_dim_size = get_chunked_dim_size(dim_size, split_size, idx)
        if (sharded_dim_size > 0):
            rank_dims = dims.copy()
            rank_offsets = ([0] * len(dims))
            rank_offsets[sharding_dim] = (split_size * idx)
            rank_dims[sharding_dim] = sharded_dim_size
            shard_metadata = ShardMetadata(rank_offsets, rank_dims, remote_device)
            shards_metadata.append(shard_metadata)
            if (current_rank == rank):
                local_shard = _create_tensor_from_params(*rank_dims, local_device=local_device, tensor_init_params=tensor_init_params)
                self._local_shards.append(Shard(local_shard, shard_metadata))
    self._metadata = ShardedTensorMetadata(shards_metadata, dims, tensor_init_params.tensor_properties)
def _init_chunked(self, dims, tensor_init_params: TensorInitParams):
    current_rank = dist.get_rank(self._process_group)
    sharding_dim = self._sharding_spec.dim
    if (not isinstance(sharding_dim, int)):
        raise ValueError(f'Sharding dim needs to be an integer, found: {sharding_dim}')
    if ((sharding_dim >= len(dims)) or (sharding_dim < (- len(dims)))):
        raise ValueError(f'Invalid sharding dim: {sharding_dim}')
    dim_size = dims[sharding_dim]
    remote_devices = self._sharding_spec.placements
    chunks = len(remote_devices)
    split_size = get_split_size(dim_size, chunks)
    shards_metadata = []
    for (idx, remote_device) in enumerate(remote_devices):
        (rank, local_device) = _parse_and_validate_remote_device(self._process_group, remote_device)
        sharded_dim_size = get_chunked_dim_size(dim_size, split_size, idx)
        if (sharded_dim_size > 0):
            rank_dims = dims.copy()
            rank_offsets = ([0] * len(dims))
            rank_offsets[sharding_dim] = (split_size * idx)
            rank_dims[sharding_dim] = sharded_dim_size
            shard_metadata = ShardMetadata(rank_offsets, rank_dims, remote_device)
            shards_metadata.append(shard_metadata)
            if (current_rank == rank):
                local_shard = _create_tensor_from_params(*rank_dims, local_device=local_device, tensor_init_params=tensor_init_params)
                self._local_shards.append(Shard(local_shard, shard_metadata))
    self._metadata = ShardedTensorMetadata(shards_metadata, dims, tensor_init_params.tensor_properties)