diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index f516901..5417855 100644 --- a/.github/workflows/unit-tests.yml +++ b/.github/workflows/unit-tests.yml @@ -57,3 +57,25 @@ jobs: for FILE in $(find docs -name '*.py'); do python $FILE done + + pytorch-2-13-compatibility: + name: PyTorch 2.13 compatibility + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v6 + + - name: Set up Python 3.12 + uses: actions/setup-python@v6 + with: + python-version: "3.12" + - name: Install dependencies + run: | + pip install --upgrade pip + pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu + pip install pytest + pip install -e .[cpu] + - name: Test named tensor removal compatibility + env: + JAX_PLATFORMS: cpu + run: pytest test/test_torch_compat.py test/test_amp.py diff --git a/test/test_torch_compat.py b/test/test_torch_compat.py new file mode 100644 index 0000000..3e1581e --- /dev/null +++ b/test/test_torch_compat.py @@ -0,0 +1,37 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +import torchax +from torchax._torch_compat import get_aten_overload +from torchax.amp import CastPolicy, autocast_policy + + +def test_get_aten_overload_returns_existing_overload(): + assert get_aten_overload("prod", "default") is torch.ops.aten.prod.default + + +def test_get_aten_overload_returns_none_for_missing_overload(): + assert get_aten_overload("prod", "not_an_overload") is None + + +def test_available_named_tensor_overload_keeps_autocast_policy(): + named_prod = get_aten_overload("prod", "dim_Dimname") + if named_prod is not None: + assert autocast_policy[named_prod] is CastPolicy.FP32 + + +def test_default_environment_loads_available_decompositions(): + assert torchax.default_env() is not None diff --git a/torchax/_torch_compat.py b/torchax/_torch_compat.py new file mode 100644 index 0000000..dbcb30d --- /dev/null +++ b/torchax/_torch_compat.py @@ -0,0 +1,21 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + + +def get_aten_overload(packet_name: str, overload_name: str): + """Returns an ATen overload when it is available in the installed PyTorch.""" + packet = getattr(torch.ops.aten, packet_name) + return getattr(packet, overload_name, None) diff --git a/torchax/amp.py b/torchax/amp.py index 2007dba..9638fd9 100644 --- a/torchax/amp.py +++ b/torchax/amp.py @@ -18,6 +18,8 @@ import torch from torch.utils import _pytree as pytree +from torchax._torch_compat import get_aten_overload + # enum class CastPolicy : uint8_t { # lower_precision_fp = 0, // Cast all inputs to lower_precision_fp before @@ -115,7 +117,6 @@ def autocast(device, dtype=torch.bfloat16, env=None): torch.ops.aten.polar.default: CastPolicy.FP32, torch.ops.aten.prod.default: CastPolicy.FP32, torch.ops.aten.prod.dim_int: CastPolicy.FP32, - torch.ops.aten.prod.dim_Dimname: CastPolicy.FP32, torch.ops.aten.quantile.default: CastPolicy.FP32, torch.ops.aten.quantile.scalar: CastPolicy.FP32, torch.ops.aten.nanquantile.default: CastPolicy.FP32, @@ -213,5 +214,13 @@ def autocast(device, dtype=torch.bfloat16, env=None): torch.ops.aten.stack.default: CastPolicy.PROMOTE, torch.ops.aten.cat.default: CastPolicy.PROMOTE, torch.ops.aten.index_copy.default: CastPolicy.PROMOTE, - torch.ops.aten.index_copy.dimname: CastPolicy.PROMOTE, } + +# Named tensor overloads were removed in PyTorch 2.13. Keep registering them +# when running against older PyTorch releases, where they remain available. +for op, policy in ( + (get_aten_overload("prod", "dim_Dimname"), CastPolicy.FP32), + (get_aten_overload("index_copy", "dimname"), CastPolicy.PROMOTE), +): + if op is not None: + autocast_policy[op] = policy diff --git a/torchax/decompositions.py b/torchax/decompositions.py index fe7c2c1..bab9221 100644 --- a/torchax/decompositions.py +++ b/torchax/decompositions.py @@ -354,8 +354,6 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens torch.ops.aten.all.out, torch.ops.aten.all.dims_out, torch.ops.aten.all.all_out, - torch.ops.aten.all.dimname, - torch.ops.aten.all.dimname_out, torch.ops.aten.aminmax.default, torch.ops.aten.aminmax.out, torch.ops.aten.arange.default, @@ -469,23 +467,16 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens torch.ops.aten.im2col.out, torch.ops.aten.index_add.default, torch.ops.aten.index_add.out, - torch.ops.aten.index_add.dimname, torch.ops.aten.index_add_.default, torch.ops.aten.index_copy.default, - torch.ops.aten.index_copy.dimname, torch.ops.aten.index_copy.out, torch.ops.aten.index_copy_.default, - torch.ops.aten.index_copy_.dimname, torch.ops.aten.index_fill.int_Tensor, torch.ops.aten.index_fill.int_Scalar, - torch.ops.aten.index_fill.Dimname_Scalar, - torch.ops.aten.index_fill.Dimname_Tensor, torch.ops.aten.index_fill.int_Scalar_out, torch.ops.aten.index_fill.int_Tensor_out, torch.ops.aten.index_fill_.int_Tensor, torch.ops.aten.index_fill_.int_Scalar, - torch.ops.aten.index_fill_.Dimname_Scalar, - torch.ops.aten.index_fill_.Dimname_Tensor, torch.ops.aten.isin.Tensor_Tensor, torch.ops.aten.isin.Tensor_Tensor_out, torch.ops.aten.isin.Tensor_Scalar, @@ -589,16 +580,12 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens torch.ops.aten.nll_loss_forward.output, torch.ops.aten.norm.Scalar, torch.ops.aten.norm.ScalarOpt_dim, - torch.ops.aten.norm.names_ScalarOpt_dim, torch.ops.aten.norm.ScalarOpt_dim_dtype, torch.ops.aten.norm.dtype_out, torch.ops.aten.norm.out, torch.ops.aten.norm.ScalarOpt_dtype, torch.ops.aten.norm.ScalarOpt_dtype_out, torch.ops.aten.norm.Scalar_out, - torch.ops.aten.norm.names_ScalarOpt_dim_dtype, - torch.ops.aten.norm.names_dtype_out, - torch.ops.aten.norm.names_out, torch.ops.aten.ones.default, torch.ops.aten.ones_like.default, torch.ops.aten.ones_like.out, @@ -701,17 +688,11 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens torch.ops.aten.std.default, torch.ops.aten.std.dim, torch.ops.aten.std.correction, - torch.ops.aten.std.names_dim, - torch.ops.aten.std.names_out, torch.ops.aten.std.out, torch.ops.aten.std.correction_out, - torch.ops.aten.std.correction_names, - torch.ops.aten.std.correction_names_out, torch.ops.aten.std_mean.default, torch.ops.aten.std_mean.dim, torch.ops.aten.std_mean.correction, - torch.ops.aten.std_mean.names_dim, - torch.ops.aten.std_mean.correction_names, torch.ops.aten.std_mean.correction_out, torch.ops.aten.stack.default, torch.ops.aten.stack.out, @@ -739,7 +720,6 @@ def get_summand(ix: torch.Tensor, iy: torch.Tensor, iz: torch.Tensor, w) -> Tens torch.ops.aten.triu.out, torch.ops.aten.triu_.default, torch.ops.aten.unbind.int, - torch.ops.aten.unbind.Dimname, torch.ops.aten.unfold_backward.default, torch.ops.aten.unfold_backward.out, torch.ops.aten.unfold_copy.default,