Use Tuple instead of tuple in scan.py (#8067)

pytorch · Sep 25, 2024 · 1717e25 · 1717e25
1 parent 3c7daa2
commit 1717e25
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torch_xla/experimental/scan.py b/torch_xla/experimental/scan.py
@@ -4,7 +4,7 @@
 
 """
 
-from typing import Callable, TypeVar
+from typing import Callable, TypeVar, Tuple
 
 import torch
 from torch.utils._pytree import tree_map, tree_iter
@@ -15,10 +15,10 @@
 
 
 def scan(
-    fn: Callable[[Carry, X], tuple[Carry, Y]],
+    fn: Callable[[Carry, X], Tuple[Carry, Y]],
     init: Carry,
     xs: X,
-) -> tuple[Carry, Y]:
+) -> Tuple[Carry, Y]:
   """Apply a function over leading dimension of tensors while carrying along state.
   
   This is similar to the JAX `jax.lax.scan` function found in [1].