Home
last modified time | relevance | path

Searched defs:DTensorSpec (Results 1 – 12 of 12) sorted by relevance

/aosp_15_r20/external/pytorch/torch/distributed/tensor/_ops/
H A Dutils.py120 def is_tensor_shardable(shape: Sequence[int], spec: DTensorSpec) -> bool:
138 def is_tensor_evenly_shardable(shape: Sequence[int], spec: DTensorSpec) -> bool:
154 def is_tensor_dim_sharded(spec: DTensorSpec, dim: int) -> bool:
159 def is_tensor_partial(spec: DTensorSpec) -> bool:
210 src_strategy: OpStrategy, dst_spec: DTensorSpec
H A D_common_rules.py23 input_specs: Tuple[DTensorSpec, ...],
/aosp_15_r20/external/pytorch/torch/distributed/tensor/
H A D_random.py152 def _distribute_region(self, spec: DTensorSpec):
174 def _distribute_region(self, spec: DTensorSpec):
214 def _set_pre_op_offset(self, spec: DTensorSpec) -> None:
306 def _set_post_op_offset(self, spec: DTensorSpec, old_offset: int) -> None:
361 def _distribute_region(self, spec: DTensorSpec):
H A D_redistribute.py32 src_spec: DTensorSpec,
33 dst_spec: DTensorSpec,
151 current_spec: DTensorSpec,
152 target_spec: DTensorSpec,
H A D_collective_utils.py227 def spec_to_bytes(spec: "dtensor_spec.DTensorSpec") -> int:
314 current_spec: "dtensor_spec.DTensorSpec",
315 target_spec: "dtensor_spec.DTensorSpec",
H A D_dtensor_spec.py25 class DTensorSpec: class
H A D_dispatch.py232 def default_tensor(spec: DTensorSpec) -> torch.Tensor:
H A D_sharding_prop.py474 spec: DTensorSpec,
H A D_api.py246 spec: DTensorSpec,
/aosp_15_r20/external/pytorch/torch/distributed/tensor/experimental/
H A D_tp_transform.py269 input_specs: Optional[Sequence[DTensorSpec]] = None,
316 def update_arg_spec(arg_spec: DTensorSpec) -> DTensorSpec:
404 def _partition_val(val: Any, spec: DTensorSpec) -> Any:
436 input_arg_spec: DTensorSpec,
437 desired_spec: DTensorSpec,
/aosp_15_r20/external/pytorch/test/distributed/_tensor/experimental/
H A Dtest_register_sharding.py31 x: DTensorSpec,
/aosp_15_r20/external/pytorch/torch/distributed/_composable/fsdp/
H A D_fsdp_common.py112 sharding_spec: DTensorSpec,