Fix PyTorch 2.13 named overload compatibility

.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

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

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)

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

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,