Searched defs:DTensor (Results 1 – 10 of 10) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/ |
| H A D | fsdp.py | 32 def _get_box(tensor: DTensor) -> Tuple[torch.Size, torch.Size]:
48 def _get_box_for(tensor: DTensor, idx: int) -> Tuple[torch.Size, torch.Size]:
53 def _get_local_box(tensor: DTensor) -> Tuple[torch.Size, torch.Size]:
60 def _create_shard_md_from_dt(dt: DTensor, current_rank: int) -> ShardMetadata:
73 dt: DTensor, dt_pg: c10d.ProcessGroup
111 def _get_dt_pg(dt: DTensor) -> c10d.ProcessGroup:
298 tensor: DTensor,
385 tensor: DTensor,
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/ |
| H A D | test_dtensor_checkpoint.py | 30 sdt: DTensor,
31 rdt: DTensor,
32 submesh_sdt: DTensor,
33 submesh_rdt: DTensor,
|
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/ |
| H A D | _redistribute.py | 279 input: "dtensor.DTensor",
309 def backward(ctx, grad_output: "dtensor.DTensor"): # type: ignore[override]
|
| H A D | _api.py | 75 input: "DTensor",
188 def backward(ctx, grad_output: "DTensor"): # type: ignore[override]
215 class DTensor(torch.Tensor): class
|
| H A D | README.md | 43 ### DTensor API
|
| H A D | _dispatch.py | 479 dtensor_arg: "dtensor.DTensor",
|
| /aosp_15_r20/external/pytorch/torch/distributed/checkpoint/ |
| H A D | planner_helpers.py | 74 def _create_write_items_for_dtensor(fqn: str, tensor: DTensor) -> WriteItem:
240 def _create_chunk_from_dtensor(tensor: DTensor) -> ChunkStorageMetadata:
|
| /aosp_15_r20/external/pytorch/torch/distributed/fsdp/ |
| H A D | _fsdp_extensions.py | 77 tensor: DTensor,
170 tensor: DTensor,
|
| H A D | _shard_utils.py | 118 tensor: DTensor,
|
| /aosp_15_r20/external/pytorch/torch/distributed/ |
| H A D | CONTRIBUTING.md | 26 #### DTensor subsubsection
|