Home
last modified time | relevance | path

Searched defs:ShardMetadata (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/_shard/sharding_spec/
H A D_internals.py7 def _check_shard_metadata_pair_overlap(shard1: ShardMetadata, shard2: ShardMetadata):
27 shards: List[ShardMetadata], sharded_dims: List[int]
58 shards: List[ShardMetadata], dim: int
72 def validate_non_overlapping_shards_metadata(shards: List[ShardMetadata]):
/aosp_15_r20/external/pytorch/torch/distributed/checkpoint/
H A Dplanner_helpers.py47 def _chunk_for_shard(shard_md: ShardMetadata) -> ChunkStorageMetadata:
55 sharded_tensor: ShardedTensor, shard_md: ShardMetadata
95 fqn: str, sharded_tensor: ShardedTensor, shard_md: ShardMetadata
/aosp_15_r20/external/pytorch/torch/distributed/_shard/
H A D_utils.py28 def narrow_tensor(tensor: torch.Tensor, metadata: ShardMetadata) -> torch.Tensor:
H A Dmetadata.py10 class ShardMetadata: class