frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Migrate from Tuple -> tuple in torch/distributed (#144258)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/144258
Approved by: https://github.com/aorenste
This commit is contained in:
bobrenjc93
2025-01-09 20:17:01 -08:00
committed by PyTorch MergeBot
parent 184549b2d7
commit 08be9ec312
103 changed files with 468 additions and 553 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
torch/distributed/pipelining/microbatch.py
View File

@ -73,7 +73,7 @@ class TensorChunkSpec:
@staticmethod
def from_tuple(
chunk_dims: Tuple[int, ...],
chunk_dims: tuple[int, ...],
):
"""
A helper for creating a tuple of `TensorChunkSpec` from a tuple of chunk
@ -242,12 +242,12 @@ def _shard_dict_of_args(
def split_args_kwargs_into_chunks(
args: Tuple[Any, ...],
args: tuple[Any, ...],
kwargs: Optional[Dict[str, Any]],
chunks: int,
args_chunk_spec: Optional[Tuple[TensorChunkSpec, ...]] = None,
args_chunk_spec: Optional[tuple[TensorChunkSpec, ...]] = None,
kwargs_chunk_spec: Optional[Dict[str, TensorChunkSpec]] = None,
) -> Tuple[List[Tuple], List[Dict]]:
) -> tuple[List[Tuple], List[Dict]]:
"""
Given a sequence of args and kwargs, split them into a number of chunks
according to their respective chunking specs.