Searched refs:unpad_tensor (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/distributed/tensor/ |
| H A D | placement_types.py | 16 unpad_tensor,
180 output = unpad_tensor(output, self.dim, pad_sizes[mesh_dim_local_rank])
215 …output = unpad_tensor(output, self.dim, pad_sizes[my_coordinate[mesh_dim]]) # type: ignore[possib…
251 result = unpad_tensor(result, self.dim, unpad_size)
324 …new_tensor = unpad_tensor(new_tensor, self.dim, old_dim_unpad_size) # type: ignore[possibly-undef…
331 …new_tensor = unpad_tensor(new_tensor, new_shard_dim, new_dim_unpad_size) # type: ignore[possibly-…
|
| H A D | _collective_utils.py | 178 def unpad_tensor(tensor: torch.Tensor, pad_dim: int, pad_size: int) -> torch.Tensor: function
|
| /aosp_15_r20/external/pytorch/test/distributed/ |
| H A D | test_device_mesh.py | 22 unpad_tensor,
770 scattered_tensor = unpad_tensor(
813 unpad_tensor(big_tensor, shard_dim, pad_sizes[i])
909 scattered_tensor = unpad_tensor(
|
| /aosp_15_r20/external/pytorch/test/distributed/_tensor/ |
| H A D | test_dtensor.py | 866 from torch.distributed._tensor._collective_utils import unpad_tensor
869 unpad_tensor(tensor, shard_placement.dim, pad_sizes[i])
|