Skip to content

Commit

Permalink
Improve some typing annotations (#8369)
Browse files Browse the repository at this point in the history
  • Loading branch information
tengyifei authored Nov 12, 2024
1 parent 5a4641b commit c4f2771
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 3 deletions.
4 changes: 2 additions & 2 deletions torch_xla/core/xla_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ def is_master_ordinal(local: bool = True) -> bool:
return ordinal == 0


def master_print(*args: Tuple[Any, ...],
def master_print(*args: Any,
fd: TextIO = sys.stdout,
local: bool = False,
flush: bool = False):
Expand Down Expand Up @@ -984,7 +984,7 @@ def _reduce_scatter_coalesced(


def add_step_closure(closure: Callable[..., Any],
args: Tuple[Any] = (),
args: Tuple[Any, ...] = (),
run_async: bool = False):
"""Adds a closure to the list of the ones to be run at the end of the step.
Expand Down
3 changes: 2 additions & 1 deletion torch_xla/distributed/spmd/xla_sharding.py
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,8 @@ def disable_manual_sharding(t: Union[torch.Tensor, XLAShardedTensor],

def mark_sharding(
t: Union[torch.Tensor, XLAShardedTensor], mesh: Mesh,
partition_spec: Tuple[Union[Tuple, int, str, None]]) -> XLAShardedTensor:
partition_spec: Tuple[Union[Tuple, int, str, None],
...]) -> XLAShardedTensor:
"""
Annotates the tensor provided with XLA partition spec. Internally,
it annotates the corresponding XLATensor as sharded for the XLA SpmdPartitioner pass.
Expand Down

0 comments on commit c4f2771

Please sign in to comment.