Support PyTorch 2.9

coremltools/_deps/__init__.py

@@ -163,7 +163,7 @@ def __get_sklearn_version(version):
 
 # ---------------------------------------------------------------------------------------
 _HAS_TORCH = True
-_TORCH_MAX_VERSION = "2.8.0"
+_TORCH_MAX_VERSION = "2.9.0"
 _HAS_TORCH_EXPORT_API = False
 _CT_OPTIMIZE_TORCH_MIN_VERSION = "2.1.0"
 _IMPORT_CT_OPTIMIZE_TORCH = False

coremltools/converters/_converters_entry.py

@@ -591,6 +591,18 @@ def forward(self, x):
                                      outputs_as_strings,
                                      outputs_as_tensor_or_image_types,
                                      outputs)
+    # Starting with torch 2.9, torch.export.export() returns an ExportedProgram
+    # in the new "TRAINING" IR dialect by default, whereas conversion expects
+    # the "ATEN" dialect (the prior default). Lower it to ATEN via
+    # run_decompositions -- the migration torch recommends -- so existing
+    # convert() calls keep working without a manual run_decompositions() step.
+    if (
+        exact_source == "pytorch"
+        and _HAS_TORCH_EXPORT_API
+        and isinstance(model, ExportedProgram)
+        and model.dialect not in ("ATEN", "EDGE")
+    ):
+        model = model.run_decompositions({})
     source_dialect = _determine_source_dialect(model, exact_source)
     exact_target = _determine_target(convert_to, minimum_deployment_target)
     _validate_conversion_arguments(

coremltools/converters/mil/frontend/torch/ops.py

@@ -3131,12 +3131,18 @@ def _parse_positional_args(context, node) -> Tuple[Var]:
 
 @register_torch_op
 def rms_norm(context, node):
-    # Parse Inputs
-    inputs = _get_inputs(context, node, expected=4)
+    # Parse Inputs. torch.export / ExecuTorch omit the optional weight/eps when
+    # they use defaults, so only window-style strict counts apply to TorchScript.
+    inputs = _get_inputs(
+        context,
+        node,
+        expected={TorchFrontend.TORCHSCRIPT: 4},
+        min_expected={TorchFrontend.TORCHEXPORT: 2, TorchFrontend.EXECUTORCH: 2},
+    )
     x = inputs[0]
     normalized_shape = inputs[1]
-    weight = inputs[2]
-    eps = inputs[3]
+    weight = inputs[2] if len(inputs) > 2 else None
+    eps = inputs[3] if len(inputs) > 3 else None
     axes = list(range(-len(normalized_shape.val), 0))
     # Store epsilon value to ensure ZeroDivisionError doesn't occur
     # while computing RMSNorm
@@ -8277,14 +8283,28 @@ def remainder(context, node):
     context.add(remainder_node)
 
 
-@register_torch_op
+@register_torch_op(torch_alias=["hann_window.periodic"])
 def hann_window(context, node):
-    inputs = _get_inputs(context, node, expected=[5, 6])
+    inputs = _get_inputs(
+        context,
+        node,
+        expected={TorchFrontend.TORCHSCRIPT: [5, 6]},
+        min_expected={TorchFrontend.TORCHEXPORT: 1, TorchFrontend.EXECUTORCH: 1},
+    )
     if inputs[0].val is None:
         raise NotImplementedError("variable 'window_length' not supported.")
 
+    # TorchScript passes the dtype/layout/device/pin_memory kwargs as graph inputs
+    # (5 for the default overload, 6 for `.periodic`); torch.export / ExecuTorch
+    # pass only window_length (and periodic for the `.periodic` overload). So the
+    # `periodic` flag is the 2nd positional input for the 6-arg TorchScript form or
+    # any >=2-arg export form.
     periodic = True
-    if len(inputs) == 6:
+    if context.frontend == TorchFrontend.TORCHSCRIPT:
+        has_periodic = len(inputs) == 6
+    else:
+        has_periodic = len(inputs) >= 2
+    if has_periodic:
         if inputs[1].val is None:
             raise NotImplementedError("variable 'periodic' not supported.")
         if not inputs[1].val:

coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py

@@ -22,6 +22,7 @@
 from coremltools._deps import (
     _HAS_HF,
     _HAS_TORCH,
+    _HAS_TORCH_EXPORT_API,
     _HAS_TORCHAO,
     MSG_TORCH_NOT_FOUND,
     MSG_TORCHAO_NOT_FOUND,
@@ -571,6 +572,41 @@ def forward(self, x):
         )
         assert isinstance(model._get_mil_internal(), ct.converters.mil.Program)
 
+    @staticmethod
+    @pytest.mark.skipif(not _HAS_TORCH_EXPORT_API, reason="torch.export API not available.")
+    def test_convert_exported_program_training_dialect():
+        # Starting with torch 2.9, torch.export.export() returns an
+        # ExportedProgram in the new "TRAINING" IR dialect by default. convert()
+        # should lower it automatically via run_decompositions instead of
+        # requiring the user to call run_decompositions() themselves.
+        class Network(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.hidden = torch.nn.Linear(8, 4)
+                self.relu = torch.nn.ReLU()
+
+            def forward(self, x):
+                return self.relu(self.hidden(x))
+
+        torch_model = Network().eval()
+        example_input = torch.rand(2, 8)
+        exported_program = torch.export.export(torch_model, (example_input,))
+
+        # With torch >= 2.9 the default export dialect is no longer ATEN, so the
+        # installed torch here actually drives convert()'s auto-lowering path.
+        # Assert it to guarantee this test exercises that path instead of
+        # silently passing as a no-op on a torch whose default is still ATEN.
+        if Version(torch.__version__) >= Version("2.9.0"):
+            assert exported_program.dialect not in ("ATEN", "EDGE")
+
+        # Convert directly, without a manual run_decompositions() step.
+        mlmodel = ct.convert(
+            exported_program,
+            inputs=[ct.TensorType(name="x", shape=example_input.shape)],
+            convert_to="mlprogram",
+        )
+        assert isinstance(mlmodel._get_mil_internal(), ct.converters.mil.Program)
+
     @staticmethod
     @pytest.mark.skipif(ct.utils._macos_version() < (12, 0), reason='Model produces specification 6.')
     @pytest.mark.parametrize(

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -12566,15 +12566,16 @@ def test_logcumsumexp(self, compute_unit, backend, frontend, axis):
 
 class TestHannWindow(TorchBaseTest):
     @pytest.mark.parametrize(
-        "compute_unit, backend, window_length, periodic",
+        "compute_unit, backend, frontend, window_length, periodic",
         itertools.product(
             compute_units,
             backends,
+            frontends,
             [1, 3, 6, 10, 12],
             [True, False],
         ),
     )
-    def test_hann_window(self, compute_unit, backend, window_length, periodic):
+    def test_hann_window(self, compute_unit, backend, frontend, window_length, periodic):
         class HannWindowModel(nn.Module):
             def forward(self, x):
                 return torch.hann_window(window_length, periodic)
@@ -12588,6 +12589,29 @@ def forward(self, x):
             model,
             expected_results=torch_out,
             input_as_shape=False,
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+        )
+
+
+class TestRMSNorm(TorchBaseTest):
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend, has_weight",
+        itertools.product(compute_units, backends, frontends, [True, False]),
+    )
+    def test_rms_norm(self, compute_unit, backend, frontend, has_weight):
+        normalized_shape = (8,)
+        weight = torch.rand(normalized_shape) if has_weight else None
+
+        class RMSNormModel(nn.Module):
+            def forward(self, x):
+                return nn.functional.rms_norm(x, normalized_shape, weight=weight, eps=1e-5)
+
+        self.run_compare_torch(
+            (2, 4, 8),
+            RMSNormModel().eval(),
+            frontend=frontend,
             backend=backend,
             compute_unit=compute_unit,
         )
@@ -13768,6 +13792,23 @@ def test_unfold(
         if frontend == TorchFrontend.EXECUTORCH:
             pytest.skip("ExecuTorch produces rank > 5 tensor")
 
+        if (
+            is_dynamic_hw
+            and frontend in TORCH_EXPORT_BASED_FRONTENDS
+            and Version(torch.__version__) >= Version("2.9.0")
+        ):
+            # torch 2.9's ExportedProgram.run_decompositions({}) raises
+            # "NameError: name 'L' is not defined" while interpreting the
+            # `_guards_fn` submodule that torch 2.9 generates for dynamic-shape
+            # exports carrying shape guards (here the H/W >= f(kernel, dilation,
+            # padding, stride) constraint). This is an upstream torch issue, not
+            # a converter bug; static-shape unfold and all other export ops are
+            # unaffected. Re-enable once the torch regression is resolved.
+            pytest.skip(
+                "torch 2.9 run_decompositions() NameError on _guards_fn "
+                "for guarded dynamic-shape exports"
+            )
+
         if isinstance(kernel_size, int):
             kernel_size = (kernel_size, kernel_size)
         if isinstance(dilation, int):

reqs/pytorch.pip

@@ -5,12 +5,12 @@ torchaudio==2.2.0; platform_machine != "arm64"
 torchvision==0.17.0; platform_machine != "arm64"
 
 # Torch dependencies for ARM
-torch==2.8.0; platform_machine == "arm64"
+torch==2.9.0; platform_machine == "arm64"
 torchaudio>=2.2.0; platform_machine == "arm64"
 torchvision>=0.17.0; platform_machine == "arm64"
 torchsr==1.0.4; platform_machine == "arm64"
 
 # TODO (rdar://141476729) support a more recent timm
 timm==0.6.13; platform_machine == "arm64"
 
-torchao==0.12.0; platform_machine == "arm64" and python_version >= '3.10'
+torchao==0.14.0; platform_machine == "arm64" and python_version >= '3.10'

reqs/test_executorch.pip

@@ -3,5 +3,5 @@
 
 # Warning: Starting from ExecuTorch 0.6.0, coremltools is added as a dependency
 # so we need to re-install built-from-source coremltools after pip install ExecuTorch
-executorch>=0.7.0; platform_machine == "arm64" and python_version >= '3.10' and python_version < '3.13'
+executorch>=1.0.0,<1.1.0; platform_machine == "arm64" and python_version >= '3.10' and python_version < '3.13'
 setuptools<82