Skip to content

Commit

Permalink
Update tuple type annotation (#7699)
Browse files Browse the repository at this point in the history
  • Loading branch information
JackCaoG authored Jul 16, 2024
1 parent a69fd51 commit e523fdf
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions torch_xla/experimental/custom_kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import torch_xla.distributed.spmd as xs
import torch_xla.debug.metrics as met

from typing import Any, List, Callable, Optional
from typing import Any, List, Callable, Optional, Tuple
from torch.library import impl
from torch_xla.core.xla_model import XLA_LIB

Expand Down Expand Up @@ -719,7 +719,7 @@ def gmm(
lhs: torch.Tensor,
rhs: torch.Tensor,
group_sizes: torch.Tensor,
tiling: tuple[int, int, int] = (512, 512, 512)
tiling: Tuple[int, int, int] = (512, 512, 512)
) -> torch.Tensor:
"""Compute lhs[sizes[i-1]:sizes[i], :] @ rhs for each group 'i'.
Expand Down Expand Up @@ -770,7 +770,7 @@ def tgmm(
lhs: torch.Tensor,
rhs: torch.Tensor,
group_sizes: torch.Tensor,
tiling: tuple[int, int, int] = (512, 512, 512)
tiling: Tuple[int, int, int] = (512, 512, 512)
) -> torch.Tensor:
"""Compute lhs[:, sizes[i-1]:sizes[i]] @ rhs[sizes[i-1]:sizes[i], :].
Expand Down

0 comments on commit e523fdf

Please sign in to comment.