[torchxla2] Drop Python 3.9 support

.github/workflows/torch_xla2.yml

@@ -22,7 +22,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        python-version: ['3.9', '3.10']
+        python-version: ['3.10']
     steps:
       - name: Checkout repo
         uses: actions/checkout@v4

experimental/torch_xla2/README.md

@@ -2,7 +2,7 @@
 
 ## Install
 
-Currently this is only source-installable. Requires Python version >= 3.9.
+Currently this is only source-installable. Requires Python version >= 3.10.
 
 ### NOTE:
 

experimental/torch_xla2/pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     # Developers should install `dev-requirements.txt` first
     "torch>=2.3.0",
 ]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 license = {file = "LICENSE"}
 dynamic = ["version"]
 

experimental/torch_xla2/test/test_ops.py

@@ -62,8 +62,6 @@
   "chalf", # Skip due to jax not support complex32 with backend: https://github.com/google/jax/issues/14180
   "__rpow__",  # NOTE: cannot fix because torch test case has undefined behavior
                # such as 0 to negative power.
-  "ceil", # only failed with python 3.9
-  "trunc", # only failed with python 3.9
   "to_sparse", # We are not supporting sparse tensors yet.
   "nn.functional.rrelu", # pure torch result match torch_xla2 test result, only OpInfo mismatch: https://gist.github.com/ManfeiBai/1a449b15f4e946bfcaa3e5ef86da20f4
 }

experimental/torch_xla2/torch_xla2/ops/jax_reimplement.py

@@ -86,7 +86,7 @@ def _scale_and_translate(x, output_shape: core.Shape,
 def scale_and_translate(image, shape: core.Shape,
                         spatial_dims: Sequence[int],
                         scale, translation,
-                        # (barney-s) use string, python 3.9 throws: E   TypeError: unsupported operand type(s) for |: 'type' and 'EnumMeta'
+                        # (barney-s) use string
                         method: str, #(barney-s) | ResizeMethod,
                         antialias: bool = True,
                         precision=lax.Precision.HIGHEST):

experimental/torch_xla2/torch_xla2/ops/op_base.py

@@ -7,12 +7,7 @@
 from torch_xla2 import types
 import sys
 
-if sys.version_info < (3, 10):
-  from typing_extensions import ParamSpec, Concatenate
-else:
-  from typing import ParamSpec, Concatenate
-
-from typing import Callable, Optional
+from typing import Callable, Optional, ParamSpec, Concatenate
 
 
 class InplaceOp:

experimental/torch_xla2/torch_xla2/types.py

@@ -1,14 +1,9 @@
-from typing import Callable, Any, Union
+from typing import Callable, Any, Union, ParamSpec, TypeAlias
 import torch
 import jax
 import jax.numpy as jnp
 import sys
 
-if sys.version_info < (3, 10):
-  from typing_extensions import ParamSpec, TypeAlias
-else:
-  from typing import ParamSpec, TypeAlias
-
 P = ParamSpec('P')
 
 TorchValue: TypeAlias = Union[torch.Tensor, torch.dtype, 'TorchCallable', Any]