Home
last modified time | relevance | path

Searched defs:ShardingSpec (Results 1 – 11 of 11) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/_shard/sharded_tensor/
H A D__init__.py23 ShardingSpec = "ShardingSpec" variable
27 sharding_spec: ShardingSpec,
82 sharding_spec: ShardingSpec,
136 sharding_spec: ShardingSpec,
190 sharding_spec: ShardingSpec,
246 sharding_spec: ShardingSpec,
302 sharding_spec: ShardingSpec,
H A Dreshard.py45 sharding_spec: shard_spec.ShardingSpec,
86 sharding_spec: shard_spec.ShardingSpec,
87 resharding_spec: shard_spec.ShardingSpec,
155 sharding_spec: shard_spec.ShardingSpec,
156 resharding_spec: shard_spec.ShardingSpec,
H A Dapi.py79 def __new__(cls, sharding_spec: shard_spec.ShardingSpec, *size, **kwargs):
240 def __new__(cls, sharding_spec: shard_spec.ShardingSpec, *size, **kwargs):
246 sharding_spec: shard_spec.ShardingSpec,
779 sharding_spec: shard_spec.ShardingSpec,
1022 def reshard(self, resharding_spec: shard_spec.ShardingSpec) -> ShardedTensor:
/aosp_15_r20/external/pytorch/torch/distributed/_shard/
H A Dapi.py17 tensor: torch.Tensor, sharding_spec: ShardingSpec, src_rank=0, process_group=None
80 sharding_spec: ShardingSpec,
166 module: torch.nn.Module, resharding_spec: ShardingSpec
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/
H A Dsharded_device_array.h175 ShardingSpec(std::vector<AvalDimSharding> sharding, in ShardingSpec() function
179 ShardingSpec(pybind11::iterable py_sharding, in ShardingSpec() function
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/proto/
H A Dlayout.proto23 message ShardingSpec { message
/aosp_15_r20/external/pytorch/test/distributed/checkpoint/
H A Dtest_file_system_checkpoint.py93 spec: ShardingSpec,
H A Dtest_file_system_checkpoint_cpu.py95 spec: ShardingSpec,
/aosp_15_r20/external/pytorch/torch/distributed/_shard/sharding_spec/
H A Dapi.py51 class ShardingSpec(ABC): class
/aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/
H A Dfsdp.py118 spec: shard_spec.ShardingSpec, tensor: torch.Tensor, rank: int
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Deinsum_spmd_expander.cc208 std::pair<std::vector<ShardingSpec>, absl::flat_hash_map<std::string, int>> in VerifyOrFixLayout()