From 9832c950a4af09327abc3a3cf26d78e1bfd743c8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 16:13:52 -0700
Subject: [PATCH 01/44] Keep Core ML inputs it cannot type, such as uint8, on
 the CPU

coremltools 9.0 cannot type a Core ML program input whose dtype is
missing from its TORCH_DTYPE_TO_MIL_DTYPE map, such as uint8 or int8.
It raises "KeyError: torch.uint8" instead of reporting the op as
unsupported. So a model that turns a uint8 camera frame into floats,
for example image.permute(2, 0, 1).float() / 255.0, failed to lower at
all (issue #11686).

After proposing partitions, the partitioner now looks for ops that read
such a tensor from outside their partition: a user input, the output of
an op left on the CPU, or a weight, constant or state it does not take
over. It keeps those ops on the CPU and partitions again, until no
partition reads one. The ops after the conversion to float still go to
Core ML. A tensor of that dtype made inside a partition is left alone,
so quantize and dequantize pairs, codebook weights and taken-over state
are delegated as before. A graph without such tensors is partitioned
once, as before. With lower_full_graph=True the partitioner now says
that the input dtype is the problem, instead of calling a supported op
unsupported.

An op kept on the CPU this way needs a CPU kernel. Some ops that Core ML
takes whole, such as aten.pad and aten.max_pool2d, have no out variant,
so ops_to_not_decompose lets them decompose, but only when a call reads
such a tensor from outside Core ML and the op has no out variant. Ops
like roll, which have a kernel, stay whole on the CPU, and a uint8
tensor made inside Core ML does not change how the graph decomposes.

test_coreml_partitioner.py gains three tests. A uint8 and an int8 camera
frame lower with the permute and the cast on the CPU and the float ops
in Core ML, and lower_full_graph=True names the dtype. A uint8 frame
read by pad, max_pool2d or roll lowers and serializes. A uint8 tensor
made inside the delegate keeps the float reflect pad whole. Each test
fails without the part of the change it covers.
---
 .../coreml/partition/coreml_partitioner.py    | 127 ++++++++++++++++--
 .../coreml/test/test_coreml_partitioner.py    |  95 +++++++++++++
 2 files changed, 210 insertions(+), 12 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 4453818f956..db25f148061 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -4,11 +4,12 @@
 
 import inspect
 import logging
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, Collection, List, Optional, Set, Tuple
 
 import coremltools as ct
 
 import torch
+from coremltools.converters.mil.frontend.torch.utils import TORCH_DTYPE_TO_MIL_DTYPE
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.compiler.coreml_preprocess import (
@@ -25,8 +26,9 @@
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.operator.convert import to_out_variant
 from torch.export.exported_program import ExportedProgram
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
 
 logger = logging.getLogger(__name__)
@@ -154,6 +156,51 @@ def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool:
     return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals)
 
 
+def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
+    """Whether node makes a tensor that Core ML cannot take as a program input.
+
+    coremltools types each program input through ``TORCH_DTYPE_TO_MIL_DTYPE``
+    and raises ``KeyError`` for a dtype missing from it, such as uint8 or int8,
+    instead of reporting the op as unsupported.
+    """
+    val = node.meta.get("val")
+    vals = val if isinstance(val, (list, tuple)) else [val]
+    return any(
+        isinstance(v, torch.Tensor) and v.dtype not in TORCH_DTYPE_TO_MIL_DTYPE
+        for v in vals
+    )
+
+
+def _find_readers_of_unsupported_inputs(
+    partitions: List[Partition], kept_inside: Collection[str]
+) -> Set[torch.fx.Node]:
+    """The nodes that read a tensor Core ML cannot take as an input from outside
+    their partition.
+
+    A tensor made inside the partition, or a weight, constant or state the
+    partitioner moves into Core ML, is not an input.
+    """
+    return {
+        node
+        for partition in partitions
+        for node in partition.nodes
+        if any(
+            arg not in partition.nodes
+            and arg.name not in kept_inside
+            and _cannot_be_coreml_input(arg)
+            for arg in node.all_input_nodes
+        )
+    }
+
+
+def _has_out_variant(op: torch._ops.OpOverload) -> bool:
+    try:
+        to_out_variant(op)
+    except RuntimeError:
+        return False
+    return True
+
+
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
     schema = op._schema
     if len(schema.arguments) == 0:
@@ -176,6 +223,7 @@ def __init__(
         self.lower_full_graph = lower_full_graph
         self._logged_msgs = set()
         self._log = log
+        self.unsupported_input_readers: Set[torch.fx.Node] = set()
 
     def log_once(self, msg: str) -> None:
         if self._log and msg not in self._logged_msgs:
@@ -286,6 +334,20 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             )
             if self.should_override_support(node):
                 is_supported = False
+            elif node in self.unsupported_input_readers:
+                msg = (
+                    f"{node_target_name} reads a tensor whose dtype CoreML cannot "
+                    "take as an input, such as uint8 or int8, from outside its "
+                    "CoreML partition."
+                )
+                if self.lower_full_graph:
+                    raise NotImplementedError(
+                        f"{msg}  Convert the tensor to a supported dtype before "
+                        "it enters the model, or set lower_full_graph=False in the "
+                        "CoreMLPartitioner to run this op outside of CoreML."
+                    )
+                self.log_once(f"{msg}  Overriding op support.")
+                is_supported = False
 
             if not is_supported:
                 if self.lower_full_graph:
@@ -393,16 +455,24 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
-        capability_partitioner = CapabilityBasedPartitioner(
-            exported_program.graph_module,
-            _OperatorsSupportedForCoreMLBackend(
-                self.skip_ops_for_coreml_delegation,
-                self.lower_full_graph,
-                log=True,
-            ),
-            allows_single_node_partition=True,
+        op_support = _OperatorsSupportedForCoreMLBackend(
+            self.skip_ops_for_coreml_delegation,
+            self.lower_full_graph,
+            log=True,
         )
-        partition_list = capability_partitioner.propose_partitions()
+        kept_inside = self._placeholders_kept_inside(exported_program)
+        # Each op moved out of Core ML can hand a tensor Core ML cannot take as
+        # an input to the next op, so repeat until no partition reads one.
+        while True:
+            partition_list = CapabilityBasedPartitioner(
+                exported_program.graph_module,
+                op_support,
+                allows_single_node_partition=True,
+            ).propose_partitions()
+            readers = _find_readers_of_unsupported_inputs(partition_list, kept_inside)
+            if not readers:
+                break
+            op_support.unsupported_input_readers |= readers
         for partition in partition_list:
             for node in partition.nodes:
                 tag = f"tag{partition.id}"
@@ -430,6 +500,21 @@ def log_once(self, msg: str) -> None:
             logging.info(msg)
             self._logged_msgs.add(msg)
 
+    def _placeholders_kept_inside(self, ep: ExportedProgram) -> Set[str]:
+        """The weights, constants and state this partitioner moves into Core ML."""
+        signature = ep.graph_signature
+        mutated = set(signature.buffers_to_mutate.values())
+        kept = set()
+        if self.take_over_constant_data:
+            kept |= set(signature.inputs_to_parameters)
+            kept |= set(signature.inputs_to_lifted_tensor_constants)
+        for name, buffer in signature.inputs_to_buffers.items():
+            if (buffer in mutated and self.take_over_mutable_buffer) or (
+                buffer not in mutated and self.take_over_constant_data
+            ):
+                kept.add(name)
+        return kept
+
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
@@ -453,6 +538,13 @@ def ops_to_not_decompose(
             torch.ops.aten.repeat_interleave.self_Tensor,
         ]
         rejected = set()
+        kept_inside = self._placeholders_kept_inside(ep)
+        # The nodes whose output stays outside Core ML.
+        outside = {
+            node
+            for node in ep.graph.nodes
+            if node.op == "placeholder" and node.name not in kept_inside
+        }
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
                 node.target, torch._ops.OpOverload
@@ -466,7 +558,18 @@ def ops_to_not_decompose(
                         f"Encountered exception when checking node support, treating node as unsupported: {e}"
                     )
                     supported = False
-                if not supported:
+                # partition() keeps an op that reads a tensor Core ML cannot take
+                # as an input on the CPU, where it stays whole only if it has an
+                # out variant.
+                reads_unsupported_input = any(
+                    arg in outside and _cannot_be_coreml_input(arg)
+                    for arg in node.all_input_nodes
+                )
+                if not supported or reads_unsupported_input:
+                    outside.add(node)
+                if not supported or (
+                    reads_unsupported_input and not _has_out_variant(node.target)
+                ):
                     rejected.add(node.target)
                 elif node.target not in do_not_decompose_blocklist and not _is_view_op(
                     node.target
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 0129eb1d33f..cbdfeae8bed 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -521,6 +521,101 @@ def forward(self, q):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2))
 
+    def test_uint8_camera_frame_is_read_outside_the_delegate(self):
+        """
+        Core ML cannot take a uint8 or int8 input, so the ops that read such a
+        camera frame stay on the CPU up to its conversion to float, and the
+        float ops after it are still delegated.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, image):
+                return torch.relu(image.permute(2, 0, 1).float() / 255.0 - 0.5)
+
+        model = Model().eval()
+        for dtype in (torch.uint8, torch.int8):
+            with self.subTest(dtype=dtype):
+                x = torch.randint(0, 100, (8, 8, 3), dtype=dtype)
+                ep = torch.export.export(model, (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                self.assertIn("aten.permute_copy.default", op_names)
+                self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names)
+                self.assertIn("executorch_call_delegate", op_names)
+                for name in ("aten.div.Tensor", "aten.sub.Tensor", "aten.relu.default"):
+                    self.assertNotIn(name, op_names)
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(edge.to_executorch().buffer)
+                    out = program.load_method("forward").execute((x,))[0]
+                    self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+
+        with self.assertRaisesRegex(NotImplementedError, "dtype CoreML cannot take"):
+            executorch.exir.to_edge_transform_and_lower(
+                ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)]
+            )
+
+    def test_uint8_input_read_by_ops_with_and_without_cpu_kernels(self):
+        """
+        An op that reads a uint8 input runs on the CPU.  Pad and max pooling
+        have no CPU kernel of their own, so they are decomposed.  Roll has one,
+        so it stays whole, since its decomposition would leave Core ML a
+        partition with no inputs.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self, op):
+                super().__init__()
+                self.op = op
+
+            def forward(self, image):
+                return self.op(image).float() / 255.0
+
+        ops = {
+            "pad": lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1)),
+            "max_pool2d": lambda x: torch.nn.functional.max_pool2d(x, 2),
+            "roll": lambda x: torch.roll(x, 1, 2),
+        }
+        x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)
+        for name, op in ops.items():
+            with self.subTest(op=name):
+                model = Model(op).eval()
+                ep = torch.export.export(model, (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(buffer)
+                    out = program.load_method("forward").execute((x,))[0]
+                    self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+
+    def test_uint8_made_inside_the_delegate_keeps_ops_whole(self):
+        """
+        A uint8 tensor made inside the delegate is not an input, so a pad that
+        reads it does not make the float reflect pad decompose.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                image = torch.nn.functional.pad((x * 255).to(torch.uint8), (1, 1))
+                return torch.nn.functional.pad(
+                    image.float() / 255.0, (1, 1), mode="reflect"
+                )
+
+        ep = torch.export.export(Model().eval(), (torch.rand(1, 3, 8),), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        delegated = format_delegated_graph(edge.exported_program().graph_module)
+        self.assertIn("aten.pad.default", delegated)
+        self.assertNotIn("aten.index.Tensor", delegated)
+
     def test_tied_embedding_is_quantized_with_its_linear(self):
         """A weight that is both an embedding table and an output projection lowers.
 

From f3a811f30842b8267e2f603c34bdc1a45d9df47e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 20:44:45 -0700
Subject: [PATCH 02/44] Share the uint8 input reader search in the Core ML
 partitioner

ops_to_not_decompose guessed which ops partition() would keep on the CPU
for reading a uint8 or int8 input, using its own walk over the graph.
The walk did not follow tensors through getitem, and it did not see
tensors that cross from a Core ML partition to another. So a uint8 frame
read through max(dim=1), split or unbind and then padded kept aten.pad
whole, and to_executorch() failed with "Missing out variants:
{'aten::pad'}".

ops_to_not_decompose and partition() now call a shared helper that
proposes partitions, finds the ops that read such a tensor from outside
their partition, and repeats until there are none. The readers are
passed to the partitioner as an extra check, so the operator support
object no longer has a public set that is changed between rounds. When
an op is decomposed because a call of it reads such a tensor, the
partitioner logs it.

The tests now check what they describe: where each op runs, with inputs
that cover the whole uint8 and int8 range so the eager comparison is not
all zeros, and the lower_full_graph error for uint8 and int8. They also
cover max(dim=1) and split before a pad, and a uint8 buffer that Core ML
does not take over. The max and split cases fail without this change,
and the buffer case fails if a buffer that is not taken over is treated
as taken over.
---
 .../coreml/partition/coreml_partitioner.py    | 133 ++++++++++--------
 .../coreml/test/test_coreml_partitioner.py    |  83 +++++++++--
 2 files changed, 142 insertions(+), 74 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index db25f148061..c405760ba77 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -29,7 +29,11 @@
 from executorch.exir.operator.convert import to_out_variant
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
-from torch.fx.passes.operator_support import OperatorSupportBase
+from torch.fx.passes.operator_support import (
+    chain,
+    create_op_support,
+    OperatorSupportBase,
+)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(get_coreml_log_level(default_level=logging.INFO))
@@ -193,6 +197,30 @@ def _find_readers_of_unsupported_inputs(
     }
 
 
+def _propose_partitions(
+    graph_module: torch.fx.GraphModule,
+    op_support: OperatorSupportBase,
+    kept_inside: Collection[str],
+) -> Tuple[List[Partition], Set[torch.fx.Node]]:
+    """Propose Core ML partitions, leaving out every op that reads a tensor
+    Core ML cannot take as an input from outside its partition.
+
+    Returns the partitions and the ops left out this way.
+    """
+    readers: Set[torch.fx.Node] = set()
+    support = chain(create_op_support(lambda _, node: node not in readers), op_support)
+    # Each op moved out of Core ML can hand such a tensor to the next op, so
+    # repeat until no partition reads one.
+    while True:
+        partitions = CapabilityBasedPartitioner(
+            graph_module, support, allows_single_node_partition=True
+        ).propose_partitions()
+        new_readers = _find_readers_of_unsupported_inputs(partitions, kept_inside)
+        if not new_readers:
+            return partitions, readers
+        readers |= new_readers
+
+
 def _has_out_variant(op: torch._ops.OpOverload) -> bool:
     try:
         to_out_variant(op)
@@ -223,7 +251,6 @@ def __init__(
         self.lower_full_graph = lower_full_graph
         self._logged_msgs = set()
         self._log = log
-        self.unsupported_input_readers: Set[torch.fx.Node] = set()
 
     def log_once(self, msg: str) -> None:
         if self._log and msg not in self._logged_msgs:
@@ -334,20 +361,6 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             )
             if self.should_override_support(node):
                 is_supported = False
-            elif node in self.unsupported_input_readers:
-                msg = (
-                    f"{node_target_name} reads a tensor whose dtype CoreML cannot "
-                    "take as an input, such as uint8 or int8, from outside its "
-                    "CoreML partition."
-                )
-                if self.lower_full_graph:
-                    raise NotImplementedError(
-                        f"{msg}  Convert the tensor to a supported dtype before "
-                        "it enters the model, or set lower_full_graph=False in the "
-                        "CoreMLPartitioner to run this op outside of CoreML."
-                    )
-                self.log_once(f"{msg}  Overriding op support.")
-                is_supported = False
 
             if not is_supported:
                 if self.lower_full_graph:
@@ -455,24 +468,27 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
-        op_support = _OperatorsSupportedForCoreMLBackend(
-            self.skip_ops_for_coreml_delegation,
-            self.lower_full_graph,
-            log=True,
+        partition_list, readers = _propose_partitions(
+            exported_program.graph_module,
+            _OperatorsSupportedForCoreMLBackend(
+                self.skip_ops_for_coreml_delegation,
+                self.lower_full_graph,
+                log=True,
+            ),
+            self._placeholders_kept_inside(exported_program),
         )
-        kept_inside = self._placeholders_kept_inside(exported_program)
-        # Each op moved out of Core ML can hand a tensor Core ML cannot take as
-        # an input to the next op, so repeat until no partition reads one.
-        while True:
-            partition_list = CapabilityBasedPartitioner(
-                exported_program.graph_module,
-                op_support,
-                allows_single_node_partition=True,
-            ).propose_partitions()
-            readers = _find_readers_of_unsupported_inputs(partition_list, kept_inside)
-            if not readers:
-                break
-            op_support.unsupported_input_readers |= readers
+        for node in readers:
+            msg = (
+                f"{node.target} reads a tensor whose dtype CoreML cannot take as "
+                "an input, such as uint8 or int8, from outside its CoreML partition."
+            )
+            if self.lower_full_graph:
+                raise NotImplementedError(
+                    f"{msg}  Convert the tensor to a supported dtype before it "
+                    "enters the model, or set lower_full_graph=False in the "
+                    "CoreMLPartitioner to run this op outside of CoreML."
+                )
+            self.log_once(f"{msg}  Running it outside of CoreML.")
         for partition in partition_list:
             for node in partition.nodes:
                 tag = f"tag{partition.id}"
@@ -537,39 +553,38 @@ def ops_to_not_decompose(
             torch.ops.aten.repeat_interleave.self_int,
             torch.ops.aten.repeat_interleave.self_Tensor,
         ]
+
+        def is_supported(submodules, node: torch.fx.Node) -> bool:
+            try:
+                return op_support.is_node_supported(submodules, node)
+            except Exception as e:
+                # CoreML's op_support.is_node_supported will sometimes throw
+                # for unsupported ops, rather than returning False
+                self.log_once(
+                    f"Encountered exception when checking node support, treating node as unsupported: {e}"
+                )
+                return False
+
+        # partition() keeps an op that reads a tensor Core ML cannot take as an
+        # input on the CPU, where it stays whole only if it has an out variant.
+        _, readers = _propose_partitions(
+            ep.graph_module,
+            create_op_support(is_supported),
+            self._placeholders_kept_inside(ep),
+        )
         rejected = set()
-        kept_inside = self._placeholders_kept_inside(ep)
-        # The nodes whose output stays outside Core ML.
-        outside = {
-            node
-            for node in ep.graph.nodes
-            if node.op == "placeholder" and node.name not in kept_inside
-        }
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
                 node.target, torch._ops.OpOverload
             ):
-                try:
-                    supported = op_support.is_node_supported(None, node)
-                except Exception as e:
-                    # CoreML's op_support.is_node_supported will sometimes throw
-                    # for unsupported ops, rather than returning False
+                if not is_supported(None, node):
+                    rejected.add(node.target)
+                elif node in readers and not _has_out_variant(node.target):
                     self.log_once(
-                        f"Encountered exception when checking node support, treating node as unsupported: {e}"
+                        f"{node.target} reads a tensor CoreML cannot take as an "
+                        "input and has no out variant, so every call of it is "
+                        "decomposed."
                     )
-                    supported = False
-                # partition() keeps an op that reads a tensor Core ML cannot take
-                # as an input on the CPU, where it stays whole only if it has an
-                # out variant.
-                reads_unsupported_input = any(
-                    arg in outside and _cannot_be_coreml_input(arg)
-                    for arg in node.all_input_nodes
-                )
-                if not supported or reads_unsupported_input:
-                    outside.add(node)
-                if not supported or (
-                    reads_unsupported_input and not _has_out_variant(node.target)
-                ):
                     rejected.add(node.target)
                 elif node.target not in do_not_decompose_blocklist and not _is_view_op(
                     node.target
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index cbdfeae8bed..4c138544d28 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -530,12 +530,13 @@ def test_uint8_camera_frame_is_read_outside_the_delegate(self):
 
         class Model(torch.nn.Module):
             def forward(self, image):
-                return torch.relu(image.permute(2, 0, 1).float() / 255.0 - 0.5)
+                return torch.relu(image.permute(2, 0, 1).float() / 127.0 - 0.25)
 
         model = Model().eval()
         for dtype in (torch.uint8, torch.int8):
             with self.subTest(dtype=dtype):
-                x = torch.randint(0, 100, (8, 8, 3), dtype=dtype)
+                info = torch.iinfo(dtype)
+                x = torch.randint(info.min, info.max + 1, (8, 8, 3), dtype=dtype)
                 ep = torch.export.export(model, (x,), strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
                     ep, partitioner=[CoreMLPartitioner()]
@@ -554,18 +555,20 @@ def forward(self, image):
                     program = Runtime.get().load_program(edge.to_executorch().buffer)
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
-
-        with self.assertRaisesRegex(NotImplementedError, "dtype CoreML cannot take"):
-            executorch.exir.to_edge_transform_and_lower(
-                ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)]
-            )
+                with self.assertRaisesRegex(
+                    NotImplementedError, "dtype CoreML cannot take"
+                ):
+                    executorch.exir.to_edge_transform_and_lower(
+                        ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)]
+                    )
 
     def test_uint8_input_read_by_ops_with_and_without_cpu_kernels(self):
         """
-        An op that reads a uint8 input runs on the CPU.  Pad and max pooling
-        have no CPU kernel of their own, so they are decomposed.  Roll has one,
-        so it stays whole, since its decomposition would leave Core ML a
-        partition with no inputs.
+        An op that reads a uint8 input runs on the CPU, also when it reads it
+        through an op with several outputs.  Pad and max pooling have no CPU
+        kernel of their own, so they are decomposed.  Roll has one, so it stays
+        whole, since its decomposition would leave Core ML a partition with no
+        inputs.
         """
 
         class Model(torch.nn.Module):
@@ -576,25 +579,75 @@ def __init__(self, op):
             def forward(self, image):
                 return self.op(image).float() / 255.0
 
+        # Each op, and the op that has to run on the CPU in its place.
         ops = {
-            "pad": lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1)),
-            "max_pool2d": lambda x: torch.nn.functional.max_pool2d(x, 2),
-            "roll": lambda x: torch.roll(x, 1, 2),
+            "pad": (
+                lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1)),
+                "aten.constant_pad_nd.default",
+            ),
+            "max_pool2d": (
+                lambda x: torch.nn.functional.max_pool2d(x, 2),
+                "aten.max_pool2d_with_indices.default",
+            ),
+            "roll": (lambda x: torch.roll(x, 1, 2), "aten.roll.default"),
+            "max_then_pad": (
+                lambda x: torch.nn.functional.pad(x.max(dim=1).values, (1, 1)),
+                "aten.constant_pad_nd.default",
+            ),
+            "split_then_pad": (
+                lambda x: torch.nn.functional.pad(x.split(1, dim=1)[0], (1, 1)),
+                "aten.constant_pad_nd.default",
+            ),
         }
         x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)
-        for name, op in ops.items():
+        for name, (op, cpu_op) in ops.items():
             with self.subTest(op=name):
                 model = Model(op).eval()
                 ep = torch.export.export(model, (x,), strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
                     ep, partitioner=[CoreMLPartitioner()]
                 )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                self.assertIn(cpu_op, op_names)
+                self.assertIn("executorch_call_delegate", op_names)
                 buffer = edge.to_executorch().buffer
                 if _TEST_RUNTIME:
                     program = Runtime.get().load_program(buffer)
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
 
+    def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self):
+        """
+        A uint8 buffer that Core ML does not take over is a program input, so
+        the op that reads it runs on the CPU.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer(
+                    "lut", torch.randint(0, 256, (8,), dtype=torch.uint8)
+                )
+
+            def forward(self, x):
+                return x * (self.lut.float() / 255.0)
+
+        model = Model().eval()
+        x = torch.randn(2, 8)
+        ep = torch.export.export(model, (x,), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner(take_over_constant_data=False)]
+        )
+        buffer = edge.to_executorch().buffer
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+
     def test_uint8_made_inside_the_delegate_keeps_ops_whole(self):
         """
         A uint8 tensor made inside the delegate is not an input, so a pad that

From c5de614cfbf1070ea36bb561c970a7ba6aa9f78b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 21:22:03 -0700
Subject: [PATCH 03/44] Handle uint8 inputs that reach a gather or a non-core
 op in Core ML

Some uint8 or int8 inputs still failed after their readers moved to the
CPU:

- A channel swap, x[:, [2, 1, 0]], or a reflect or replicate pad moved
  the gather to the CPU and left Core ML a partition that reads nothing
  at runtime: only the index math, or a copy of the index constant.
  Below iOS 18 Core ML rejects such a partition with "Deploying a model
  with no inputs in CoreML requires ... iOS18".
- zeros_like and ones_like on such an input have an out variant, so they
  were kept whole on the CPU, but no portable kernel implements them,
  and the program failed to load with "kernel 'aten::zeros_like.out' not
  found".

Once an op has been moved out for reading such a tensor, the partitioner
now also keeps on the CPU every partition that reads no tensor except
the weights, constants and state it takes over. Such a partition reads
only symbolic sizes, which Core ML cannot take as inputs, or nothing,
which Core ML rejects below iOS 18. From iOS 18 a partition with no
inputs stays delegated, which also keeps any state it updates in Core
ML. A graph without such tensors is not affected.

An op left on the CPU this way is now never preserved. A core ATen op
stays whole anyway, and any other op, such as pad, max_pool2d, roll or
zeros_like, is decomposed into core ATen ops, the op set the portable
kernels target. Before, an op with an out variant was kept whole, but an
out variant does not mean a kernel exists. Not preserving core ops also
matters: masked_fill decomposes into a scalar_tensor that runs on the
CPU, and preserving it would make lowering fail with "was not decomposed
or delegated". Roll used to stay whole to avoid the empty partition its
decomposition leaves. That partition now runs on the CPU, so roll is
decomposed like the others. The decision is per op, so a float roll
elsewhere in the model is decomposed too, and the partitioner logs each
op it decomposes this way.

test_uint8_input_read_by_ops_with_and_without_cpu_kernels gains a
reflect pad, a channel swap and zeros_like, which fail without this
change, and masked_fill. roll and masked_fill now check that their
decomposed pieces run on the CPU. New tests cover a roll with a dynamic
height and a state update at iOS 18. The eager comparisons use
atol=1e-2, since Core ML computes in float16 and 1e-3 failed for some
random inputs. The partitioning and decomposition of models without
uint8 or int8 tensors are unchanged, checked on a transformer block and
on a conv model with and without skipped ops.
---
 .../coreml/partition/coreml_partitioner.py    | 109 +++++++++++++-----
 .../coreml/test/test_coreml_partitioner.py    |  90 +++++++++++++--
 2 files changed, 159 insertions(+), 40 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index c405760ba77..c39c3beb7eb 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -26,7 +26,6 @@
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.operator.convert import to_out_variant
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import (
@@ -197,36 +196,67 @@ def _find_readers_of_unsupported_inputs(
     }
 
 
+def _makes_tensor(node: torch.fx.Node) -> bool:
+    val = node.meta.get("val")
+    vals = val if isinstance(val, (list, tuple)) else [val]
+    return any(isinstance(v, torch.Tensor) for v in vals)
+
+
+def _find_nodes_of_partitions_without_tensor_inputs(
+    partitions: List[Partition],
+    kept_inside: Collection[str],
+    allow_partitions_without_inputs: bool,
+) -> Set[torch.fx.Node]:
+    """The nodes of the partitions that read no tensor from outside, except the
+    weights, constants and state the partitioner moves into Core ML.
+
+    Core ML cannot take a symbolic size as an input, and below iOS 18 it
+    rejects a model with no inputs, so such a partition cannot be delegated.
+    """
+    nodes = set()
+    for partition in partitions:
+        inputs = [
+            arg
+            for node in partition.nodes
+            for arg in node.all_input_nodes
+            if arg not in partition.nodes and arg.name not in kept_inside
+        ]
+        reads_tensor = any(_makes_tensor(arg) for arg in inputs)
+        if not reads_tensor and (inputs or not allow_partitions_without_inputs):
+            nodes.update(partition.nodes)
+    return nodes
+
+
 def _propose_partitions(
     graph_module: torch.fx.GraphModule,
     op_support: OperatorSupportBase,
     kept_inside: Collection[str],
+    allow_partitions_without_inputs: bool,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
     """Propose Core ML partitions, leaving out every op that reads a tensor
-    Core ML cannot take as an input from outside its partition.
+    Core ML cannot take as an input from outside its partition, and then every
+    partition this leaves without a tensor input.
 
-    Returns the partitions and the ops left out this way.
+    Returns the partitions and the ops left out.
     """
-    readers: Set[torch.fx.Node] = set()
-    support = chain(create_op_support(lambda _, node: node not in readers), op_support)
+    left_out: Set[torch.fx.Node] = set()
+    support = chain(create_op_support(lambda _, node: node not in left_out), op_support)
     # Each op moved out of Core ML can hand such a tensor to the next op, so
     # repeat until no partition reads one.
     while True:
         partitions = CapabilityBasedPartitioner(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
-        new_readers = _find_readers_of_unsupported_inputs(partitions, kept_inside)
-        if not new_readers:
-            return partitions, readers
-        readers |= new_readers
-
-
-def _has_out_variant(op: torch._ops.OpOverload) -> bool:
-    try:
-        to_out_variant(op)
-    except RuntimeError:
-        return False
-    return True
+        new = _find_readers_of_unsupported_inputs(partitions, kept_inside)
+        if new or left_out:
+            # Moving those ops out can leave a partition that reads no tensor,
+            # such as the index math of a gather.
+            new |= _find_nodes_of_partitions_without_tensor_inputs(
+                partitions, kept_inside, allow_partitions_without_inputs
+            )
+        if not new:
+            return partitions, left_out
+        left_out |= new
 
 
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
@@ -468,7 +498,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
-        partition_list, readers = _propose_partitions(
+        partition_list, left_out = _propose_partitions(
             exported_program.graph_module,
             _OperatorsSupportedForCoreMLBackend(
                 self.skip_ops_for_coreml_delegation,
@@ -476,19 +506,24 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 log=True,
             ),
             self._placeholders_kept_inside(exported_program),
+            self._allows_partitions_without_inputs(),
         )
-        for node in readers:
+        if left_out:
             msg = (
-                f"{node.target} reads a tensor whose dtype CoreML cannot take as "
-                "an input, such as uint8 or int8, from outside its CoreML partition."
+                "The model reads a tensor whose dtype CoreML cannot take as an "
+                "input, such as uint8 or int8, from outside a CoreML partition."
             )
             if self.lower_full_graph:
                 raise NotImplementedError(
                     f"{msg}  Convert the tensor to a supported dtype before it "
                     "enters the model, or set lower_full_graph=False in the "
-                    "CoreMLPartitioner to run this op outside of CoreML."
+                    "CoreMLPartitioner to run the ops that read it outside of "
+                    "CoreML."
                 )
-            self.log_once(f"{msg}  Running it outside of CoreML.")
+            self.log_once(
+                f"{msg}  Running these ops outside of CoreML: "
+                + ", ".join(sorted(str(node.target) for node in left_out))
+            )
         for partition in partition_list:
             for node in partition.nodes:
                 tag = f"tag{partition.id}"
@@ -516,6 +551,13 @@ def log_once(self, msg: str) -> None:
             logging.info(msg)
             self._logged_msgs.add(msg)
 
+    def _allows_partitions_without_inputs(self) -> bool:
+        """Whether Core ML accepts a partition with no inputs, which needs iOS 18."""
+        target = CoreMLBackend.min_deployment_target_from_compile_specs(
+            self.delegation_spec.compile_specs
+        )
+        return target is not None and target >= ct.target.iOS18
+
     def _placeholders_kept_inside(self, ep: ExportedProgram) -> Set[str]:
         """The weights, constants and state this partitioner moves into Core ML."""
         signature = ep.graph_signature
@@ -565,12 +607,14 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 )
                 return False
 
-        # partition() keeps an op that reads a tensor Core ML cannot take as an
-        # input on the CPU, where it stays whole only if it has an out variant.
-        _, readers = _propose_partitions(
+        # An op that partition() leaves out runs on the CPU, so it is never
+        # preserved: a core ATen op stays whole anyway, and any other op is
+        # decomposed into core ATen ops, the op set the portable kernels target.
+        _, left_out = _propose_partitions(
             ep.graph_module,
             create_op_support(is_supported),
             self._placeholders_kept_inside(ep),
+            self._allows_partitions_without_inputs(),
         )
         rejected = set()
         for node in ep.graph.nodes:
@@ -579,12 +623,13 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
             ):
                 if not is_supported(None, node):
                     rejected.add(node.target)
-                elif node in readers and not _has_out_variant(node.target):
-                    self.log_once(
-                        f"{node.target} reads a tensor CoreML cannot take as an "
-                        "input and has no out variant, so every call of it is "
-                        "decomposed."
-                    )
+                elif node in left_out:
+                    if torch.Tag.core not in node.target.tags:
+                        self.log_once(
+                            f"{node.target} runs outside of CoreML here because "
+                            "of a tensor CoreML cannot take as an input, so every "
+                            "call of it is decomposed."
+                        )
                     rejected.add(node.target)
                 elif node.target not in do_not_decompose_blocklist and not _is_view_op(
                     node.target
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 4c138544d28..1da4272cfad 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -554,7 +554,7 @@ def forward(self, image):
                 if _TEST_RUNTIME:
                     program = Runtime.get().load_program(edge.to_executorch().buffer)
                     out = program.load_method("forward").execute((x,))[0]
-                    self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+                    self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
                 with self.assertRaisesRegex(
                     NotImplementedError, "dtype CoreML cannot take"
                 ):
@@ -565,10 +565,10 @@ def forward(self, image):
     def test_uint8_input_read_by_ops_with_and_without_cpu_kernels(self):
         """
         An op that reads a uint8 input runs on the CPU, also when it reads it
-        through an op with several outputs.  Pad and max pooling have no CPU
-        kernel of their own, so they are decomposed.  Roll has one, so it stays
-        whole, since its decomposition would leave Core ML a partition with no
-        inputs.
+        through an op with several outputs.  An op that is not core ATen, such
+        as pad, max pooling, roll or zeros_like, may have no CPU kernel, so it
+        is decomposed.  The index math of a gather also runs on the CPU, since
+        it would leave Core ML a partition with no inputs.
         """
 
         class Model(torch.nn.Module):
@@ -589,7 +589,17 @@ def forward(self, image):
                 lambda x: torch.nn.functional.max_pool2d(x, 2),
                 "aten.max_pool2d_with_indices.default",
             ),
-            "roll": (lambda x: torch.roll(x, 1, 2), "aten.roll.default"),
+            "reflect_pad": (
+                lambda x: torch.nn.functional.pad(x, (1, 1, 1, 1), mode="reflect"),
+                "aten.arange.start_step",
+            ),
+            "roll": (lambda x: torch.roll(x, 1, 2), "aten.arange.start_step"),
+            "zeros_like": (
+                lambda x: torch.where(x > 100, x, torch.zeros_like(x)),
+                "aten.full_like.default",
+            ),
+            "channel_swap": (lambda x: x[:, [2, 1, 0]], "aten.index.Tensor"),
+            "masked_fill": (lambda x: x.masked_fill(x < 10, 0), "aten.where.self"),
             "max_then_pad": (
                 lambda x: torch.nn.functional.pad(x.max(dim=1).values, (1, 1)),
                 "aten.constant_pad_nd.default",
@@ -618,7 +628,71 @@ def forward(self, image):
                 if _TEST_RUNTIME:
                     program = Runtime.get().load_program(buffer)
                     out = program.load_method("forward").execute((x,))[0]
-                    self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+                    self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
+
+    def test_uint8_input_with_dynamic_shape(self):
+        """
+        With a dynamic height, the index math of a decomposed roll reads only
+        a symbolic size, which Core ML cannot take as an input, so it runs on
+        the CPU too.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, image):
+                return torch.roll(image, 1, 2).float() / 255.0
+
+        model = Model().eval()
+        x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)
+        ep = torch.export.export(
+            model,
+            (x,),
+            dynamic_shapes=({2: torch.export.Dim("height", min=4, max=64)},),
+            strict=True,
+        )
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        buffer = edge.to_executorch().buffer
+        if _TEST_RUNTIME:
+            x = torch.randint(0, 256, (1, 3, 12, 8), dtype=torch.uint8)
+            program = Runtime.get().load_program(buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
+
+    def test_uint8_input_keeps_state_in_the_delegate_from_ios18(self):
+        """
+        From iOS 18 Core ML takes a partition with no inputs, so a partition
+        that only updates state stays in the delegate with that state instead
+        of moving to the CPU, which would leave the delegate a stale copy.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("frames", torch.zeros(1))
+
+            def forward(self, image):
+                out = image.float() / 255.0 * self.frames
+                frames = self.frames + 1
+                self.frames.copy_(frames)
+                # acosh runs on the CPU, so the update gets its own partition.
+                return out * torch.acosh(frames + 1)
+
+        x = torch.randint(0, 256, (1, 3, 2, 2), dtype=torch.uint8)
+        ep = torch.export.export(Model().eval(), (x,), strict=True)
+        compile_specs = CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=ct.target.iOS18
+        )
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertNotIn("aten.add.Tensor", op_names)
+        self.assertIn("aten.acosh.default", op_names)
 
     def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self):
         """
@@ -646,7 +720,7 @@ def forward(self, x):
         if _TEST_RUNTIME:
             program = Runtime.get().load_program(buffer)
             out = program.load_method("forward").execute((x,))[0]
-            self.assertTrue(torch.allclose(out, model(x), atol=1e-3))
+            self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
     def test_uint8_made_inside_the_delegate_keeps_ops_whole(self):
         """

From 296aab7c4321b5000279191bb21395677e821f7c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 23:49:15 -0700
Subject: [PATCH 04/44] Keep a Core ML partition that uses state in the
 delegate

Below iOS 18 the partitioner moves a partition with no inputs to the
CPU, because Core ML refuses such a model there. A partition that only
updates state, such as a frame counter next to a uint8 camera input, has
no inputs either, so it moved to the CPU. The state's other readers stayed
in Core ML with a frozen copy of it, and from the second call on the
model returned wrong results with no error.

Never move a partition that uses state. Core ML state needs iOS 18
anyway, so below it the model now fails at lowering with Core ML's own
message, as any model with state does. The test checks that error at the
default deployment target.
---
 .../coreml/partition/coreml_partitioner.py    | 40 ++++++++++++++-----
 .../coreml/test/test_coreml_partitioner.py    | 13 ++++--
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index c39c3beb7eb..a251334bbf8 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -205,6 +205,7 @@ def _makes_tensor(node: torch.fx.Node) -> bool:
 def _find_nodes_of_partitions_without_tensor_inputs(
     partitions: List[Partition],
     kept_inside: Collection[str],
+    state: Collection[str],
     allow_partitions_without_inputs: bool,
 ) -> Set[torch.fx.Node]:
     """The nodes of the partitions that read no tensor from outside, except the
@@ -212,15 +213,21 @@ def _find_nodes_of_partitions_without_tensor_inputs(
 
     Core ML cannot take a symbolic size as an input, and below iOS 18 it
     rejects a model with no inputs, so such a partition cannot be delegated.
+    A partition that uses state stays: moving it to the CPU would leave the
+    state's other readers in Core ML a stale copy, and Core ML state needs
+    iOS 18 anyway.
     """
     nodes = set()
     for partition in partitions:
-        inputs = [
+        outside = [
             arg
             for node in partition.nodes
             for arg in node.all_input_nodes
-            if arg not in partition.nodes and arg.name not in kept_inside
+            if arg not in partition.nodes
         ]
+        if any(arg.name in state for arg in outside):
+            continue
+        inputs = [arg for arg in outside if arg.name not in kept_inside]
         reads_tensor = any(_makes_tensor(arg) for arg in inputs)
         if not reads_tensor and (inputs or not allow_partitions_without_inputs):
             nodes.update(partition.nodes)
@@ -231,6 +238,7 @@ def _propose_partitions(
     graph_module: torch.fx.GraphModule,
     op_support: OperatorSupportBase,
     kept_inside: Collection[str],
+    state: Collection[str],
     allow_partitions_without_inputs: bool,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
     """Propose Core ML partitions, leaving out every op that reads a tensor
@@ -252,7 +260,7 @@ def _propose_partitions(
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
             new |= _find_nodes_of_partitions_without_tensor_inputs(
-                partitions, kept_inside, allow_partitions_without_inputs
+                partitions, kept_inside, state, allow_partitions_without_inputs
             )
         if not new:
             return partitions, left_out
@@ -506,6 +514,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 log=True,
             ),
             self._placeholders_kept_inside(exported_program),
+            self._state_kept_inside(exported_program),
             self._allows_partitions_without_inputs(),
         )
         if left_out:
@@ -562,17 +571,29 @@ def _placeholders_kept_inside(self, ep: ExportedProgram) -> Set[str]:
         """The weights, constants and state this partitioner moves into Core ML."""
         signature = ep.graph_signature
         mutated = set(signature.buffers_to_mutate.values())
-        kept = set()
+        kept = self._state_kept_inside(ep)
         if self.take_over_constant_data:
             kept |= set(signature.inputs_to_parameters)
             kept |= set(signature.inputs_to_lifted_tensor_constants)
-        for name, buffer in signature.inputs_to_buffers.items():
-            if (buffer in mutated and self.take_over_mutable_buffer) or (
-                buffer not in mutated and self.take_over_constant_data
-            ):
-                kept.add(name)
+            kept |= {
+                name
+                for name, buffer in signature.inputs_to_buffers.items()
+                if buffer not in mutated
+            }
         return kept
 
+    def _state_kept_inside(self, ep: ExportedProgram) -> Set[str]:
+        """The mutable buffers this partitioner moves into Core ML as state."""
+        if not self.take_over_mutable_buffer:
+            return set()
+        signature = ep.graph_signature
+        mutated = set(signature.buffers_to_mutate.values())
+        return {
+            name
+            for name, buffer in signature.inputs_to_buffers.items()
+            if buffer in mutated
+        }
+
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
@@ -614,6 +635,7 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
             ep.graph_module,
             create_op_support(is_supported),
             self._placeholders_kept_inside(ep),
+            self._state_kept_inside(ep),
             self._allows_partitions_without_inputs(),
         )
         rejected = set()
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 1da4272cfad..eacf8ca9419 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -659,11 +659,12 @@ def forward(self, image):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
-    def test_uint8_input_keeps_state_in_the_delegate_from_ios18(self):
+    def test_uint8_input_keeps_state_in_the_delegate(self):
         """
-        From iOS 18 Core ML takes a partition with no inputs, so a partition
-        that only updates state stays in the delegate with that state instead
-        of moving to the CPU, which would leave the delegate a stale copy.
+        A partition that only updates state stays in the delegate with that
+        state instead of moving to the CPU, which would leave the delegate a
+        stale copy.  Below iOS 18 Core ML then refuses the state at lowering,
+        as it does for any model with state.
         """
 
         class Model(torch.nn.Module):
@@ -680,6 +681,10 @@ def forward(self, image):
 
         x = torch.randint(0, 256, (1, 3, 2, 2), dtype=torch.uint8)
         ep = torch.export.export(Model().eval(), (x,), strict=True)
+        with self.assertRaisesRegex(ValueError, "State model is supported only"):
+            executorch.exir.to_edge_transform_and_lower(
+                ep, partitioner=[CoreMLPartitioner()]
+            )
         compile_specs = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
         )

From 70bf4c8d4119768cf986efd4144331109c17ff2e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 23:52:41 -0700
Subject: [PATCH 05/44] Run Core ML ops that read a size from a CPU op on the
 CPU too

With a dynamic uint8 input, the op that reads the input's size runs on
the CPU, since it reads the uint8 input. A delegated op that used that
size, such as a reshape to image.shape[2], then read a symbolic size
from outside Core ML, and lowering failed with "Placeholder val must be a
tensor or fake tensor, but got type <class 'torch.SymInt'>". A float
computed from it, such as image.shape[2] / 8, failed the same way.

Treat a symbolic value made by an op left out of Core ML like a tensor
Core ML cannot take, so the ops that read it run on the CPU too. The
dynamic shape test now also reshapes to such a size and scales by such a
float, and checks where the ops run without needing the macOS runtime.

Such an op can also read state that Core ML takes over, for example a
slice of it to that size. On the CPU it would read a stale copy of the
state, so lowering now refuses it and suggests
take_over_mutable_buffer=False, with which the same model lowers and
runs.
---
 .../coreml/partition/coreml_partitioner.py    | 31 +++++++--
 .../coreml/test/test_coreml_partitioner.py    | 66 ++++++++++++++++++-
 2 files changed, 90 insertions(+), 7 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index a251334bbf8..ebfeee38fbd 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -175,10 +175,12 @@ def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
 
 
 def _find_readers_of_unsupported_inputs(
-    partitions: List[Partition], kept_inside: Collection[str]
+    partitions: List[Partition],
+    kept_inside: Collection[str],
+    left_out: Collection[torch.fx.Node],
 ) -> Set[torch.fx.Node]:
     """The nodes that read a tensor Core ML cannot take as an input from outside
-    their partition.
+    their partition, or a symbolic value computed by an op left out of Core ML.
 
     A tensor made inside the partition, or a weight, constant or state the
     partitioner moves into Core ML, is not an input.
@@ -190,7 +192,16 @@ def _find_readers_of_unsupported_inputs(
         if any(
             arg not in partition.nodes
             and arg.name not in kept_inside
-            and _cannot_be_coreml_input(arg)
+            and (
+                _cannot_be_coreml_input(arg)
+                or (
+                    arg in left_out
+                    and isinstance(
+                        arg.meta.get("val"),
+                        (torch.SymInt, torch.SymFloat, torch.SymBool),
+                    )
+                )
+            )
             for arg in node.all_input_nodes
         )
     }
@@ -255,7 +266,7 @@ def _propose_partitions(
         partitions = CapabilityBasedPartitioner(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
-        new = _find_readers_of_unsupported_inputs(partitions, kept_inside)
+        new = _find_readers_of_unsupported_inputs(partitions, kept_inside, left_out)
         if new or left_out:
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
@@ -506,6 +517,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
+        state = self._state_kept_inside(exported_program)
         partition_list, left_out = _propose_partitions(
             exported_program.graph_module,
             _OperatorsSupportedForCoreMLBackend(
@@ -514,9 +526,18 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 log=True,
             ),
             self._placeholders_kept_inside(exported_program),
-            self._state_kept_inside(exported_program),
+            state,
             self._allows_partitions_without_inputs(),
         )
+        for node in left_out:
+            read = [arg.name for arg in node.all_input_nodes if arg.name in state]
+            if read:
+                raise NotImplementedError(
+                    f"{node} has to run outside of CoreML, but it reads "
+                    f"{', '.join(read)}, which CoreML takes over as state, so it "
+                    "would read a stale copy.  Set take_over_mutable_buffer=False "
+                    "in the CoreMLPartitioner to keep the state outside of CoreML."
+                )
         if left_out:
             msg = (
                 "The model reads a tensor whose dtype CoreML cannot take as an "
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index eacf8ca9419..1af7aa4b44e 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -634,12 +634,15 @@ def test_uint8_input_with_dynamic_shape(self):
         """
         With a dynamic height, the index math of a decomposed roll reads only
         a symbolic size, which Core ML cannot take as an input, so it runs on
-        the CPU too.
+        the CPU too.  So do a reshape to a size read from the uint8 input, and
+        a scale computed from it.
         """
 
         class Model(torch.nn.Module):
             def forward(self, image):
-                return torch.roll(image, 1, 2).float() / 255.0
+                rolled = torch.roll(image, 1, 2).float() / 255.0
+                scaled = rolled * (image.shape[2] / 8)
+                return scaled.reshape(1, 3 * image.shape[2], 8)
 
         model = Model().eval()
         x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)
@@ -652,6 +655,14 @@ def forward(self, image):
         edge = executorch.exir.to_edge_transform_and_lower(
             ep, partitioner=[CoreMLPartitioner()]
         )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        for name in ("aten.arange.start_step", "aten.view_copy.default"):
+            self.assertIn(name, op_names)
+        self.assertIn("executorch_call_delegate", op_names)
         buffer = edge.to_executorch().buffer
         if _TEST_RUNTIME:
             x = torch.randint(0, 256, (1, 3, 12, 8), dtype=torch.uint8)
@@ -699,6 +710,57 @@ def forward(self, image):
         self.assertNotIn("aten.add.Tensor", op_names)
         self.assertIn("aten.acosh.default", op_names)
 
+    def test_uint8_input_size_that_slices_state(self):
+        """
+        A slice of state to a size read from a uint8 input has to run on the
+        CPU, where it would read a stale copy of state that Core ML takes over,
+        so lowering refuses it unless the state stays outside of Core ML.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("hist", torch.zeros(64))
+
+            def forward(self, image, y):
+                out = y * self.hist + self.hist[: image.shape[2]].sum()
+                self.hist.add_(1)
+                return out
+
+        model = Model().eval()
+        inputs = (
+            torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8),
+            torch.randn(64),
+        )
+        ep = torch.export.export(
+            model,
+            inputs,
+            dynamic_shapes=({2: torch.export.Dim("height", min=2, max=64)}, None),
+            strict=True,
+        )
+        compile_specs = CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=ct.target.iOS18
+        )
+        with self.assertRaisesRegex(NotImplementedError, "take_over_mutable_buffer"):
+            executorch.exir.to_edge_transform_and_lower(
+                ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+            )
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep,
+            partitioner=[
+                CoreMLPartitioner(
+                    compile_specs=compile_specs, take_over_mutable_buffer=False
+                )
+            ],
+        )
+        if _TEST_RUNTIME:
+            method = Runtime.get().load_program(edge.to_executorch().buffer)
+            method = method.load_method("forward")
+            eager = Model().eval()
+            for _ in range(2):
+                out = method.execute(inputs)[0]
+                self.assertTrue(torch.allclose(out, eager(*inputs), atol=1e-2))
+
     def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self):
         """
         A uint8 buffer that Core ML does not take over is a program input, so

From 4d60df11b53e7dd89207d9f2cfbbff76a6d4f486 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 23:52:41 -0700
Subject: [PATCH 06/44] Check the uint8 Core ML tests on Linux and from iOS 18
 too

The uint8 buffer test only checked results, under the macOS runtime, so
on Linux it only checked that lowering does not fail. It now also checks
that the op reading the buffer runs on the CPU.

The test of ops that read a uint8 input now also lowers and runs the
gather cases with an iOS 18 deployment target, and checks that their
index math, which has no inputs, then stays in Core ML.
---
 .../coreml/test/test_coreml_partitioner.py    | 32 +++++++++++++++----
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 1af7aa4b44e..e4aac754b26 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -567,8 +567,8 @@ def test_uint8_input_read_by_ops_with_and_without_cpu_kernels(self):
         An op that reads a uint8 input runs on the CPU, also when it reads it
         through an op with several outputs.  An op that is not core ATen, such
         as pad, max pooling, roll or zeros_like, may have no CPU kernel, so it
-        is decomposed.  The index math of a gather also runs on the CPU, since
-        it would leave Core ML a partition with no inputs.
+        is decomposed.  Below iOS 18 the index math of a gather also runs on
+        the CPU, since it would leave Core ML a partition with no inputs.
         """
 
         class Model(torch.nn.Module):
@@ -610,19 +610,32 @@ def forward(self, image):
             ),
         }
         x = torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8)
-        for name, (op, cpu_op) in ops.items():
-            with self.subTest(op=name):
+        ios18 = CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=ct.target.iOS18
+        )
+        # From iOS 18 index math with no inputs stays in Core ML instead.
+        cases = [(name, op, cpu_op, None) for name, (op, cpu_op) in ops.items()]
+        cases += [
+            (name, op, cpu_op, ios18)
+            for name, (op, cpu_op) in ops.items()
+            if cpu_op == "aten.arange.start_step"
+        ]
+        for name, op, cpu_op, compile_specs in cases:
+            with self.subTest(op=name, ios18=compile_specs is not None):
                 model = Model(op).eval()
                 ep = torch.export.export(model, (x,), strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
-                    ep, partitioner=[CoreMLPartitioner()]
+                    ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
                 )
                 op_names = [
                     n.target.__name__
                     for n in edge.exported_program().graph.nodes
                     if n.op == "call_function"
                 ]
-                self.assertIn(cpu_op, op_names)
+                if compile_specs is None:
+                    self.assertIn(cpu_op, op_names)
+                else:
+                    self.assertNotIn(cpu_op, op_names)
                 self.assertIn("executorch_call_delegate", op_names)
                 buffer = edge.to_executorch().buffer
                 if _TEST_RUNTIME:
@@ -783,6 +796,13 @@ def forward(self, x):
         edge = executorch.exir.to_edge_transform_and_lower(
             ep, partitioner=[CoreMLPartitioner(take_over_constant_data=False)]
         )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names)
+        self.assertIn("executorch_call_delegate", op_names)
         buffer = edge.to_executorch().buffer
         if _TEST_RUNTIME:
             program = Runtime.get().load_program(buffer)

From f7aafb8d29905ad75dfa5fdbd33fc3a2ae025d26 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 23:53:38 -0700
Subject: [PATCH 07/44] Name the uint8 input in the Core ML partitioner
 messages

When the partitioner runs ops outside Core ML because they read a uint8
or int8 tensor, its log listed the full repr of each op once per node,
and the lower_full_graph=True error did not say which input it was
about. Both now name the input and its dtype, and the log lists each op
name once.

CoreMLPartitioner.log_once also logs through the module's logger rather
than the root logger, so ET_COREML_LOG_LEVEL controls these messages like
the partitioner's other ones.
---
 .../coreml/partition/coreml_partitioner.py    | 29 ++++++++++++++++---
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index ebfeee38fbd..de8ba15cbdb 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -517,6 +517,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         logger.info("CoreMLPartitioner::partition")
         partition_tags = {}
 
+        kept_inside = self._placeholders_kept_inside(exported_program)
         state = self._state_kept_inside(exported_program)
         partition_list, left_out = _propose_partitions(
             exported_program.graph_module,
@@ -525,7 +526,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 self.lower_full_graph,
                 log=True,
             ),
-            self._placeholders_kept_inside(exported_program),
+            kept_inside,
             state,
             self._allows_partitions_without_inputs(),
         )
@@ -539,9 +540,22 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                     "in the CoreMLPartitioner to keep the state outside of CoreML."
                 )
         if left_out:
+            inputs = sorted(
+                {
+                    f"{arg.name} ({arg.meta['val'].dtype})"
+                    for node in left_out
+                    for arg in node.all_input_nodes
+                    if arg not in left_out
+                    and arg.name not in kept_inside
+                    and isinstance(arg.meta.get("val"), torch.Tensor)
+                    and _cannot_be_coreml_input(arg)
+                }
+            )
             msg = (
                 "The model reads a tensor whose dtype CoreML cannot take as an "
-                "input, such as uint8 or int8, from outside a CoreML partition."
+                "input, such as uint8 or int8, from outside a CoreML partition: "
+                + ", ".join(inputs)
+                + "."
             )
             if self.lower_full_graph:
                 raise NotImplementedError(
@@ -552,7 +566,14 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 )
             self.log_once(
                 f"{msg}  Running these ops outside of CoreML: "
-                + ", ".join(sorted(str(node.target) for node in left_out))
+                + ", ".join(
+                    sorted(
+                        {
+                            getattr(node.target, "__name__", str(node.target))
+                            for node in left_out
+                        }
+                    )
+                )
             )
         for partition in partition_list:
             for node in partition.nodes:
@@ -578,7 +599,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
     def log_once(self, msg: str) -> None:
         if msg not in self._logged_msgs:
-            logging.info(msg)
+            logger.info(msg)
             self._logged_msgs.add(msg)
 
     def _allows_partitions_without_inputs(self) -> bool:

From 6e2012c72fc7a3b8ba3c8887053493dfbbfbbb6f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Thu, 29 Oct 2026 23:53:38 -0700
Subject: [PATCH 08/44] Skip the extra Core ML partitioning when no tensor
 needs it

ops_to_not_decompose partitioned the graph a second time to find the ops
that read a uint8 or int8 tensor, even in a graph without one, and it is
called several times per export. Ops are only left out because of such a
tensor, so skip that work when the graph has none.
---
 .../coreml/partition/coreml_partitioner.py     | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index de8ba15cbdb..a003621209e 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -673,13 +673,17 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
         # An op that partition() leaves out runs on the CPU, so it is never
         # preserved: a core ATen op stays whole anyway, and any other op is
         # decomposed into core ATen ops, the op set the portable kernels target.
-        _, left_out = _propose_partitions(
-            ep.graph_module,
-            create_op_support(is_supported),
-            self._placeholders_kept_inside(ep),
-            self._state_kept_inside(ep),
-            self._allows_partitions_without_inputs(),
-        )
+        # Ops are only left out because of a tensor Core ML cannot take as an
+        # input, so a graph without one needs no extra partitioning.
+        left_out = set()
+        if any(_cannot_be_coreml_input(node) for node in ep.graph.nodes):
+            _, left_out = _propose_partitions(
+                ep.graph_module,
+                create_op_support(is_supported),
+                self._placeholders_kept_inside(ep),
+                self._state_kept_inside(ep),
+                self._allows_partitions_without_inputs(),
+            )
         rejected = set()
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(

From 48b4f89835ad9b484bd21555b9ecc4d174515758 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 06:11:18 -0700
Subject: [PATCH 09/44] Leave a Core ML min or max over a bool tensor on the
 CPU

coremltools 9.0 computes min and max only over fp16, fp32 and int32. For
a bool input it raises "expects tensor or scalar of dtype from type domain
['fp16', 'fp32', 'int32'] but got tensor[...,bool]" while converting,
instead of reporting the op as unsupported, so lowering failed. The
argument check of torch.distributions.MultivariateNormal is one such
reduction.

The partitioner now refuses min, max, amin, amax, minimum and maximum
over a bool tensor, like it refuses bfloat16 ops, so they run on the
portable kernels and the rest is still delegated. any, all, sum, cumsum
and logical_not over bool convert and match eager, so they stay
delegated.
---
 .../coreml/partition/coreml_partitioner.py    | 43 +++++++++++++++++++
 .../coreml/test/test_coreml_partitioner.py    | 34 +++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index a003621209e..da32e0db902 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -159,6 +159,42 @@ def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool:
     return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals)
 
 
+_MIN_MAX_TARGETS = (
+    torch.ops.aten.min.default,
+    torch.ops.aten.min.dim,
+    torch.ops.aten.min.other,
+    torch.ops.aten.max.default,
+    torch.ops.aten.max.dim,
+    torch.ops.aten.max.other,
+    torch.ops.aten.amin.default,
+    torch.ops.aten.amax.default,
+    torch.ops.aten.minimum.default,
+    torch.ops.aten.maximum.default,
+    exir_ops.edge.aten.min.default,
+    exir_ops.edge.aten.min.dim,
+    exir_ops.edge.aten.min.other,
+    exir_ops.edge.aten.max.default,
+    exir_ops.edge.aten.max.dim,
+    exir_ops.edge.aten.max.other,
+    exir_ops.edge.aten.amin.default,
+    exir_ops.edge.aten.amax.default,
+    exir_ops.edge.aten.minimum.default,
+    exir_ops.edge.aten.maximum.default,
+)
+
+
+def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
+    """coremltools 9.0 computes min and max only over fp16, fp32 and int32.
+
+    It raises ``ValueError`` for a bool input, such as the all-true check that
+    torch.distributions exports, instead of reporting the op as unsupported.
+    """
+    if node.target not in _MIN_MAX_TARGETS:
+        return False
+    val = node.args[0].meta.get("val")
+    return isinstance(val, torch.Tensor) and val.dtype == torch.bool
+
+
 def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
     """Whether node makes a tensor that Core ML cannot take as a program input.
 
@@ -373,6 +409,13 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _is_min_max_of_bool(node):
+            self.log_once(
+                f"{node.target} over a bool tensor is not supported by CoreML.  "
+                "Overriding op support."
+            )
+            return True
+
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
         # # in the placeholders due to partitioning, which CoreML does not support
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index e4aac754b26..936d467135a 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -521,6 +521,40 @@ def forward(self, q):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2))
 
+    def test_min_max_of_bool_is_skipped(self):
+        """
+        coremltools raises on a min or max over a bool tensor, such as the
+        all-true check torch.distributions exports.  The partitioner must leave
+        those ops to the portable backend and still delegate the rest.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                ok = (x.tril() == x).view(1, -1)
+                flags = ok.min(-1)[0].to(x.dtype) + ok.amax(-1).to(x.dtype)
+                both = torch.maximum(ok, ok.flip(-1)).view(x.shape).to(x.dtype)
+                return x * 2 + flags + both
+
+        model = Model().eval()
+        x = torch.rand(14, 14)
+        ep = torch.export.export(model, (x,), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertIn("aten.min.dim", op_names)
+        self.assertIn("aten.amax.default", op_names)
+        self.assertIn("aten.maximum.default", op_names)
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(edge.to_executorch().buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
+
     def test_uint8_camera_frame_is_read_outside_the_delegate(self):
         """
         Core ML cannot take a uint8 or int8 input, so the ops that read such a

From 1cca6059f54d1a9ae7d87e4e353b59250d47465b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 06:13:19 -0700
Subject: [PATCH 10/44] Do not keep a Core ML op whole when its decomposed
 calls are rejected

ops_to_not_decompose decides per op, from the graph before
decomposition, which ops to keep whole. One float32 matmul that Core ML
takes, such as a RoPE frequency product, kept aten.matmul whole. A
bfloat16 linear is rejected and decomposed, and that makes a new matmul
call, which Core ML rejects too. It then stayed whole outside Core ML,
and the program failed to load with "kernel 'aten::matmul.out' not
found", since the portable kernels have no matmul.

When some call is rejected and some op will be decomposed, the
partitioner now decomposes the program as lowering will, with its kept
ops, and gives back every non-core kept op that has a rejected call
there. Core ops stay kept, as before. exir's per-node filter for kept
ops was not used: it switches the program to the EDGE_DO_NOT_DECOMP
path, which the MLX partitioner notes fails on ordinary attention
blocks.
---
 .../coreml/partition/coreml_partitioner.py    | 37 ++++++++++++++++++-
 .../coreml/test/test_coreml_partitioner.py    | 37 +++++++++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index da32e0db902..a1541a090fa 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -4,7 +4,7 @@
 
 import inspect
 import logging
-from typing import Callable, Collection, List, Optional, Set, Tuple
+from typing import Any, Callable, Collection, List, Optional, Set, Tuple
 
 import coremltools as ct
 
@@ -26,6 +26,7 @@
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.tracer import _default_decomposition_table
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import (
@@ -314,6 +315,35 @@ def _propose_partitions(
         left_out |= new
 
 
+def _drop_ops_rejected_after_decomposition(
+    ep: ExportedProgram,
+    keep: List[torch._ops.OpOverload],
+    is_supported: Callable[[Any, torch.fx.Node], bool],
+) -> List[torch._ops.OpOverload]:
+    """Decomposing a rejected call can make new calls of a kept op, such as the
+    matmul of a bfloat16 linear, that are rejected too. A non-core op may have
+    no CPU kernel, so such ops go back to decomposition.
+    """
+    while any(torch.Tag.core not in op.tags for op in keep):
+        table = _default_decomposition_table()
+        for op in keep:
+            table.pop(op, None)
+        # Only a decomposition can make new calls.
+        if not any(node.target in table for node in ep.graph.nodes):
+            break
+        late = {
+            node.target
+            for node in ep.run_decompositions(table).graph.nodes
+            if node.target in keep
+            and torch.Tag.core not in node.target.tags
+            and not is_supported(None, node)
+        }
+        if not late:
+            break
+        keep = [op for op in keep if op not in late]
+    return keep
+
+
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
     schema = op._schema
     if len(schema.arguments) == 0:
@@ -748,4 +778,7 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                     do_not_decompose.append(node.target)
         # Without a per-node filter, a preserved op stays whole at every call, so
         # one rejected call would be left with no kernel.
-        return [op for op in do_not_decompose if op not in rejected], None
+        keep = [op for op in do_not_decompose if op not in rejected]
+        if rejected:
+            keep = _drop_ops_rejected_after_decomposition(ep, keep, is_supported)
+        return keep, None
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 936d467135a..039e72ce00f 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -521,6 +521,43 @@ def forward(self, q):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2))
 
+    def test_op_made_by_decomposing_a_rejected_op_is_not_kept_whole(self):
+        """
+        A float32 matmul keeps aten.matmul whole.  A bfloat16 linear is
+        rejected and decomposes into another matmul, which Core ML rejects too.
+        That matmul must be decomposed as well, since the portable kernels
+        have no matmul, or the program fails to load.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = torch.nn.Linear(8, 8, bias=False).to(torch.bfloat16)
+                self.register_buffer("inv_freq", torch.rand(1, 4, 1))
+
+            def forward(self, x, pos):
+                freqs = self.inv_freq @ pos[:, None, :].float()
+                return self.proj(x).float(), freqs.cos()
+
+        model = Model().eval()
+        inputs = (torch.rand(1, 6, 8, dtype=torch.bfloat16), torch.arange(6)[None])
+        ep = torch.export.export(model, inputs, strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertNotIn("aten.matmul.default", op_names)
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(edge.to_executorch().buffer)
+            outputs = program.load_method("forward").execute(inputs)
+            for out, ref in zip(outputs, model(*inputs)):
+                self.assertTrue(torch.allclose(out, ref, atol=1e-2))
+
     def test_min_max_of_bool_is_skipped(self):
         """
         coremltools raises on a min or max over a bool tensor, such as the

From 26385c4f2b3b3d038eb8705d691eb60dd549120a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:04:37 -0700
Subject: [PATCH 11/44] Move only uint8 and int8 Core ML inputs to the CPU

The partitioner moved to the CPU every op that reads a tensor of a dtype
coremltools cannot type as a program input. That also covered uint16,
uint32, uint64 and complex, which the portable kernels do not all
handle. A uint16 depth frame read as depth.float() / 1000 then exported,
and failed at the first run with "Unhandled dtype UInt16 for
dim_order_ops::_to_dim_order_copy.out", where it used to fail at export.

Move only uint8 and int8, the camera and quantized dtypes the portable
kernels handle. Other dtypes fail at export again, as before.
---
 .../apple/coreml/partition/coreml_partitioner.py     | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index a1541a090fa..1e119f78b08 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -9,7 +9,6 @@
 import coremltools as ct
 
 import torch
-from coremltools.converters.mil.frontend.torch.utils import TORCH_DTYPE_TO_MIL_DTYPE
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.compiler.coreml_preprocess import (
@@ -197,16 +196,19 @@ def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
 
 
 def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
-    """Whether node makes a tensor that Core ML cannot take as a program input.
+    """Whether node makes a uint8 or int8 tensor, which Core ML cannot take as a
+    program input.
 
     coremltools types each program input through ``TORCH_DTYPE_TO_MIL_DTYPE``
-    and raises ``KeyError`` for a dtype missing from it, such as uint8 or int8,
-    instead of reporting the op as unsupported.
+    and raises ``KeyError`` for a dtype missing from it instead of reporting the
+    op as unsupported. Only uint8 and int8 move to the CPU: the portable kernels
+    handle them, while other missing dtypes, such as uint16, would only fail
+    later, at run time.
     """
     val = node.meta.get("val")
     vals = val if isinstance(val, (list, tuple)) else [val]
     return any(
-        isinstance(v, torch.Tensor) and v.dtype not in TORCH_DTYPE_TO_MIL_DTYPE
+        isinstance(v, torch.Tensor) and v.dtype in (torch.uint8, torch.int8)
         for v in vals
     )
 

From 386fa63dc3585cb9d27d6ae9405db8de382b4b01 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:05:47 -0700
Subject: [PATCH 12/44] Keep an op that makes an int8 tensor for the CPU out of
 Core ML

The Core ML runtime cannot return an int8 tensor ("DataType=1 is not
supported"). Moving the readers of an int8 input to the CPU could leave
an op that makes an int8 tensor for them in Core ML, such as the cast in
torch.maximum(image, x.to(torch.int8)), so the model exported and then
failed at run time. A plain float32 model that returns an int8 tensor
failed the same way.

The partitioner now also runs on the CPU an op that makes an int8 tensor
read outside its partition.
---
 .../coreml/partition/coreml_partitioner.py    | 21 +++++++++++++
 .../coreml/test/test_coreml_partitioner.py    | 31 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 1e119f78b08..5c76427984e 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -246,6 +246,26 @@ def _find_readers_of_unsupported_inputs(
     }
 
 
+def _find_makers_of_int8_outputs(partitions: List[Partition]) -> Set[torch.fx.Node]:
+    """The nodes that make an int8 tensor read outside their partition.
+
+    Such a tensor becomes a Core ML output, and the Core ML runtime cannot
+    return int8.
+    """
+    return {
+        node
+        for partition in partitions
+        for node in partition.nodes
+        if _makes_int8(node) and any(u not in partition.nodes for u in node.users)
+    }
+
+
+def _makes_int8(node: torch.fx.Node) -> bool:
+    val = node.meta.get("val")
+    vals = val if isinstance(val, (list, tuple)) else [val]
+    return any(isinstance(v, torch.Tensor) and v.dtype == torch.int8 for v in vals)
+
+
 def _makes_tensor(node: torch.fx.Node) -> bool:
     val = node.meta.get("val")
     vals = val if isinstance(val, (list, tuple)) else [val]
@@ -306,6 +326,7 @@ def _propose_partitions(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
         new = _find_readers_of_unsupported_inputs(partitions, kept_inside, left_out)
+        new |= _find_makers_of_int8_outputs(partitions)
         if new or left_out:
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 039e72ce00f..f9fd23ad1c4 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -592,6 +592,37 @@ def forward(self, x):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
+    def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
+        """
+        The Core ML runtime cannot return an int8 tensor, so an op that makes
+        one that is read outside Core ML, here the model output, runs on the
+        CPU.
+        """
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return (x * 2).to(torch.int8)
+
+        model = Model().eval()
+        # Small integers, which Core ML's float16 math keeps exact.
+        x = torch.randint(-50, 50, (1, 3, 4, 4)).float()
+        ep = torch.export.export(model, (x,), strict=True)
+        edge = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner()]
+        )
+        op_names = [
+            n.target.__name__
+            for n in edge.exported_program().graph.nodes
+            if n.op == "call_function"
+        ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names)
+        buffer = edge.to_executorch().buffer
+        if _TEST_RUNTIME:
+            program = Runtime.get().load_program(buffer)
+            out = program.load_method("forward").execute((x,))[0]
+            self.assertTrue(torch.equal(out, model(x)))
+
     def test_uint8_camera_frame_is_read_outside_the_delegate(self):
         """
         Core ML cannot take a uint8 or int8 input, so the ops that read such a

From 891f97c2e4138161520652fc8cbf05631e5da888 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:11:58 -0700
Subject: [PATCH 13/44] Predict Core ML CPU ops for decomposition from uint8
 and int8 inputs only

ops_to_not_decompose predicts on the graph before decomposition which
ops partition() will run on the CPU, so it can decompose the ones that
are not core ATen. It treated every uint8 tensor at a partition
boundary as one Core ML could not take. But before decomposition a
boundary can still go away: for a float input cast to uint8, then
rot90 and roll, rot90 is outside Core ML only until it decomposes into
ops Core ML takes. The roll was predicted to run on the CPU and
decomposed, and lowering failed in both modes, where main lowered the
whole model as one delegate.

The prediction now follows only uint8 and int8 graph inputs, through
the ops that run outside Core ML. partition() still checks every
boundary of the decomposed graph. The made-inside test also covers the
rot90 and roll case and checks that the top-level graph only calls the
delegate.
---
 .../coreml/partition/coreml_partitioner.py    | 58 ++++++++++++++++---
 .../coreml/test/test_coreml_partitioner.py    | 39 +++++++++----
 2 files changed, 78 insertions(+), 19 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 5c76427984e..4cbf1e429eb 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -217,12 +217,14 @@ def _find_readers_of_unsupported_inputs(
     partitions: List[Partition],
     kept_inside: Collection[str],
     left_out: Collection[torch.fx.Node],
+    sources: Optional[Collection[torch.fx.Node]] = None,
 ) -> Set[torch.fx.Node]:
     """The nodes that read a tensor Core ML cannot take as an input from outside
     their partition, or a symbolic value computed by an op left out of Core ML.
 
     A tensor made inside the partition, or a weight, constant or state the
-    partitioner moves into Core ML, is not an input.
+    partitioner moves into Core ML, is not an input. With sources, only tensors
+    made by those nodes count.
     """
     return {
         node
@@ -231,6 +233,7 @@ def _find_readers_of_unsupported_inputs(
         if any(
             arg not in partition.nodes
             and arg.name not in kept_inside
+            and (sources is None or arg in sources)
             and (
                 _cannot_be_coreml_input(arg)
                 or (
@@ -266,6 +269,24 @@ def _makes_int8(node: torch.fx.Node) -> bool:
     return any(isinstance(v, torch.Tensor) and v.dtype == torch.int8 for v in vals)
 
 
+def _cpu_nodes_from_graph_inputs(
+    graph: torch.fx.Graph, partitions: List[Partition], kept_inside: Collection[str]
+) -> Set[torch.fx.Node]:
+    """uint8 and int8 graph inputs, and the nodes outside Core ML that read
+    them, directly or through other such nodes."""
+    in_partition = {node for partition in partitions for node in partition.nodes}
+    found = set()
+    for node in graph.nodes:
+        if node.op == "placeholder":
+            if node.name not in kept_inside and _cannot_be_coreml_input(node):
+                found.add(node)
+        elif node not in in_partition and any(
+            arg in found for arg in node.all_input_nodes
+        ):
+            found.add(node)
+    return found
+
+
 def _makes_tensor(node: torch.fx.Node) -> bool:
     val = node.meta.get("val")
     vals = val if isinstance(val, (list, tuple)) else [val]
@@ -310,11 +331,17 @@ def _propose_partitions(
     kept_inside: Collection[str],
     state: Collection[str],
     allow_partitions_without_inputs: bool,
+    graph_inputs_only: bool = False,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
     """Propose Core ML partitions, leaving out every op that reads a tensor
     Core ML cannot take as an input from outside its partition, and then every
     partition this leaves without a tensor input.
 
+    graph_inputs_only follows only the uint8 and int8 tensors that come from
+    graph inputs through ops outside Core ML. On a graph before decomposition,
+    any other boundary can still go away, for example when an op Core ML
+    rejects decomposes into ops it takes.
+
     Returns the partitions and the ops left out.
     """
     left_out: Set[torch.fx.Node] = set()
@@ -325,8 +352,17 @@ def _propose_partitions(
         partitions = CapabilityBasedPartitioner(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
-        new = _find_readers_of_unsupported_inputs(partitions, kept_inside, left_out)
-        new |= _find_makers_of_int8_outputs(partitions)
+        if graph_inputs_only:
+            sources = _cpu_nodes_from_graph_inputs(
+                graph_module.graph, partitions, kept_inside
+            )
+            new = _find_readers_of_unsupported_inputs(
+                partitions, kept_inside, left_out, sources
+            )
+        else:
+            new = _find_readers_of_unsupported_inputs(
+                partitions, kept_inside, left_out
+            ) | _find_makers_of_int8_outputs(partitions)
         if new or left_out:
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
@@ -766,19 +802,23 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 )
                 return False
 
-        # An op that partition() leaves out runs on the CPU, so it is never
-        # preserved: a core ATen op stays whole anyway, and any other op is
-        # decomposed into core ATen ops, the op set the portable kernels target.
-        # Ops are only left out because of a tensor Core ML cannot take as an
-        # input, so a graph without one needs no extra partitioning.
+        # An op that partition() leaves out runs on the CPU, so it is not
+        # preserved: it is decomposed into core ATen ops, the op set the
+        # portable kernels target, unless it is one already. Here only uint8
+        # and int8 graph inputs leave ops out, so a graph without one needs no
+        # extra partitioning.
         left_out = set()
-        if any(_cannot_be_coreml_input(node) for node in ep.graph.nodes):
+        if any(
+            node.op == "placeholder" and _cannot_be_coreml_input(node)
+            for node in ep.graph.nodes
+        ):
             _, left_out = _propose_partitions(
                 ep.graph_module,
                 create_op_support(is_supported),
                 self._placeholders_kept_inside(ep),
                 self._state_kept_inside(ep),
                 self._allows_partitions_without_inputs(),
+                graph_inputs_only=True,
             )
         rejected = set()
         for node in ep.graph.nodes:
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index f9fd23ad1c4..9797dc42500 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -913,24 +913,43 @@ def forward(self, x):
 
     def test_uint8_made_inside_the_delegate_keeps_ops_whole(self):
         """
-        A uint8 tensor made inside the delegate is not an input, so a pad that
-        reads it does not make the float reflect pad decompose.
+        A uint8 tensor made inside the delegate is not an input, so the ops
+        that read it stay in the delegate, whole: a pad, and a rot90 that only
+        Core ML's decomposition supports followed by a roll.
         """
 
-        class Model(torch.nn.Module):
+        class Pad(torch.nn.Module):
             def forward(self, x):
                 image = torch.nn.functional.pad((x * 255).to(torch.uint8), (1, 1))
                 return torch.nn.functional.pad(
                     image.float() / 255.0, (1, 1), mode="reflect"
                 )
 
-        ep = torch.export.export(Model().eval(), (torch.rand(1, 3, 8),), strict=True)
-        edge = executorch.exir.to_edge_transform_and_lower(
-            ep, partitioner=[CoreMLPartitioner()]
-        )
-        delegated = format_delegated_graph(edge.exported_program().graph_module)
-        self.assertIn("aten.pad.default", delegated)
-        self.assertNotIn("aten.index.Tensor", delegated)
+        class RotateRoll(torch.nn.Module):
+            def forward(self, x):
+                image = torch.rot90((x * 255).to(torch.uint8), 1, (2, 3))
+                return torch.roll(image, 1, 2).float() / 255.0
+
+        for model, x in (
+            (Pad(), torch.rand(1, 3, 8)),
+            (RotateRoll(), torch.rand(1, 3, 8, 8)),
+        ):
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                op_names = {
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                }
+                self.assertEqual(op_names, {"executorch_call_delegate", "getitem"})
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(buffer)
+                    out = program.load_method("forward").execute((x,))[0]
+                    self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
     def test_tied_embedding_is_quantized_with_its_linear(self):
         """A weight that is both an embedding table and an output projection lowers.

From 04b00c589730c03d5428c1d23d30d0465a24a6f5 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:15:21 -0700
Subject: [PATCH 14/44] Make the uint8 Core ML tests check more without the
 macOS runtime

- The camera frame test serializes the program on every platform, like
  the other tests, so a missing CPU kernel shows up on Linux too.
- The dynamic shape test also lowers with an iOS 18 target, where only a
  partition with no inputs at all may stay in Core ML, not one that
  reads only a size.
- The state test splits the update off with a skipped op instead of
  acosh, which runs on the CPU only until coremltools supports it. It
  also checks that a delegate exists and serializes the program, and its
  docstring no longer claims the reader of the state gets a live copy.
- The test of a uint8 buffer Core ML does not take over says that one it
  takes over still fails to lower.
---
 .../coreml/test/test_coreml_partitioner.py    | 74 ++++++++++++-------
 1 file changed, 47 insertions(+), 27 deletions(-)

diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 9797dc42500..cead6d92093 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -653,8 +653,9 @@ def forward(self, image):
                 self.assertIn("executorch_call_delegate", op_names)
                 for name in ("aten.div.Tensor", "aten.sub.Tensor", "aten.relu.default"):
                     self.assertNotIn(name, op_names)
+                buffer = edge.to_executorch().buffer
                 if _TEST_RUNTIME:
-                    program = Runtime.get().load_program(edge.to_executorch().buffer)
+                    program = Runtime.get().load_program(buffer)
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
                 with self.assertRaisesRegex(
@@ -767,30 +768,37 @@ def forward(self, image):
             dynamic_shapes=({2: torch.export.Dim("height", min=4, max=64)},),
             strict=True,
         )
-        edge = executorch.exir.to_edge_transform_and_lower(
-            ep, partitioner=[CoreMLPartitioner()]
+        ios18 = CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=ct.target.iOS18
         )
-        op_names = [
-            n.target.__name__
-            for n in edge.exported_program().graph.nodes
-            if n.op == "call_function"
-        ]
-        for name in ("aten.arange.start_step", "aten.view_copy.default"):
-            self.assertIn(name, op_names)
-        self.assertIn("executorch_call_delegate", op_names)
-        buffer = edge.to_executorch().buffer
-        if _TEST_RUNTIME:
-            x = torch.randint(0, 256, (1, 3, 12, 8), dtype=torch.uint8)
-            program = Runtime.get().load_program(buffer)
-            out = program.load_method("forward").execute((x,))[0]
-            self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
+        # Also from iOS 18, where only a partition with no inputs at all may
+        # stay in Core ML.
+        for compile_specs in (None, ios18):
+            with self.subTest(ios18=compile_specs is not None):
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+                )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                for name in ("aten.arange.start_step", "aten.view_copy.default"):
+                    self.assertIn(name, op_names)
+                self.assertIn("executorch_call_delegate", op_names)
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    y = torch.randint(0, 256, (1, 3, 12, 8), dtype=torch.uint8)
+                    program = Runtime.get().load_program(buffer)
+                    out = program.load_method("forward").execute((y,))[0]
+                    self.assertTrue(torch.allclose(out, model(y), atol=1e-2))
 
     def test_uint8_input_keeps_state_in_the_delegate(self):
         """
-        A partition that only updates state stays in the delegate with that
-        state instead of moving to the CPU, which would leave the delegate a
-        stale copy.  Below iOS 18 Core ML then refuses the state at lowering,
-        as it does for any model with state.
+        A partition that only updates state, which has no inputs, stays in
+        Core ML with that state instead of moving to the CPU.  Below iOS 18
+        Core ML then refuses the state at lowering, as it does for any model
+        with state.
         """
 
         class Model(torch.nn.Module):
@@ -802,28 +810,38 @@ def forward(self, image):
                 out = image.float() / 255.0 * self.frames
                 frames = self.frames + 1
                 self.frames.copy_(frames)
-                # acosh runs on the CPU, so the update gets its own partition.
-                return out * torch.acosh(frames + 1)
+                # The skipped sqrt runs on the CPU, so the update gets its own
+                # partition.
+                return out * torch.sqrt(frames + 1)
 
+        skip = ["aten.sqrt.default"]
         x = torch.randint(0, 256, (1, 3, 2, 2), dtype=torch.uint8)
         ep = torch.export.export(Model().eval(), (x,), strict=True)
         with self.assertRaisesRegex(ValueError, "State model is supported only"):
             executorch.exir.to_edge_transform_and_lower(
-                ep, partitioner=[CoreMLPartitioner()]
+                ep,
+                partitioner=[CoreMLPartitioner(skip_ops_for_coreml_delegation=skip)],
             )
         compile_specs = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
         )
         edge = executorch.exir.to_edge_transform_and_lower(
-            ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+            ep,
+            partitioner=[
+                CoreMLPartitioner(
+                    compile_specs=compile_specs, skip_ops_for_coreml_delegation=skip
+                )
+            ],
         )
         op_names = [
             n.target.__name__
             for n in edge.exported_program().graph.nodes
             if n.op == "call_function"
         ]
+        self.assertIn("executorch_call_delegate", op_names)
+        self.assertIn("aten.sqrt.default", op_names)
         self.assertNotIn("aten.add.Tensor", op_names)
-        self.assertIn("aten.acosh.default", op_names)
+        edge.to_executorch()
 
     def test_uint8_input_size_that_slices_state(self):
         """
@@ -879,7 +897,9 @@ def forward(self, image, y):
     def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self):
         """
         A uint8 buffer that Core ML does not take over is a program input, so
-        the op that reads it runs on the CPU.
+        the op that reads it runs on the CPU.  A uint8 buffer Core ML takes
+        over, the default, still fails to lower: coremltools cannot cast a
+        uint8 constant.
         """
 
         class Model(torch.nn.Module):

From 7d01d82d0702c4d045a9239937d25edb3b849a81 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:15:51 -0700
Subject: [PATCH 15/44] Log every Core ML op left on the CPU that gets
 decomposed

ops_to_not_decompose logged that an op left on the CPU is decomposed
only when the op is not core ATen. A core op with a default
decomposition, such as reflection_pad1d, is decomposed too, without a
message, and the comment said core ops stay whole. Log it too, and say
in the comment which core ops stay whole.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 4cbf1e429eb..91cc427f496 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -804,9 +804,9 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
 
         # An op that partition() leaves out runs on the CPU, so it is not
         # preserved: it is decomposed into core ATen ops, the op set the
-        # portable kernels target, unless it is one already. Here only uint8
-        # and int8 graph inputs leave ops out, so a graph without one needs no
-        # extra partitioning.
+        # portable kernels target, unless it is a core op with no default
+        # decomposition. Here only uint8 and int8 graph inputs leave ops out,
+        # so a graph without one needs no extra partitioning.
         left_out = set()
         if any(
             node.op == "placeholder" and _cannot_be_coreml_input(node)
@@ -820,6 +820,7 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 self._allows_partitions_without_inputs(),
                 graph_inputs_only=True,
             )
+        decomposed_by_default = _default_decomposition_table() if left_out else {}
         rejected = set()
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
@@ -828,7 +829,9 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 if not is_supported(None, node):
                     rejected.add(node.target)
                 elif node in left_out:
-                    if torch.Tag.core not in node.target.tags:
+                    if torch.Tag.core not in node.target.tags or (
+                        node.target in decomposed_by_default
+                    ):
                         self.log_once(
                             f"{node.target} runs outside of CoreML here because "
                             "of a tensor CoreML cannot take as an input, so every "

From eebf0f683e39767ec150ca6175f2e878cabd4e8a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 07:38:01 -0700
Subject: [PATCH 16/44] Also predict Core ML CPU ops behind ops that stay off
 Core ML

The decomposition prediction followed only uint8 and int8 graph inputs.
A uint8 tensor made inside Core ML and passed through an op that stays
on the CPU, such as diagonal, then reached a pad that partition() later
ran on the CPU. The pad was kept whole, and to_executorch failed with
"Missing out variants: {'aten::pad'}". A pad that makes an int8 model
output failed the same way, because the prediction did not apply the
int8 output rule.

The prediction now also starts from ops outside Core ML that have no
default decomposition, since those stay outside after decomposition,
and applies the int8 output rule to the graph output and to those ops.
Ops with a default decomposition, such as rot90, can still turn into
ops Core ML takes, so their boundaries are still not followed. The
messages now say uint8 or int8, the only dtypes this handles.
---
 .../coreml/partition/coreml_partitioner.py    | 70 +++++++++++--------
 .../coreml/test/test_coreml_partitioner.py    | 42 ++++++++++-
 2 files changed, 82 insertions(+), 30 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 91cc427f496..0af4705bb86 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -249,8 +249,12 @@ def _find_readers_of_unsupported_inputs(
     }
 
 
-def _find_makers_of_int8_outputs(partitions: List[Partition]) -> Set[torch.fx.Node]:
-    """The nodes that make an int8 tensor read outside their partition.
+def _find_makers_of_int8_outputs(
+    partitions: List[Partition],
+    readers: Optional[Collection[torch.fx.Node]] = None,
+) -> Set[torch.fx.Node]:
+    """The nodes that make an int8 tensor read outside their partition, by the
+    graph output or, with readers, by one of those nodes.
 
     Such a tensor becomes a Core ML output, and the Core ML runtime cannot
     return int8.
@@ -259,7 +263,12 @@ def _find_makers_of_int8_outputs(partitions: List[Partition]) -> Set[torch.fx.No
         node
         for partition in partitions
         for node in partition.nodes
-        if _makes_int8(node) and any(u not in partition.nodes for u in node.users)
+        if _makes_int8(node)
+        and any(
+            u not in partition.nodes
+            and (readers is None or u.op == "output" or u in readers)
+            for u in node.users
+        )
     }
 
 
@@ -269,21 +278,27 @@ def _makes_int8(node: torch.fx.Node) -> bool:
     return any(isinstance(v, torch.Tensor) and v.dtype == torch.int8 for v in vals)
 
 
-def _cpu_nodes_from_graph_inputs(
-    graph: torch.fx.Graph, partitions: List[Partition], kept_inside: Collection[str]
+def _lasting_cpu_sources(
+    graph: torch.fx.Graph,
+    partitions: List[Partition],
+    kept_inside: Collection[str],
+    decomposed: Collection[torch._ops.OpOverload],
 ) -> Set[torch.fx.Node]:
-    """uint8 and int8 graph inputs, and the nodes outside Core ML that read
-    them, directly or through other such nodes."""
+    """The nodes whose outputs still come from outside Core ML once the graph is
+    decomposed: uint8 and int8 graph inputs, ops outside Core ML that have no
+    default decomposition, and ops outside Core ML that read such a node.
+    """
     in_partition = {node for partition in partitions for node in partition.nodes}
     found = set()
     for node in graph.nodes:
         if node.op == "placeholder":
             if node.name not in kept_inside and _cannot_be_coreml_input(node):
                 found.add(node)
-        elif node not in in_partition and any(
-            arg in found for arg in node.all_input_nodes
-        ):
-            found.add(node)
+        elif node.op == "call_function" and node not in in_partition:
+            if node.target not in decomposed or any(
+                arg in found for arg in node.all_input_nodes
+            ):
+                found.add(node)
     return found
 
 
@@ -331,19 +346,19 @@ def _propose_partitions(
     kept_inside: Collection[str],
     state: Collection[str],
     allow_partitions_without_inputs: bool,
-    graph_inputs_only: bool = False,
+    before_decomposition: bool = False,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
     """Propose Core ML partitions, leaving out every op that reads a tensor
     Core ML cannot take as an input from outside its partition, and then every
     partition this leaves without a tensor input.
 
-    graph_inputs_only follows only the uint8 and int8 tensors that come from
-    graph inputs through ops outside Core ML. On a graph before decomposition,
-    any other boundary can still go away, for example when an op Core ML
-    rejects decomposes into ops it takes.
+    On a graph before decomposition, a boundary can still go away, for example
+    when an op Core ML rejects decomposes into ops it takes, so then only the
+    boundaries that last are followed.
 
     Returns the partitions and the ops left out.
     """
+    decomposed = _default_decomposition_table() if before_decomposition else {}
     left_out: Set[torch.fx.Node] = set()
     support = chain(create_op_support(lambda _, node: node not in left_out), op_support)
     # Each op moved out of Core ML can hand such a tensor to the next op, so
@@ -352,13 +367,13 @@ def _propose_partitions(
         partitions = CapabilityBasedPartitioner(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
-        if graph_inputs_only:
-            sources = _cpu_nodes_from_graph_inputs(
-                graph_module.graph, partitions, kept_inside
+        if before_decomposition:
+            sources = _lasting_cpu_sources(
+                graph_module.graph, partitions, kept_inside, decomposed
             )
             new = _find_readers_of_unsupported_inputs(
                 partitions, kept_inside, left_out, sources
-            )
+            ) | _find_makers_of_int8_outputs(partitions, sources)
         else:
             new = _find_readers_of_unsupported_inputs(
                 partitions, kept_inside, left_out
@@ -684,8 +699,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 }
             )
             msg = (
-                "The model reads a tensor whose dtype CoreML cannot take as an "
-                "input, such as uint8 or int8, from outside a CoreML partition: "
+                "The model reads a uint8 or int8 tensor, which CoreML cannot take "
+                "as an input, from outside a CoreML partition: "
                 + ", ".join(inputs)
                 + "."
             )
@@ -805,20 +820,17 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
         # An op that partition() leaves out runs on the CPU, so it is not
         # preserved: it is decomposed into core ATen ops, the op set the
         # portable kernels target, unless it is a core op with no default
-        # decomposition. Here only uint8 and int8 graph inputs leave ops out,
-        # so a graph without one needs no extra partitioning.
+        # decomposition. Ops are only left out because of a uint8 or int8
+        # tensor, so a graph without one needs no extra partitioning.
         left_out = set()
-        if any(
-            node.op == "placeholder" and _cannot_be_coreml_input(node)
-            for node in ep.graph.nodes
-        ):
+        if any(_cannot_be_coreml_input(node) for node in ep.graph.nodes):
             _, left_out = _propose_partitions(
                 ep.graph_module,
                 create_op_support(is_supported),
                 self._placeholders_kept_inside(ep),
                 self._state_kept_inside(ep),
                 self._allows_partitions_without_inputs(),
-                graph_inputs_only=True,
+                before_decomposition=True,
             )
         decomposed_by_default = _default_decomposition_table() if left_out else {}
         rejected = set()
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index cead6d92093..f7630a8535f 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -659,7 +659,7 @@ def forward(self, image):
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
                 with self.assertRaisesRegex(
-                    NotImplementedError, "dtype CoreML cannot take"
+                    NotImplementedError, "image \\(torch\\.u?int8\\)"
                 ):
                     executorch.exir.to_edge_transform_and_lower(
                         ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)]
@@ -971,6 +971,46 @@ def forward(self, x):
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
+    def test_ops_after_a_cpu_op_on_a_tensor_made_inside(self):
+        """
+        A uint8 tensor made in Core ML and passed through an op that stays on
+        the CPU, such as diagonal, sends the pad after it to the CPU too.  The
+        pad is decomposed there, since the portable kernels have no pad, and so
+        is a pad that makes an int8 model output.
+        """
+
+        class DiagonalPad(torch.nn.Module):
+            def forward(self, x):
+                d = torch.diagonal((x * 255).to(torch.uint8), 0, 2, 3)
+                return torch.nn.functional.pad(d, (1, 1)).float() / 255.0
+
+        class PadInt8(torch.nn.Module):
+            def forward(self, x):
+                return torch.nn.functional.pad(x.to(torch.int8), (1, 1))
+
+        for model, x in (
+            (DiagonalPad(), torch.rand(1, 3, 8, 8)),
+            (PadInt8(), torch.randint(-50, 50, (1, 3, 8)).float()),
+        ):
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                self.assertIn("aten.constant_pad_nd.default", op_names)
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(buffer)
+                    out = program.load_method("forward").execute((x,))[0]
+                    self.assertTrue(
+                        torch.allclose(out.float(), model(x).float(), atol=1e-2)
+                    )
+
     def test_tied_embedding_is_quantized_with_its_linear(self):
         """A weight that is both an embedding table and an output projection lowers.
 

From 8816121defe04c0ffa749279cd210123de091b91 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:47 -0700
Subject: [PATCH 17/44] Stop the Core ML int8 output search from looping on
 several outputs

An op with several outputs, such as split or max(dim), hands each
output out through a getitem. When such an int8 output was read outside
Core ML, the int8 output rule left out the getitem. The fx partitioner
always puts a getitem back next to its producer, so the next round found
the same getitem again, and to_edge_transform_and_lower never returned.

The rule now leaves out the op that makes the output, and the search
stops once a round leaves out nothing new, so a node the partitioner
puts back can never make it spin.
---
 .../coreml/partition/coreml_partitioner.py    |  9 +++-
 .../coreml/test/test_coreml_partitioner.py    | 44 +++++++++++--------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 0af4705bb86..3eb7591eea6 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -4,6 +4,7 @@
 
 import inspect
 import logging
+import operator
 from typing import Any, Callable, Collection, List, Optional, Set, Tuple
 
 import coremltools as ct
@@ -257,10 +258,11 @@ def _find_makers_of_int8_outputs(
     graph output or, with readers, by one of those nodes.
 
     Such a tensor becomes a Core ML output, and the Core ML runtime cannot
-    return int8.
+    return int8. An op with several outputs hands each one out through a
+    getitem, which always stays with the op, so the op itself is returned.
     """
     return {
-        node
+        node.args[0] if node.target is operator.getitem else node
         for partition in partitions
         for node in partition.nodes
         if _makes_int8(node)
@@ -384,6 +386,9 @@ def _propose_partitions(
             new |= _find_nodes_of_partitions_without_tensor_inputs(
                 partitions, kept_inside, state, allow_partitions_without_inputs
             )
+        # The partitioner can put a node that was left out back, such as a
+        # getitem next to its producer, so only a new node makes progress.
+        new -= left_out
         if not new:
             return partitions, left_out
         left_out |= new
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index f7630a8535f..ed3ecdfaafb 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -596,32 +596,38 @@ def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
         """
         The Core ML runtime cannot return an int8 tensor, so an op that makes
         one that is read outside Core ML, here the model output, runs on the
-        CPU.
+        CPU.  So does an op that hands one out through a getitem, which always
+        stays with the op.
         """
 
         class Model(torch.nn.Module):
             def forward(self, x):
                 return (x * 2).to(torch.int8)
 
-        model = Model().eval()
+        class Split(torch.nn.Module):
+            def forward(self, x):
+                return (x * 2).to(torch.int8).split(2, dim=1)[0]
+
         # Small integers, which Core ML's float16 math keeps exact.
-        x = torch.randint(-50, 50, (1, 3, 4, 4)).float()
-        ep = torch.export.export(model, (x,), strict=True)
-        edge = executorch.exir.to_edge_transform_and_lower(
-            ep, partitioner=[CoreMLPartitioner()]
-        )
-        op_names = [
-            n.target.__name__
-            for n in edge.exported_program().graph.nodes
-            if n.op == "call_function"
-        ]
-        self.assertIn("executorch_call_delegate", op_names)
-        self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names)
-        buffer = edge.to_executorch().buffer
-        if _TEST_RUNTIME:
-            program = Runtime.get().load_program(buffer)
-            out = program.load_method("forward").execute((x,))[0]
-            self.assertTrue(torch.equal(out, model(x)))
+        x = torch.randint(-50, 50, (1, 4, 4, 4)).float()
+        for model in (Model().eval(), Split().eval()):
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model, (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                self.assertIn("executorch_call_delegate", op_names)
+                self.assertIn("dim_order_ops._to_dim_order_copy.default", op_names)
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(buffer)
+                    out = program.load_method("forward").execute((x,))[0]
+                    self.assertTrue(torch.equal(out, model(x)))
 
     def test_uint8_camera_frame_is_read_outside_the_delegate(self):
         """

From e013b356d407fa1f8bb153f8b51dc014e9ebbe30 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:48 -0700
Subject: [PATCH 18/44] Check where Core ML state ends up after partitioning

The partitioner refused a CPU op that reads state Core ML takes over,
and never moved a partition that uses such state. That missed state
whose readers end up in two Core ML partitions, or in Core ML and on the
CPU, when a CPU op that reads a uint8 input sits between them. Each
place then works on its own copy, so one of them reads stale values. It
also refused state that only CPU ops use, which works.

After tagging, the partitioner now checks every user of each state
buffer, and the op that writes its new value. If they are in more than
one place, lowering fails and points to take_over_mutable_buffer=False.
State used only on the CPU stays there, so a partition that only updates
state can move to the CPU below iOS 18 like any partition with no
inputs.
---
 .../coreml/partition/coreml_partitioner.py    |  56 ++++---
 .../coreml/test/test_coreml_partitioner.py    | 146 +++++++++++-------
 2 files changed, 119 insertions(+), 83 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 3eb7591eea6..aefc58e4d68 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -28,6 +28,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.tracer import _default_decomposition_table
 from torch.export.exported_program import ExportedProgram
+from torch.export.graph_signature import OutputKind
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import (
     chain,
@@ -313,7 +314,6 @@ def _makes_tensor(node: torch.fx.Node) -> bool:
 def _find_nodes_of_partitions_without_tensor_inputs(
     partitions: List[Partition],
     kept_inside: Collection[str],
-    state: Collection[str],
     allow_partitions_without_inputs: bool,
 ) -> Set[torch.fx.Node]:
     """The nodes of the partitions that read no tensor from outside, except the
@@ -321,21 +321,15 @@ def _find_nodes_of_partitions_without_tensor_inputs(
 
     Core ML cannot take a symbolic size as an input, and below iOS 18 it
     rejects a model with no inputs, so such a partition cannot be delegated.
-    A partition that uses state stays: moving it to the CPU would leave the
-    state's other readers in Core ML a stale copy, and Core ML state needs
-    iOS 18 anyway.
     """
     nodes = set()
     for partition in partitions:
-        outside = [
+        inputs = [
             arg
             for node in partition.nodes
             for arg in node.all_input_nodes
-            if arg not in partition.nodes
+            if arg not in partition.nodes and arg.name not in kept_inside
         ]
-        if any(arg.name in state for arg in outside):
-            continue
-        inputs = [arg for arg in outside if arg.name not in kept_inside]
         reads_tensor = any(_makes_tensor(arg) for arg in inputs)
         if not reads_tensor and (inputs or not allow_partitions_without_inputs):
             nodes.update(partition.nodes)
@@ -346,7 +340,6 @@ def _propose_partitions(
     graph_module: torch.fx.GraphModule,
     op_support: OperatorSupportBase,
     kept_inside: Collection[str],
-    state: Collection[str],
     allow_partitions_without_inputs: bool,
     before_decomposition: bool = False,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
@@ -384,7 +377,7 @@ def _propose_partitions(
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
             new |= _find_nodes_of_partitions_without_tensor_inputs(
-                partitions, kept_inside, state, allow_partitions_without_inputs
+                partitions, kept_inside, allow_partitions_without_inputs
             )
         # The partitioner can put a node that was left out back, such as a
         # getitem next to its producer, so only a new node makes progress.
@@ -670,7 +663,6 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         partition_tags = {}
 
         kept_inside = self._placeholders_kept_inside(exported_program)
-        state = self._state_kept_inside(exported_program)
         partition_list, left_out = _propose_partitions(
             exported_program.graph_module,
             _OperatorsSupportedForCoreMLBackend(
@@ -679,18 +671,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 log=True,
             ),
             kept_inside,
-            state,
             self._allows_partitions_without_inputs(),
         )
-        for node in left_out:
-            read = [arg.name for arg in node.all_input_nodes if arg.name in state]
-            if read:
-                raise NotImplementedError(
-                    f"{node} has to run outside of CoreML, but it reads "
-                    f"{', '.join(read)}, which CoreML takes over as state, so it "
-                    "would read a stale copy.  Set take_over_mutable_buffer=False "
-                    "in the CoreMLPartitioner to keep the state outside of CoreML."
-                )
         if left_out:
             inputs = sorted(
                 {
@@ -732,6 +714,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 tag = f"tag{partition.id}"
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
+        self._check_state_is_used_in_one_place(exported_program)
 
         if self.take_over_constant_data:
             tag_constant_data(exported_program)
@@ -754,6 +737,34 @@ def log_once(self, msg: str) -> None:
             logger.info(msg)
             self._logged_msgs.add(msg)
 
+    def _check_state_is_used_in_one_place(self, ep: ExportedProgram) -> None:
+        """Raises if ops in more than one place, CoreML partitions or the CPU,
+        read or update a buffer CoreML takes over as state, since each place
+        would then work on its own copy of it.
+        """
+        signature = ep.graph_signature
+        nodes = {node.name: node for node in ep.graph.nodes}
+        updates = {
+            spec.target: nodes[spec.arg.name]
+            for spec in signature.output_specs
+            if spec.kind == OutputKind.BUFFER_MUTATION
+        }
+        for name in self._state_kept_inside(ep):
+            users = [user for user in nodes[name].users if user.op != "output"]
+            users.append(updates[signature.inputs_to_buffers[name]])
+            if len({user.meta.get("delegation_tag") for user in users}) > 1:
+                hint = (
+                    ""
+                    if self.lower_full_graph
+                    else "  Set take_over_mutable_buffer=False in the "
+                    "CoreMLPartitioner to keep the state outside of CoreML."
+                )
+                raise NotImplementedError(
+                    f"CoreML takes over {name} as state, but it is used in more "
+                    "than one place: in two CoreML partitions, or both in and "
+                    f"outside of CoreML.  Each place would use its own copy.{hint}"
+                )
+
     def _allows_partitions_without_inputs(self) -> bool:
         """Whether Core ML accepts a partition with no inputs, which needs iOS 18."""
         target = CoreMLBackend.min_deployment_target_from_compile_specs(
@@ -833,7 +844,6 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 ep.graph_module,
                 create_op_support(is_supported),
                 self._placeholders_kept_inside(ep),
-                self._state_kept_inside(ep),
                 self._allows_partitions_without_inputs(),
                 before_decomposition=True,
             )
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index ed3ecdfaafb..5402e286a7a 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -799,61 +799,69 @@ def forward(self, image):
                     out = program.load_method("forward").execute((y,))[0]
                     self.assertTrue(torch.allclose(out, model(y), atol=1e-2))
 
-    def test_uint8_input_keeps_state_in_the_delegate(self):
+    def test_uint8_input_with_state(self):
         """
-        A partition that only updates state, which has no inputs, stays in
-        Core ML with that state instead of moving to the CPU.  Below iOS 18
-        Core ML then refuses the state at lowering, as it does for any model
-        with state.
+        A CPU op that reads a uint8 input can separate a read of state from its
+        update, so two places would each use their own copy of the state.
+        Lowering refuses that, at any deployment target, unless the state stays
+        outside of Core ML.  State that only a partition with no inputs uses
+        moves to the CPU with it below iOS 18.
         """
 
-        class Model(torch.nn.Module):
+        class Split(torch.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.register_buffer("frames", torch.zeros(1))
+                self.register_buffer("state", torch.zeros(1, 3, 2, 2))
 
             def forward(self, image):
-                out = image.float() / 255.0 * self.frames
-                frames = self.frames + 1
-                self.frames.copy_(frames)
-                # The skipped sqrt runs on the CPU, so the update gets its own
-                # partition.
-                return out * torch.sqrt(frames + 1)
-
-        skip = ["aten.sqrt.default"]
-        x = torch.randint(0, 256, (1, 3, 2, 2), dtype=torch.uint8)
-        ep = torch.export.export(Model().eval(), (x,), strict=True)
-        with self.assertRaisesRegex(ValueError, "State model is supported only"):
-            executorch.exir.to_edge_transform_and_lower(
-                ep,
-                partitioner=[CoreMLPartitioner(skip_ops_for_coreml_delegation=skip)],
-            )
-        compile_specs = CoreMLBackend.generate_compile_specs(
+                y = (image + (self.state * 2).to(torch.uint8)).float()
+                self.state.add_(y)
+                return y + 1
+
+        class Counter(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("count", torch.zeros(1))
+
+            def forward(self, image):
+                self.count.add_(1)
+                return (image + self.count.to(torch.uint8)).float()
+
+        x = torch.randint(0, 3, (1, 3, 2, 2), dtype=torch.uint8)
+        ep = torch.export.export(Split().eval(), (x,), strict=True)
+        ios18 = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
         )
-        edge = executorch.exir.to_edge_transform_and_lower(
-            ep,
-            partitioner=[
-                CoreMLPartitioner(
-                    compile_specs=compile_specs, skip_ops_for_coreml_delegation=skip
-                )
-            ],
-        )
-        op_names = [
-            n.target.__name__
-            for n in edge.exported_program().graph.nodes
-            if n.op == "call_function"
+        for compile_specs in (None, ios18):
+            with self.subTest(ios18=compile_specs is not None):
+                with self.assertRaisesRegex(NotImplementedError, "b_state as state"):
+                    executorch.exir.to_edge_transform_and_lower(
+                        ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+                    )
+        lowered = [
+            (Split, CoreMLPartitioner(take_over_mutable_buffer=False)),
+            (Counter, CoreMLPartitioner()),
         ]
-        self.assertIn("executorch_call_delegate", op_names)
-        self.assertIn("aten.sqrt.default", op_names)
-        self.assertNotIn("aten.add.Tensor", op_names)
-        edge.to_executorch()
+        for model_class, partitioner in lowered:
+            with self.subTest(model=model_class.__name__):
+                ep = torch.export.export(model_class().eval(), (x,), strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[partitioner]
+                )
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    method = Runtime.get().load_program(buffer).load_method("forward")
+                    eager = model_class().eval()
+                    for _ in range(2):
+                        out = method.execute((x,))[0]
+                        self.assertTrue(torch.allclose(out, eager(x), atol=1e-2))
 
     def test_uint8_input_size_that_slices_state(self):
         """
         A slice of state to a size read from a uint8 input has to run on the
-        CPU, where it would read a stale copy of state that Core ML takes over,
-        so lowering refuses it unless the state stays outside of Core ML.
+        CPU.  When Core ML uses that state too, the slice would read a stale
+        copy, so lowering refuses it unless the state stays outside of Core ML.
+        When only CPU ops use it, the state stays on the CPU.
         """
 
         class Model(torch.nn.Module):
@@ -866,39 +874,57 @@ def forward(self, image, y):
                 self.hist.add_(1)
                 return out
 
-        model = Model().eval()
+        class SliceUpdate(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("hist", torch.zeros(64))
+
+            def forward(self, image, y):
+                self.hist[: image.shape[2]] += 1
+                return y * 2 + self.hist
+
         inputs = (
             torch.randint(0, 256, (1, 3, 8, 8), dtype=torch.uint8),
             torch.randn(64),
         )
-        ep = torch.export.export(
-            model,
-            inputs,
-            dynamic_shapes=({2: torch.export.Dim("height", min=2, max=64)}, None),
-            strict=True,
-        )
+        dynamic_shapes = ({2: torch.export.Dim("height", min=2, max=64)}, None)
         compile_specs = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
         )
+        ep = torch.export.export(
+            Model().eval(), inputs, dynamic_shapes=dynamic_shapes, strict=True
+        )
         with self.assertRaisesRegex(NotImplementedError, "take_over_mutable_buffer"):
             executorch.exir.to_edge_transform_and_lower(
                 ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
             )
-        edge = executorch.exir.to_edge_transform_and_lower(
-            ep,
-            partitioner=[
+        lowered = [
+            (
+                Model,
                 CoreMLPartitioner(
                     compile_specs=compile_specs, take_over_mutable_buffer=False
+                ),
+            ),
+            (SliceUpdate, CoreMLPartitioner(compile_specs=compile_specs)),
+        ]
+        for model_class, partitioner in lowered:
+            with self.subTest(model=model_class.__name__):
+                ep = torch.export.export(
+                    model_class().eval(),
+                    inputs,
+                    dynamic_shapes=dynamic_shapes,
+                    strict=True,
                 )
-            ],
-        )
-        if _TEST_RUNTIME:
-            method = Runtime.get().load_program(edge.to_executorch().buffer)
-            method = method.load_method("forward")
-            eager = Model().eval()
-            for _ in range(2):
-                out = method.execute(inputs)[0]
-                self.assertTrue(torch.allclose(out, eager(*inputs), atol=1e-2))
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[partitioner]
+                )
+                buffer = edge.to_executorch().buffer
+                if _TEST_RUNTIME:
+                    method = Runtime.get().load_program(buffer).load_method("forward")
+                    eager = model_class().eval()
+                    for _ in range(2):
+                        out = method.execute(inputs)[0]
+                        self.assertTrue(torch.allclose(out, eager(*inputs), atol=1e-2))
 
     def test_uint8_buffer_not_taken_over_is_read_outside_the_delegate(self):
         """

From d9bc00dcf5b1996846c6239278c3c31f9196b85e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:48 -0700
Subject: [PATCH 19/44] Predict Core ML CPU ops on the program as lowering
 decomposes it

ops_to_not_decompose ran its own search before decomposition to guess
which ops partition() would run on the CPU, and that guess could
disagree with partition(). A pad after a skipped flip, or after an add
that reads a uint8 tensor from another Core ML partition, was kept whole
and then ran on the CPU, where the portable kernels have no pad.

It now decomposes the program the way lowering will, runs the same
search partition() runs, and stops keeping every op with a call that
search leaves out or that Core ML rejects, until nothing changes. Skip
names are matched in their edge form there, as partition() sees them.
The work is skipped when nothing Core ML rejects gets decomposed and no
tensor outside the taken-over weights is uint8 or int8, so other graphs
are handled as before. The log now says an op is not kept whole instead
of claiming a reason and a decomposition that may not apply.
---
 .../coreml/partition/coreml_partitioner.py    | 233 ++++++++----------
 .../coreml/test/test_coreml_partitioner.py    |  42 +++-
 2 files changed, 135 insertions(+), 140 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index aefc58e4d68..ce37ea1cfd0 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -26,7 +26,12 @@
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes.dim_order_ops_registry import DimOrderOpsMap
+from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge
 from executorch.exir.tracer import _default_decomposition_table
+from torch._export.passes.replace_view_ops_with_view_copy_ops_pass import (
+    get_view_copy_of_view_op,
+)
 from torch.export.exported_program import ExportedProgram
 from torch.export.graph_signature import OutputKind
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
@@ -197,6 +202,18 @@ def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
     return isinstance(val, torch.Tensor) and val.dtype == torch.bool
 
 
+def _makes_tensor(
+    node: torch.fx.Node, dtypes: Optional[Collection[torch.dtype]] = None
+) -> bool:
+    """Whether node makes a tensor, of one of dtypes if given."""
+    val = node.meta.get("val")
+    vals = val if isinstance(val, (list, tuple)) else [val]
+    return any(
+        isinstance(v, torch.Tensor) and (dtypes is None or v.dtype in dtypes)
+        for v in vals
+    )
+
+
 def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
     """Whether node makes a uint8 or int8 tensor, which Core ML cannot take as a
     program input.
@@ -207,26 +224,19 @@ def _cannot_be_coreml_input(node: torch.fx.Node) -> bool:
     handle them, while other missing dtypes, such as uint16, would only fail
     later, at run time.
     """
-    val = node.meta.get("val")
-    vals = val if isinstance(val, (list, tuple)) else [val]
-    return any(
-        isinstance(v, torch.Tensor) and v.dtype in (torch.uint8, torch.int8)
-        for v in vals
-    )
+    return _makes_tensor(node, (torch.uint8, torch.int8))
 
 
 def _find_readers_of_unsupported_inputs(
     partitions: List[Partition],
     kept_inside: Collection[str],
     left_out: Collection[torch.fx.Node],
-    sources: Optional[Collection[torch.fx.Node]] = None,
 ) -> Set[torch.fx.Node]:
     """The nodes that read a tensor Core ML cannot take as an input from outside
     their partition, or a symbolic value computed by an op left out of Core ML.
 
     A tensor made inside the partition, or a weight, constant or state the
-    partitioner moves into Core ML, is not an input. With sources, only tensors
-    made by those nodes count.
+    partitioner moves into Core ML, is not an input.
     """
     return {
         node
@@ -235,7 +245,6 @@ def _find_readers_of_unsupported_inputs(
         if any(
             arg not in partition.nodes
             and arg.name not in kept_inside
-            and (sources is None or arg in sources)
             and (
                 _cannot_be_coreml_input(arg)
                 or (
@@ -251,12 +260,8 @@ def _find_readers_of_unsupported_inputs(
     }
 
 
-def _find_makers_of_int8_outputs(
-    partitions: List[Partition],
-    readers: Optional[Collection[torch.fx.Node]] = None,
-) -> Set[torch.fx.Node]:
-    """The nodes that make an int8 tensor read outside their partition, by the
-    graph output or, with readers, by one of those nodes.
+def _find_makers_of_int8_outputs(partitions: List[Partition]) -> Set[torch.fx.Node]:
+    """The nodes that make an int8 tensor read outside their partition.
 
     Such a tensor becomes a Core ML output, and the Core ML runtime cannot
     return int8. An op with several outputs hands each one out through a
@@ -266,51 +271,11 @@ def _find_makers_of_int8_outputs(
         node.args[0] if node.target is operator.getitem else node
         for partition in partitions
         for node in partition.nodes
-        if _makes_int8(node)
-        and any(
-            u not in partition.nodes
-            and (readers is None or u.op == "output" or u in readers)
-            for u in node.users
-        )
+        if _makes_tensor(node, (torch.int8,))
+        and any(user not in partition.nodes for user in node.users)
     }
 
 
-def _makes_int8(node: torch.fx.Node) -> bool:
-    val = node.meta.get("val")
-    vals = val if isinstance(val, (list, tuple)) else [val]
-    return any(isinstance(v, torch.Tensor) and v.dtype == torch.int8 for v in vals)
-
-
-def _lasting_cpu_sources(
-    graph: torch.fx.Graph,
-    partitions: List[Partition],
-    kept_inside: Collection[str],
-    decomposed: Collection[torch._ops.OpOverload],
-) -> Set[torch.fx.Node]:
-    """The nodes whose outputs still come from outside Core ML once the graph is
-    decomposed: uint8 and int8 graph inputs, ops outside Core ML that have no
-    default decomposition, and ops outside Core ML that read such a node.
-    """
-    in_partition = {node for partition in partitions for node in partition.nodes}
-    found = set()
-    for node in graph.nodes:
-        if node.op == "placeholder":
-            if node.name not in kept_inside and _cannot_be_coreml_input(node):
-                found.add(node)
-        elif node.op == "call_function" and node not in in_partition:
-            if node.target not in decomposed or any(
-                arg in found for arg in node.all_input_nodes
-            ):
-                found.add(node)
-    return found
-
-
-def _makes_tensor(node: torch.fx.Node) -> bool:
-    val = node.meta.get("val")
-    vals = val if isinstance(val, (list, tuple)) else [val]
-    return any(isinstance(v, torch.Tensor) for v in vals)
-
-
 def _find_nodes_of_partitions_without_tensor_inputs(
     partitions: List[Partition],
     kept_inside: Collection[str],
@@ -341,19 +306,13 @@ def _propose_partitions(
     op_support: OperatorSupportBase,
     kept_inside: Collection[str],
     allow_partitions_without_inputs: bool,
-    before_decomposition: bool = False,
 ) -> Tuple[List[Partition], Set[torch.fx.Node]]:
     """Propose Core ML partitions, leaving out every op that reads a tensor
-    Core ML cannot take as an input from outside its partition, and then every
-    partition this leaves without a tensor input.
-
-    On a graph before decomposition, a boundary can still go away, for example
-    when an op Core ML rejects decomposes into ops it takes, so then only the
-    boundaries that last are followed.
+    Core ML cannot take as an input from outside its partition or makes one it
+    cannot return, and then every partition this leaves without a tensor input.
 
     Returns the partitions and the ops left out.
     """
-    decomposed = _default_decomposition_table() if before_decomposition else {}
     left_out: Set[torch.fx.Node] = set()
     support = chain(create_op_support(lambda _, node: node not in left_out), op_support)
     # Each op moved out of Core ML can hand such a tensor to the next op, so
@@ -362,17 +321,9 @@ def _propose_partitions(
         partitions = CapabilityBasedPartitioner(
             graph_module, support, allows_single_node_partition=True
         ).propose_partitions()
-        if before_decomposition:
-            sources = _lasting_cpu_sources(
-                graph_module.graph, partitions, kept_inside, decomposed
-            )
-            new = _find_readers_of_unsupported_inputs(
-                partitions, kept_inside, left_out, sources
-            ) | _find_makers_of_int8_outputs(partitions, sources)
-        else:
-            new = _find_readers_of_unsupported_inputs(
-                partitions, kept_inside, left_out
-            ) | _find_makers_of_int8_outputs(partitions)
+        new = _find_readers_of_unsupported_inputs(
+            partitions, kept_inside, left_out
+        ) | _find_makers_of_int8_outputs(partitions)
         if new or left_out:
             # Moving those ops out can leave a partition that reads no tensor,
             # such as the index math of a gather.
@@ -387,33 +338,15 @@ def _propose_partitions(
         left_out |= new
 
 
-def _drop_ops_rejected_after_decomposition(
-    ep: ExportedProgram,
-    keep: List[torch._ops.OpOverload],
-    is_supported: Callable[[Any, torch.fx.Node], bool],
-) -> List[torch._ops.OpOverload]:
-    """Decomposing a rejected call can make new calls of a kept op, such as the
-    matmul of a bfloat16 linear, that are rejected too. A non-core op may have
-    no CPU kernel, so such ops go back to decomposition.
+def _edge_op_names(target: Any) -> Set[str]:
+    """The lowercase names target can have once lowered to the edge dialect,
+    which partition() matches skip names against: an ATen op gains its
+    namespace, a view op becomes its copy, and some ops become dim order ops.
     """
-    while any(torch.Tag.core not in op.tags for op in keep):
-        table = _default_decomposition_table()
-        for op in keep:
-            table.pop(op, None)
-        # Only a decomposition can make new calls.
-        if not any(node.target in table for node in ep.graph.nodes):
-            break
-        late = {
-            node.target
-            for node in ep.run_decompositions(table).graph.nodes
-            if node.target in keep
-            and torch.Tag.core not in node.target.tags
-            and not is_supported(None, node)
-        }
-        if not late:
-            break
-        keep = [op for op in keep if op not in late]
-    return keep
+    if not isinstance(target, torch._ops.OpOverload):
+        return {getattr(target, "__name__", "").lower()}
+    edge = aten_to_edge(get_view_copy_of_view_op(target._schema) or target)
+    return {edge.__name__.lower(), DimOrderOpsMap.get(edge, edge).__name__.lower()}
 
 
 def _is_view_op(op: torch._ops.OpOverload) -> bool:
@@ -833,21 +766,6 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
                 )
                 return False
 
-        # An op that partition() leaves out runs on the CPU, so it is not
-        # preserved: it is decomposed into core ATen ops, the op set the
-        # portable kernels target, unless it is a core op with no default
-        # decomposition. Ops are only left out because of a uint8 or int8
-        # tensor, so a graph without one needs no extra partitioning.
-        left_out = set()
-        if any(_cannot_be_coreml_input(node) for node in ep.graph.nodes):
-            _, left_out = _propose_partitions(
-                ep.graph_module,
-                create_op_support(is_supported),
-                self._placeholders_kept_inside(ep),
-                self._allows_partitions_without_inputs(),
-                before_decomposition=True,
-            )
-        decomposed_by_default = _default_decomposition_table() if left_out else {}
         rejected = set()
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
@@ -855,16 +773,6 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
             ):
                 if not is_supported(None, node):
                     rejected.add(node.target)
-                elif node in left_out:
-                    if torch.Tag.core not in node.target.tags or (
-                        node.target in decomposed_by_default
-                    ):
-                        self.log_once(
-                            f"{node.target} runs outside of CoreML here because "
-                            "of a tensor CoreML cannot take as an input, so every "
-                            "call of it is decomposed."
-                        )
-                    rejected.add(node.target)
                 elif node.target not in do_not_decompose_blocklist and not _is_view_op(
                     node.target
                 ):
@@ -872,6 +780,69 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
         # Without a per-node filter, a preserved op stays whole at every call, so
         # one rejected call would be left with no kernel.
         keep = [op for op in do_not_decompose if op not in rejected]
-        if rejected:
-            keep = _drop_ops_rejected_after_decomposition(ep, keep, is_supported)
-        return keep, None
+        return self._drop_ops_left_outside(ep, keep, rejected, is_supported), None
+
+    def _drop_ops_left_outside(
+        self,
+        ep: ExportedProgram,
+        keep: List[torch._ops.OpOverload],
+        rejected: Set[torch._ops.OpOverload],
+        is_supported: Callable[[Any, torch.fx.Node], bool],
+    ) -> List[torch._ops.OpOverload]:
+        """Drops every kept op that partition() leaves outside of CoreML at some
+        call, since the portable kernels may have no kernel for it.
+
+        partition() sees the program once lowering decomposes the ops that are
+        not kept. That can make new calls of a kept op that CoreML rejects, such
+        as the matmul of a bfloat16 linear, or change which ops read a uint8 or
+        int8 tensor from outside of CoreML. So this decomposes the program the
+        same way and runs the same search, again after each drop.
+        """
+        kept_inside = self._placeholders_kept_inside(ep)
+        skip = set(self.skip_ops_for_coreml_delegation)
+        not_kept = set(rejected)
+        while keep:
+            table = _default_decomposition_table()
+            for op in keep:
+                table.pop(op, None)
+            # Decomposing only views and blocklisted ops changes nothing here.
+            decomposed = any(
+                node.target in table and node.target in not_kept
+                for node in ep.graph.nodes
+            )
+            graph_module = (
+                ep.run_decompositions(table).graph_module
+                if decomposed
+                else ep.graph_module
+            )
+            left_out = set()
+            if any(
+                node.name not in kept_inside and _cannot_be_coreml_input(node)
+                for node in graph_module.graph.nodes
+            ):
+                _, left_out = _propose_partitions(
+                    graph_module,
+                    create_op_support(
+                        lambda submodules, node: is_supported(submodules, node)
+                        and not (skip and _edge_op_names(node.target) & skip)
+                    ),
+                    kept_inside,
+                    self._allows_partitions_without_inputs(),
+                )
+            kept = set(keep)
+            late = {
+                node.target
+                for node in graph_module.graph.nodes
+                if node.target in kept
+                and (node in left_out or (decomposed and not is_supported(None, node)))
+            }
+            if not late:
+                break
+            for op in late:
+                self.log_once(
+                    f"Not keeping {op} whole, because some call of it runs outside "
+                    "of CoreML."
+                )
+            not_kept |= late
+            keep = [op for op in keep if op not in late]
+        return keep
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 5402e286a7a..d7a6ccdc8fe 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -1006,9 +1006,10 @@ def forward(self, x):
     def test_ops_after_a_cpu_op_on_a_tensor_made_inside(self):
         """
         A uint8 tensor made in Core ML and passed through an op that stays on
-        the CPU, such as diagonal, sends the pad after it to the CPU too.  The
-        pad is decomposed there, since the portable kernels have no pad, and so
-        is a pad that makes an int8 model output.
+        the CPU, such as diagonal, a skipped flip or view, or an add that
+        reads it from another partition past acosh, sends the pad after it to
+        the CPU too.  The pad is decomposed there, since the portable kernels have no
+        pad, and so is a pad that makes an int8 model output.
         """
 
         class DiagonalPad(torch.nn.Module):
@@ -1016,18 +1017,40 @@ def forward(self, x):
                 d = torch.diagonal((x * 255).to(torch.uint8), 0, 2, 3)
                 return torch.nn.functional.pad(d, (1, 1)).float() / 255.0
 
-        class PadInt8(torch.nn.Module):
+        class FlipPad(torch.nn.Module):
             def forward(self, x):
-                return torch.nn.functional.pad(x.to(torch.int8), (1, 1))
+                d = (x * 255).to(torch.uint8).flip(-1)
+                return torch.nn.functional.pad(d, (1, 1)).float() / 255.0
 
-        for model, x in (
-            (DiagonalPad(), torch.rand(1, 3, 8, 8)),
-            (PadInt8(), torch.randint(-50, 50, (1, 3, 8)).float()),
+        class ViewPad(torch.nn.Module):
+            def forward(self, x):
+                d = (x * 255).to(torch.uint8).view(1, 3, 8)
+                return torch.nn.functional.pad(d, (1, 1)).float() / 255.0
+
+        class AcoshPad(torch.nn.Module):
+            def forward(self, x):
+                d = (x * 255).to(torch.uint8)
+                d = d + torch.acosh(d.float() + 2).to(torch.uint8)
+                return torch.nn.functional.pad(d, (1, 1)).float() / 255.0
+
+        class PadInt8(torch.nn.Module):
+            def forward(self, x):
+                return torch.nn.functional.pad((x * 2).to(torch.int8), (1, 1))
+
+        for model, x, skip in (
+            (DiagonalPad(), torch.rand(1, 3, 8, 8), None),
+            (FlipPad(), torch.rand(1, 3, 8), ["aten.flip.default"]),
+            (ViewPad(), torch.rand(1, 3, 2, 4), ["aten.view_copy.default"]),
+            (AcoshPad(), torch.rand(1, 3, 8), None),
+            (PadInt8(), torch.randint(-25, 25, (1, 3, 8)).float(), None),
         ):
             with self.subTest(model=type(model).__name__):
                 ep = torch.export.export(model.eval(), (x,), strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
-                    ep, partitioner=[CoreMLPartitioner()]
+                    ep,
+                    partitioner=[
+                        CoreMLPartitioner(skip_ops_for_coreml_delegation=skip)
+                    ],
                 )
                 op_names = [
                     n.target.__name__
@@ -1035,6 +1058,7 @@ def forward(self, x):
                     if n.op == "call_function"
                 ]
                 self.assertIn("aten.constant_pad_nd.default", op_names)
+                self.assertIn("executorch_call_delegate", op_names)
                 buffer = edge.to_executorch().buffer
                 if _TEST_RUNTIME:
                     program = Runtime.get().load_program(buffer)

From 8af1ced6408363c99837a11edf63b07753bd690c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:48 -0700
Subject: [PATCH 20/44] Name int8 outputs in the Core ML partitioner messages

When ops ran on the CPU only because they make an int8 tensor Core ML
cannot return, the message listed no inputs ("...partition: .") and,
with lower_full_graph=True, told the user to convert an input. The
messages now name such int8 outputs, list inputs only when there are
some, and ask for a dtype Core ML supports.
---
 .../coreml/partition/coreml_partitioner.py    | 31 ++++++++++++++-----
 .../coreml/test/test_coreml_partitioner.py    |  4 +++
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index ce37ea1cfd0..4b00e2f9d57 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -618,18 +618,33 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                     and _cannot_be_coreml_input(arg)
                 }
             )
+            outputs = sorted(
+                node.name
+                for node in left_out
+                if _makes_tensor(node, (torch.int8,))
+                and any(user not in left_out for user in node.users)
+            )
+            reasons = []
+            if inputs:
+                reasons.append(
+                    "reads a uint8 or int8 tensor, which CoreML cannot take as an "
+                    "input, from outside a CoreML partition: " + ", ".join(inputs)
+                )
+            if outputs:
+                reasons.append(
+                    "makes an int8 tensor, which CoreML cannot return, for use "
+                    "outside of CoreML: " + ", ".join(outputs)
+                )
             msg = (
-                "The model reads a uint8 or int8 tensor, which CoreML cannot take "
-                "as an input, from outside a CoreML partition: "
-                + ", ".join(inputs)
-                + "."
+                f"The model {', and '.join(reasons)}."
+                if reasons
+                else "The model uses a uint8 or int8 tensor CoreML cannot handle."
             )
             if self.lower_full_graph:
                 raise NotImplementedError(
-                    f"{msg}  Convert the tensor to a supported dtype before it "
-                    "enters the model, or set lower_full_graph=False in the "
-                    "CoreMLPartitioner to run the ops that read it outside of "
-                    "CoreML."
+                    f"{msg}  Use a dtype CoreML supports for these tensors, or set "
+                    "lower_full_graph=False in the CoreMLPartitioner to run the ops "
+                    "that use them outside of CoreML."
                 )
             self.log_once(
                 f"{msg}  Running these ops outside of CoreML: "
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index d7a6ccdc8fe..ea644410fb8 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -628,6 +628,10 @@ def forward(self, x):
                     program = Runtime.get().load_program(buffer)
                     out = program.load_method("forward").execute((x,))[0]
                     self.assertTrue(torch.equal(out, model(x)))
+                with self.assertRaisesRegex(NotImplementedError, "makes an int8"):
+                    executorch.exir.to_edge_transform_and_lower(
+                        ep, partitioner=[CoreMLPartitioner(lower_full_graph=True)]
+                    )
 
     def test_uint8_camera_frame_is_read_outside_the_delegate(self):
         """

From a7fdd80739f5f0af236671a6a6f354c551962dc3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:48 -0700
Subject: [PATCH 21/44] Note that a Core ML uint8 buffer needs
 take_over_constant_data=False

coremltools cannot cast a uint8 constant, so a uint8 buffer that the
model casts, such as a lookup table, fails to lower with the default
take_over_constant_data=True. Say so in the partitioner docs, and pin
the failure in the test that lowers it with False.
---
 backends/apple/coreml/test/test_coreml_partitioner.py | 4 ++++
 docs/source/backends/coreml/coreml-partitioner.md     | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index ea644410fb8..02760e16002 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -951,6 +951,10 @@ def forward(self, x):
         model = Model().eval()
         x = torch.randn(2, 8)
         ep = torch.export.export(model, (x,), strict=True)
+        with self.assertRaisesRegex(ValueError, "tensor\\[8,uint8\\]"):
+            executorch.exir.to_edge_transform_and_lower(
+                ep, partitioner=[CoreMLPartitioner()]
+            )
         edge = executorch.exir.to_edge_transform_and_lower(
             ep, partitioner=[CoreMLPartitioner(take_over_constant_data=False)]
         )
diff --git a/docs/source/backends/coreml/coreml-partitioner.md b/docs/source/backends/coreml/coreml-partitioner.md
index fa7a1de8277..7c249d539fa 100644
--- a/docs/source/backends/coreml/coreml-partitioner.md
+++ b/docs/source/backends/coreml/coreml-partitioner.md
@@ -6,7 +6,7 @@ The Core ML partitioner API allows for configuration of the model delegation to
  - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by Core ML.  By default, all ops that Core ML supports will be delegated.  See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation.
 - `compile_specs`: A list of `CompileSpec`s for the Core ML backend.  These control low-level details of Core ML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32).  These are discussed more below.
 - `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate).  If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood.  Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0.
-- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate.  If set to False, constant data is passed to the Core ML delegate as inputs.  By default, take_over_constant_data=True.
+- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate.  If set to False, constant data is passed to the Core ML delegate as inputs.  By default, take_over_constant_data=True.  A uint8 buffer that the model casts, such as a lookup table, needs take_over_constant_data=False: coremltools cannot cast a uint8 constant, while with False the cast runs on the CPU.
 - `lower_full_graph`: A boolean that indicates whether the entire graph must be lowered to Core ML.  If set to True and Core ML does not support an op, an error is raised during lowering.  If set to False and Core ML does not support an op, the op is executed on the CPU by ExecuTorch.  Although setting `lower_full_graph`=False can allow a model to lower where it would otherwise fail, it can introduce performance overhead in the model when there are unsupported ops.  You will see warnings about unsupported ops during lowering if there are any.  By default, `lower_full_graph`=False.
 
 

From ba6a891a86060765db7aeb9b6af7112002998fc9 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 11:35:49 -0700
Subject: [PATCH 22/44] Say which min and max inputs coremltools rejects

The docstring said coremltools computes min and max only over fp16,
fp32 and int32. Only a bool input fails: uint8, int8, int16, int32 and
int64 all lower and match.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 4b00e2f9d57..8924e5f95c1 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -191,10 +191,11 @@ def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool:
 
 
 def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
-    """coremltools 9.0 computes min and max only over fp16, fp32 and int32.
+    """coremltools 9.0 cannot compute a min or max over a bool tensor.
 
-    It raises ``ValueError`` for a bool input, such as the all-true check that
+    It raises ``ValueError`` for one, such as the all-true check that
     torch.distributions exports, instead of reporting the op as unsupported.
+    Integer tensors work.
     """
     if node.target not in _MIN_MAX_TARGETS:
         return False

From e8b2fb004c918b0c0200db3289ab458e04eb05a7 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 12:38:51 -0700
Subject: [PATCH 23/44] Say which ops the Core ML keep-whole check drops

The docstring said it drops every kept op that partition() leaves
outside of Core ML. It drops only ops moved out for a uint8 or int8
tensor, or rejected once the program is decomposed. An op skipped by
name stays kept, as on main, since it may have a portable kernel.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 8924e5f95c1..7bb405baf89 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -805,8 +805,10 @@ def _drop_ops_left_outside(
         rejected: Set[torch._ops.OpOverload],
         is_supported: Callable[[Any, torch.fx.Node], bool],
     ) -> List[torch._ops.OpOverload]:
-        """Drops every kept op that partition() leaves outside of CoreML at some
-        call, since the portable kernels may have no kernel for it.
+        """Drops every kept op that partition() moves out of CoreML at some call
+        for a uint8 or int8 tensor, or that CoreML rejects once the program is
+        decomposed, since the portable kernels may have no kernel for it. An op
+        skipped by name stays kept, as before.
 
         partition() sees the program once lowering decomposes the ops that are
         not kept. That can make new calls of a kept op that CoreML rejects, such

From 1e4972839c262446756efc1765d3b8a5dcc6a9f8 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 12:38:51 -0700
Subject: [PATCH 24/44] Test Core ML state updated on the CPU and a skipped
 cast

Two cases had no test that fails without them:
- State read in Core ML whose new value comes from a CPU op, here a
  cast of a uint8 input. Without the update in the state check,
  lowering fails later with an unclear error.
- A skipped cast, which lowering turns into a dim order op, so its
  skip name only matches that edge name.
---
 .../coreml/test/test_coreml_partitioner.py    | 34 +++++++++++++++----
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 02760e16002..fa3717da9fa 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -806,7 +806,8 @@ def forward(self, image):
     def test_uint8_input_with_state(self):
         """
         A CPU op that reads a uint8 input can separate a read of state from its
-        update, so two places would each use their own copy of the state.
+        update, or make the update itself, so two places would each use their
+        own copy of the state.
         Lowering refuses that, at any deployment target, unless the state stays
         outside of Core ML.  State that only a partition with no inputs uses
         moves to the CPU with it below iOS 18.
@@ -822,6 +823,16 @@ def forward(self, image):
                 self.state.add_(y)
                 return y + 1
 
+        class CopyIn(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1, 3, 2, 2))
+
+            def forward(self, image):
+                y = self.state * 2 + 1
+                self.state.copy_(image.float())
+                return y
+
         class Counter(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -832,12 +843,18 @@ def forward(self, image):
                 return (image + self.count.to(torch.uint8)).float()
 
         x = torch.randint(0, 3, (1, 3, 2, 2), dtype=torch.uint8)
-        ep = torch.export.export(Split().eval(), (x,), strict=True)
         ios18 = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
         )
-        for compile_specs in (None, ios18):
-            with self.subTest(ios18=compile_specs is not None):
+        for model_class, compile_specs in (
+            (Split, None),
+            (Split, ios18),
+            (CopyIn, ios18),
+        ):
+            with self.subTest(
+                model=model_class.__name__, ios18=compile_specs is not None
+            ):
+                ep = torch.export.export(model_class().eval(), (x,), strict=True)
                 with self.assertRaisesRegex(NotImplementedError, "b_state as state"):
                     executorch.exir.to_edge_transform_and_lower(
                         ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
@@ -1014,7 +1031,7 @@ def forward(self, x):
     def test_ops_after_a_cpu_op_on_a_tensor_made_inside(self):
         """
         A uint8 tensor made in Core ML and passed through an op that stays on
-        the CPU, such as diagonal, a skipped flip or view, or an add that
+        the CPU, such as diagonal, a skipped flip, view or cast, or an add that
         reads it from another partition past acosh, sends the pad after it to
         the CPU too.  The pad is decomposed there, since the portable kernels have no
         pad, and so is a pad that makes an int8 model output.
@@ -1049,10 +1066,15 @@ def forward(self, x):
             (DiagonalPad(), torch.rand(1, 3, 8, 8), None),
             (FlipPad(), torch.rand(1, 3, 8), ["aten.flip.default"]),
             (ViewPad(), torch.rand(1, 3, 2, 4), ["aten.view_copy.default"]),
+            (
+                FlipPad(),
+                torch.rand(1, 3, 8),
+                ["dim_order_ops._to_dim_order_copy.default"],
+            ),
             (AcoshPad(), torch.rand(1, 3, 8), None),
             (PadInt8(), torch.randint(-25, 25, (1, 3, 8)).float(), None),
         ):
-            with self.subTest(model=type(model).__name__):
+            with self.subTest(model=type(model).__name__, skip=skip):
                 ep = torch.export.export(model.eval(), (x,), strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
                     ep,

From 51dba926269cf1b017068463dba336f0191b706f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 12:46:44 -0700
Subject: [PATCH 25/44] Explain why the Core ML partition search drops nodes
 already left out

The comment said the partitioner can put a left-out getitem back next to
its producer. Since the int8 rule now leaves out the producer, that no
longer happens. The line only guarantees that the loop ends, so say so.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 7bb405baf89..135b160f496 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -331,8 +331,7 @@ def _propose_partitions(
             new |= _find_nodes_of_partitions_without_tensor_inputs(
                 partitions, kept_inside, allow_partitions_without_inputs
             )
-        # The partitioner can put a node that was left out back, such as a
-        # getitem next to its producer, so only a new node makes progress.
+        # Stop once a round leaves out nothing new, so the loop always ends.
         new -= left_out
         if not new:
             return partitions, left_out

From 3aad2f8540c334f5cc874d8e7d15e547dbd434ac Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Fri, 30 Oct 2026 16:58:35 -0700
Subject: [PATCH 26/44] Stop Core ML from making int and bool fills with the
 wrong dtype

coremltools 9.0 often ignores the dtype of an op that fills a tensor
with one value. torch.full((1, 16), 0) has no dtype, so torch makes it
int64, but coremltools makes it float. zeros and new_zeros come out
float even with an int dtype, ones_like does from iOS 16, and
full_like, fill and scalar_tensor take the type of their fill value.

An int or bool tensor made as float breaks in two ways. An op that does
not promote types, such as index_put through a bool mask, fails to
lower ("updates has dtype int32 whereas data has dtype fp32"). The
SmolVLM vision embeddings hit this. Other ops run at float16 precision
and silently lose integers above 2048.

The partitioner now refuses those fills when they make an int or bool
tensor. full, full_like, fill and scalar_tensor run on the portable
kernels. zeros, new_zeros, zeros_like, ones_like and new_full are no
longer kept whole, so lowering decomposes them into a full or full_like
with the dtype given, which Core ML handles. Float fills are unchanged.
With lower_full_graph=True, a model with such a fill now raises instead
of lowering with the wrong dtype.
---
 .../coreml/partition/coreml_partitioner.py    | 63 +++++++++++++++++++
 .../coreml/test/test_coreml_partitioner.py    | 55 ++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 135b160f496..b856ac6d7e5 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -203,6 +203,62 @@ def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
     return isinstance(val, torch.Tensor) and val.dtype == torch.bool
 
 
+_FULL_TARGETS = (torch.ops.aten.full.default, exir_ops.edge.aten.full.default)
+
+
+# Where the fill value is, for ops that coremltools types by it.
+_FILL_VALUE_INDEX = {
+    torch.ops.aten.full_like.default: 1,
+    torch.ops.aten.fill.Scalar: 1,
+    torch.ops.aten.scalar_tensor.default: 0,
+    exir_ops.edge.aten.full_like.default: 1,
+    exir_ops.edge.aten.fill.Scalar: 1,
+    exir_ops.edge.aten.scalar_tensor.default: 0,
+}
+
+
+_DECOMPOSED_FILL_TARGETS = (
+    torch.ops.aten.zeros.default,
+    torch.ops.aten.new_zeros.default,
+    torch.ops.aten.zeros_like.default,
+    torch.ops.aten.ones_like.default,
+    torch.ops.aten.new_full.default,
+    exir_ops.edge.aten.zeros.default,
+    exir_ops.edge.aten.new_zeros.default,
+    exir_ops.edge.aten.zeros_like.default,
+    exir_ops.edge.aten.ones_like.default,
+    exir_ops.edge.aten.new_full.default,
+)
+
+
+def _fills_int_or_bool_tensor_with_wrong_type(node: torch.fx.Node) -> bool:
+    """coremltools 9.0 often ignores the dtype of an op that fills a tensor.
+
+    It makes ``full`` without a dtype, ``zeros``, ``new_zeros`` and, from iOS
+    16, ``ones_like`` as float, ``zeros_like`` with the dtype of its input, and
+    ``full_like``, ``new_full``, ``fill`` and ``scalar_tensor`` with the type
+    of their fill value. An int or bool tensor made as float silently loses
+    values above 2048 at float16 precision, and fails to lower in an op that
+    does not promote types, such as ``index_put``. Ops in
+    ``_DECOMPOSED_FILL_TARGETS`` that make one are refused even when they would
+    work, so lowering decomposes them into ``full`` or ``full_like`` with the
+    dtype given.
+    """
+    val = node.meta.get("val")
+    if not isinstance(val, torch.Tensor) or val.dtype.is_floating_point:
+        return False
+    if node.target in _DECOMPOSED_FILL_TARGETS:
+        return True
+    if node.target in _FULL_TARGETS:
+        return node.kwargs.get("dtype") is None
+    if node.target in _FILL_VALUE_INDEX:
+        value = node.args[_FILL_VALUE_INDEX[node.target]]
+        return isinstance(value, (bool, int, float)) and type(value) is not (
+            bool if val.dtype == torch.bool else int
+        )
+    return False
+
+
 def _makes_tensor(
     node: torch.fx.Node, dtypes: Optional[Collection[torch.dtype]] = None
 ) -> bool:
@@ -451,6 +507,13 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _fills_int_or_bool_tensor_with_wrong_type(node):
+            self.log_once(
+                f"{node.target} making an int or bool tensor gets the wrong dtype "
+                "in CoreML.  Overriding op support."
+            )
+            return True
+
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
         # # in the placeholders due to partitioning, which CoreML does not support
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index fa3717da9fa..84c13e4c2f1 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -592,6 +592,61 @@ def forward(self, x):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
+    def test_int_and_bool_fills_get_their_dtype(self):
+        """
+        coremltools makes a filled int or bool tensor with the wrong dtype in
+        some cases, such as ``full`` without a dtype or ``zeros`` with one.
+        Index math with such a tensor then fails to lower, or loses integers
+        above 2048 to float16 precision.  The partitioner must leave those
+        fills to the portable backend, or let lowering decompose them into a
+        ``full`` with the dtype given.
+        """
+
+        class Positions(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.emb = torch.nn.Embedding(16, 4)
+
+            def forward(self, x, mask):
+                pos = torch.full((1, 16), fill_value=0)
+                pos[mask] = torch.arange(16)[None][mask]
+                return x + self.emb(pos)
+
+        class Counts(torch.nn.Module):
+            def forward(self, ids):
+                return (ids + torch.zeros((2, 3), dtype=torch.long)) * 1001
+
+        class Flags(torch.nn.Module):
+            def forward(self, mask):
+                return mask & torch.full_like(mask, 1)
+
+        mask = torch.arange(16)[None] % 3 != 0
+        cases = [
+            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default"),
+            (Counts(), (torch.arange(6).reshape(2, 3),), None),
+            (Flags(), (mask,), "aten.full_like.default"),
+        ]
+        for model, inputs, cpu_op in cases:
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), inputs, strict=True)
+                edge = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner()]
+                )
+                op_names = [
+                    n.target.__name__
+                    for n in edge.exported_program().graph.nodes
+                    if n.op == "call_function"
+                ]
+                self.assertIn("executorch_call_delegate", op_names)
+                if cpu_op:
+                    self.assertIn(cpu_op, op_names)
+                if _TEST_RUNTIME:
+                    program = Runtime.get().load_program(edge.to_executorch().buffer)
+                    out = program.load_method("forward").execute(inputs)[0]
+                    ref = model(*inputs)
+                    self.assertEqual(out.dtype, ref.dtype)
+                    self.assertTrue(torch.allclose(out, ref, atol=1e-2))
+
     def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
         """
         The Core ML runtime cannot return an int8 tensor, so an op that makes

From b9bdbf42679aacfc7c31575959a09d9124d87947 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 12:39:25 -0700
Subject: [PATCH 27/44] Give Core ML the eps torch uses for an RMSNorm without
 one

torch.nn.RMSNorm with its default eps=None exports rms_norm without the
eps argument. coremltools 9.0 requires that argument, so lowering fails
with "node aten_rms_norm_default (rms_norm) got 3 input(s), expected
[4]", and it turns a None eps into 1e-5. Torch computes RMSNorm in
float32 and uses the float32 machine epsilon, about 1.19e-7, also for
float16 and bfloat16 inputs.

The Core ML backend now gives coremltools that eps when the model leaves
it out, so float32 models with such an RMSNorm lower and match torch. A
float16 model still fails in coremltools' rms_norm, with or without eps.
The branch for an older coremltools, which has no rms_norm conversion,
is unchanged.
---
 backends/apple/coreml/compiler/torch_ops.py  | 12 +++++++
 backends/apple/coreml/test/test_torch_ops.py | 37 +++++++++++++++++---
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index 7f6abe980e1..dab49e477e5 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -58,6 +58,7 @@ def is_fbcode():
         _empty_dim_order,
         _to_dim_order_copy,
     )
+    from coremltools.converters.mil.frontend.torch.ops import rms_norm as _rms_norm
 
     # This is a temporary hack to register the alias "dim_order_ops._to_dim_order_copy",
     # which was missed by coremltools
@@ -71,6 +72,17 @@ def _to_dim_order_copy_TMP_EXECUTORCH_ALIAS_HACK(context, node):
     def _empty_dim_order_TMP_EXECUTORCH_ALIAS_HACK(context, node):
         _empty_dim_order(context, node)
 
+    # coremltools 9.0 requires the eps that torch.nn.RMSNorm(eps=None) leaves out,
+    # and https://github.com/apple/coremltools/pull/2743 would use 1e-5 for it.
+    # Torch computes in float32 with its machine epsilon, also for float16 inputs.
+    @register_torch_op(override=True)
+    def rms_norm(context, node):
+        inputs = list(node.inputs) + [None] * (4 - len(node.inputs))
+        if inputs[3] is None:
+            inputs[3] = float(np.finfo(np.float32).eps)
+        node.inputs = inputs
+        _rms_norm(context, node)
+
 else:
     # TODO: remove this case when fbcode updates to coremltools 9.0
     @register_torch_op(
diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py
index 10c3f01a585..2b954ff01d6 100644
--- a/backends/apple/coreml/test/test_torch_ops.py
+++ b/backends/apple/coreml/test/test_torch_ops.py
@@ -40,7 +40,9 @@ def _get_test_model(self):
         example_inputs = (torch.LongTensor([0]),)
         return model, example_inputs
 
-    def _compare_outputs(self, executorch_program, eager_program, example_inputs):
+    def _compare_outputs(
+        self, executorch_program, eager_program, example_inputs, atol=1e-2, rtol=1e-2
+    ):
         if not IS_VALID_TEST_RUNTIME:
             return
         runtime = Runtime.get()
@@ -48,9 +50,7 @@ def _compare_outputs(self, executorch_program, eager_program, example_inputs):
         method = program.load_method("forward")
         et_outputs = method.execute(example_inputs)[0]
         eager_outputs = eager_program(*example_inputs)
-        self.assertTrue(
-            torch.allclose(et_outputs, eager_outputs, atol=1e-02, rtol=1e-02)
-        )
+        self.assertTrue(torch.allclose(et_outputs, eager_outputs, atol=atol, rtol=rtol))
 
     def test_dequantize_affine_b4w_embedding(self):
         model, example_inputs = self._get_test_model()
@@ -317,6 +317,35 @@ def forward(self, x):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
 
+    def test_rms_norm_without_eps(self):
+        """
+        torch.nn.RMSNorm(eps=None) leaves eps out of rms_norm, which
+        coremltools requires.  Torch then uses the float32 machine epsilon,
+        which the output for a one-hot row of 65536 values depends on.
+        """
+        x = torch.zeros(1, 65536)
+        x[0, 0] = 1
+        compile_specs = CoreMLBackend.generate_compile_specs(
+            compute_precision=ct.precision.FLOAT32
+        )
+        for model in (
+            torch.nn.RMSNorm(65536),
+            torch.nn.RMSNorm(65536, elementwise_affine=False),
+        ):
+            with self.subTest(elementwise_affine=model.elementwise_affine):
+                ep = torch.export.export(model.eval(), (x,))
+                delegated_program = executorch.exir.to_edge_transform_and_lower(
+                    ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+                )
+                for node in delegated_program.exported_program().graph.nodes:
+                    if node.op == "call_function":
+                        self.assertIn(
+                            node.target.__name__,
+                            ["executorch_call_delegate", "getitem"],
+                        )
+                et_prog = delegated_program.to_executorch()
+                self._compare_outputs(et_prog, model, (x,), rtol=1e-3)
+
     def test_dequantize_affine_below_ios18_raises_with_hint(self):
         """
         Regression test for https://github.com/pytorch/executorch/issues/13122.

From 0f44136f6f8c29a6d22573791d115635662a87dd Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 12:40:19 -0700
Subject: [PATCH 28/44] Convert gelu with approximate="tanh" to the tanh gelu
 in Core ML

torch.export passes the approximate argument of gelu as a keyword. The
coremltools 9.0 converter reads it only from the positional inputs, so
gelu(approximate="tanh") silently becomes the exact gelu. The two
differ by up to 4.7e-4, at float32 and float16 precision alike, which
looks like poor float32 accuracy in Core ML.

The Core ML backend now registers a gelu conversion that also reads the
keyword. A tanh gelu written out with pow was not affected, since
coremltools fuses that formula into the right mode.
---
 backends/apple/coreml/compiler/torch_ops.py  | 16 ++++++++++++++
 backends/apple/coreml/test/test_torch_ops.py | 23 ++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index dab49e477e5..6da56142faa 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -11,6 +11,7 @@
 import numpy as np
 import torch as _torch
 from coremltools import _logger
+from coremltools.converters.mil import Builder as mb
 from coremltools.converters.mil.frontend import _utils
 from coremltools.converters.mil.frontend.torch.ops import (
     _get_inputs,
@@ -49,6 +50,21 @@ def split_copy(context, node):
     split(context, node)
 
 
+# https://github.com/apple/coremltools/pull/2863
+# coremltools reads approximate only from the positional inputs, but torch.export
+# passes it as a keyword, so gelu(approximate="tanh") became the exact gelu.
+@register_torch_op(override=True)
+def gelu(context, node):
+    inputs = _get_inputs(context, node, expected=[1, 2])
+    if len(inputs) == 2:
+        approximate = inputs[1].val
+    else:
+        # Read the keyword as is: _get_kwinputs would bind "tanh" to a node named tanh.
+        approximate = (node.kwinputs or {}).get("approximate", ["none"])[0]
+    mode = {"none": "EXACT", "tanh": "TANH_APPROXIMATION"}[approximate]
+    context.add(mb.gelu(x=inputs[0], mode=mode, name=node.name))
+
+
 def is_fbcode():
     return not hasattr(_torch.version, "git_version")
 
diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py
index 2b954ff01d6..54e489941aa 100644
--- a/backends/apple/coreml/test/test_torch_ops.py
+++ b/backends/apple/coreml/test/test_torch_ops.py
@@ -317,6 +317,29 @@ def forward(self, x):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs)
 
+    def test_gelu_tanh(self):
+        """
+        torch.export passes approximate="tanh" to gelu as a keyword, and
+        coremltools reads it only from the positional inputs.  If the keyword
+        is dropped, the exact gelu runs, which differs by up to 4.7e-4.
+        """
+        model = torch.nn.GELU(approximate="tanh").eval()
+        example_inputs = (torch.linspace(-8, 8, 4096),)
+        compile_specs = CoreMLBackend.generate_compile_specs(
+            compute_precision=ct.precision.FLOAT32
+        )
+        ep = torch.export.export(model, example_inputs)
+        delegated_program = executorch.exir.to_edge_transform_and_lower(
+            ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
+        )
+        for node in delegated_program.exported_program().graph.nodes:
+            if node.op == "call_function":
+                self.assertIn(
+                    node.target.__name__, ["executorch_call_delegate", "getitem"]
+                )
+        et_prog = delegated_program.to_executorch()
+        self._compare_outputs(et_prog, model, example_inputs, atol=1e-5, rtol=1e-5)
+
     def test_rms_norm_without_eps(self):
         """
         torch.nn.RMSNorm(eps=None) leaves eps out of rms_norm, which

From e6f7c68a9c5492d5823bcebfd9657f5a02c9d522 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:18:42 -0700
Subject: [PATCH 29/44] Give each Core ML output a value of its own

coremltools converts ops such as clone into nothing, so a Core ML
partition can return one of its own inputs, as in
"return x.clone(), x * 2", or in "y = x.clone()" right before an op that
runs on the CPU. It can also return one value twice, as in
"return y, y.clone()". When the partition has other outputs, Core ML
hands back such an output wrong: all zeros, or the bytes of another
output.

The Core ML backend now runs one more coremltools pass, which gives
each output that repeats an input or an earlier output an identity op of
its own, so every output comes from an op and has its own name.

The clone test now returns such outputs. It checks them at run time,
and checks the pass on the converted program, which needs no runtime.
---
 .../coreml/compiler/coreml_preprocess.py      | 31 +++++++++++++++++++
 backends/apple/coreml/test/test_torch_ops.py  | 28 ++++++++++++++---
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py
index 931d1552606..8afb334577d 100644
--- a/backends/apple/coreml/compiler/coreml_preprocess.py
+++ b/backends/apple/coreml/compiler/coreml_preprocess.py
@@ -17,6 +17,9 @@
 
 import coremltools as ct
 import coremltools.optimize as cto
+from coremltools.converters.mil import Builder as mb
+from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass
+from coremltools.converters.mil.mil.passes.pass_registry import register_pass
 from executorch.backends.apple.coreml import executorchcoreml
 from executorch.backends.apple.coreml.compiler.enumerated_shape_utils import (
     _get_ct_inputs,
@@ -67,6 +70,33 @@ def _gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]:
     return shared
 
 
+@register_pass(namespace="executorch", name="copy_reused_outputs")
+class _CopyReusedOutputs(AbstractGraphPass):
+    """
+    Gives each output of the program a value of its own.
+
+    coremltools converts ops such as clone into nothing, so a partition can
+    return one of its own inputs, or one value twice. Core ML then hands back
+    such an output wrong, such as all zeros or the bytes of another output.
+    """
+
+    def apply(self, prog):
+        for function in prog.functions.values():
+            seen = set(function.inputs.values())
+            outputs = []
+            for output in function.outputs:
+                if output in seen:
+                    with function:
+                        copy = mb.identity(x=output)
+                    # coremltools wants every op to carry the source program's
+                    # scopes, which set_outputs, unlike a var replacement, skips.
+                    function._copy_scope_info(output, copy)
+                    output = copy
+                seen.add(output)
+                outputs.append(output)
+            function.set_outputs(outputs)
+
+
 class COMPILE_SPEC_KEYS(Enum):
     COMPUTE_UNITS = "compute_units"
     MODEL_TYPE = "model_type"
@@ -565,6 +595,7 @@ def _convert_to_mlmodel(
         )
         compute_units = CoreMLBackend.compute_unit_from_compile_specs(compile_specs)
         pass_pipeline = CoreMLBackend.pass_pipeline_from_compile_specs(compile_specs)
+        pass_pipeline.append_pass("executorch::copy_reused_outputs")
         enumerated_shapes = CoreMLBackend.enumerated_shapes_from_compile_specs(
             compile_specs
         )
diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py
index 54e489941aa..aed29d86cce 100644
--- a/backends/apple/coreml/test/test_torch_ops.py
+++ b/backends/apple/coreml/test/test_torch_ops.py
@@ -9,6 +9,7 @@
 import executorch.exir
 
 import torch
+from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
@@ -48,9 +49,15 @@ def _compare_outputs(
         runtime = Runtime.get()
         program = runtime.load_program(executorch_program.buffer)
         method = program.load_method("forward")
-        et_outputs = method.execute(example_inputs)[0]
+        et_outputs = method.execute(example_inputs)
         eager_outputs = eager_program(*example_inputs)
-        self.assertTrue(torch.allclose(et_outputs, eager_outputs, atol=atol, rtol=rtol))
+        if isinstance(eager_outputs, torch.Tensor):
+            eager_outputs = (eager_outputs,)
+        self.assertEqual(len(et_outputs), len(eager_outputs))
+        for et_output, eager_output in zip(et_outputs, eager_outputs):
+            self.assertTrue(
+                torch.allclose(et_output, eager_output, atol=atol, rtol=rtol)
+            )
 
     def test_dequantize_affine_b4w_embedding(self):
         model, example_inputs = self._get_test_model()
@@ -296,14 +303,27 @@ def test_dequantize_codebook_embedding_per_grouped_row(self):
         self._compare_outputs(et_prog, model, example_inputs)
 
     def test__clone_dim_order_contiguous(self):
+        """
+        The clones convert to nothing, so the Core ML program returns its own
+        input, and one value twice, which Core ML used to return wrong.
+        """
+
         class Model(torch.nn.Module):
             def forward(self, x):
-                return torch.ops.dim_order_ops._clone_dim_order(
-                    x, dim_order=[0, 1, 2, 3]
+                y = x * 2
+                return (
+                    torch.ops.dim_order_ops._clone_dim_order(x, dim_order=[0, 1, 2, 3]),
+                    y,
+                    torch.ops.dim_order_ops._clone_dim_order(y, dim_order=[0, 1, 2, 3]),
                 )
 
         model, example_inputs = Model(), (torch.randn(1, 3, 8, 8),)
         ep = torch.export.export(model, example_inputs)
+        program = ct.convert(ep.run_decompositions({}), convert_to="milinternal")
+        PASS_REGISTRY["executorch::copy_reused_outputs"](program)
+        function = program.functions["main"]
+        self.assertEqual(len(set(function.outputs)), len(function.outputs))
+        self.assertFalse(set(function.outputs) & set(function.inputs.values()))
         delegated_program = executorch.exir.to_edge_transform_and_lower(
             ep,
             partitioner=[self._coreml_partitioner()],

From 6ba764d4e196335744c103b6006708cb750090b9 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 22:43:50 -0700
Subject: [PATCH 30/44] Duplicate a shared constant also where an op reads it
 from a list

When a constant has users in different partitions, or in a partition
and on the CPU, to_backend gives each user its own copy of the
placeholder. duplicate_constant_node only replaced the constant in the
top-level args and kwargs of a user, so an op that reads it from a list,
such as cat or the indices of index_put, kept the original. Lowering
then failed with "constant data node ... is tagged with ... but has user
... which has tag (None)". For example, a Core ML model whose index
buffer is read by a gather in Core ML and by an index_put on the CPU
failed this way.

The pass now replaces the constant wherever the user reads it.
---
 .../duplicate_constant_node_pass.py           | 21 ++------------
 exir/backend/test/test_passes.py              | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
index 07c3eb5e57a..1d76e541a04 100644
--- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
+++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
@@ -96,26 +96,9 @@ def duplicate_constant_node(
                     copied_constant_or_attribute_node.meta["val"] = (
                         constant_or_attribute_node.meta["val"]
                     )
-                    new_args = tuple(
-                        [
-                            (
-                                arg
-                                if arg != constant_or_attribute_node
-                                else copied_constant_or_attribute_node
-                            )
-                            for arg in users[ith].args
-                        ]
+                    users[ith].replace_input_with(
+                        constant_or_attribute_node, copied_constant_or_attribute_node
                     )
-                    new_kwargs = {
-                        key: (
-                            value
-                            if value != constant_or_attribute_node
-                            else copied_constant_or_attribute_node
-                        )
-                        for key, value in users[ith].kwargs.items()
-                    }
-                    users[ith].args = new_args
-                    users[ith].kwargs = new_kwargs
                     # The copy keeps the original target, so it reads the
                     # original tensor and needs no data of its own.
                     new_input_specs.append(
diff --git a/exir/backend/test/test_passes.py b/exir/backend/test/test_passes.py
index 815a5b5bbed..3ded45e7ce2 100644
--- a/exir/backend/test/test_passes.py
+++ b/exir/backend/test/test_passes.py
@@ -86,3 +86,31 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         self.assertEqual(len(copied_nodes), 1)
         self.assertNotEqual(user.kwargs["other"], buffer_node)
+
+    def test_duplicate_constant_node_in_a_list_argument(self) -> None:
+        """
+        An op such as cat or index_put reads its tensors from a list, and the
+        copy has to replace the constant there too.
+        """
+
+        class CatConst(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("const", torch.ones(2))
+
+            def forward(self, x):
+                return torch.cat([x, self.const]), torch.cat([self.const, x])
+
+        edge = exir.to_edge(export(CatConst(), (torch.ones(2),), strict=True))
+        program = edge.exported_program()
+        name = next(
+            node.name
+            for node in program.graph.nodes
+            if node.op == "placeholder" and is_buffer(program, node)
+        )
+
+        copied_nodes = duplicate_constant_node(program, name)
+
+        users = {node.name: len(node.users) for node in program.graph.nodes}
+        self.assertEqual(users[name], 1)
+        self.assertEqual([users[copy] for copy in copied_nodes], [1])

From d50afddcedddd2c300bfece9361fb5e519fb7f06 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:22:33 -0700
Subject: [PATCH 31/44] Keep index_put forms coremltools cannot convert out of
 Core ML

coremltools 9.0 converts index_put with a whole slice, such as
y[:, idx] = v, by writing v into the part of y the indices select, with
each indexed dim kept as size 1. That only works with accumulate=False,
int indices of shape () or (1,), a static size in every other dim, a v
that broadcasts to that part, and a non-bool y below iOS 18. With a 0-d
index or several indices, torch lays out that part differently, so v
must not vary before the last index, or the values land in the wrong
place. With two entries in idx, as in the X-VLA policy, lowering fails
with "Input squeeze_0 has rank 1 != other inputs rank 0". Without a
whole slice, it fails when several index tensors have different shapes,
as in y[rows[:, None], cols] = v, or include a bool mask.

The partitioner now refuses those forms, so they run on the portable
kernels. The forms coremltools handles stay in Core ML, such as a
single-entry index next to a slice for a cache update, a scalar written
through several single-entry indices, or index tensors of one shape
without a slice.

A negative index next to a slice still stays in Core ML, since the check
cannot see index values: it fails, or gives wrong values when it comes
from an input.

The fill test now checks how many calls of an op stay on the CPU, since
its index_put model has several forms that must all be refused.
---
 .../coreml/partition/coreml_partitioner.py    | 91 ++++++++++++++++++-
 .../coreml/test/test_coreml_partitioner.py    | 61 +++++++++++--
 2 files changed, 138 insertions(+), 14 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index b856ac6d7e5..98ba0b4079b 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -34,6 +34,7 @@
 )
 from torch.export.exported_program import ExportedProgram
 from torch.export.graph_signature import OutputKind
+from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import (
     chain,
@@ -259,6 +260,76 @@ def _fills_int_or_bool_tensor_with_wrong_type(node: torch.fx.Node) -> bool:
     return False
 
 
+_INDEX_PUT_TARGETS = (
+    torch.ops.aten.index_put.default,
+    exir_ops.edge.aten.index_put.default,
+)
+
+
+def _is_index_put_coremltools_cannot_convert(node: torch.fx.Node, ios18: bool) -> bool:
+    """coremltools 9.0 converts only some forms of ``index_put``.
+
+    With a whole slice, such as ``y[:, idx] = v``, it writes ``v`` into the
+    part of ``y`` the indices select, keeping each indexed dim as size 1. That
+    needs ``accumulate=False``, int indices of shape ``()`` or ``(1,)``, a
+    static size in every other dim, a ``v`` that broadcasts to that part, and
+    a non-bool ``y`` below iOS 18. With a 0-d index or several indices, torch
+    lays out that part differently, so ``v`` may only vary in the dims after
+    the last index. Without a whole slice, several indices need the same shape
+    and no bool mask. It raises or writes the wrong place for any other form,
+    such as ``idx`` with two entries, instead of reporting the op as
+    unsupported.
+
+    This check cannot see index values, so a negative index next to a slice
+    stays in Core ML: it fails, or gives wrong values when it comes from an
+    input.
+    """
+    if node.target not in _INDEX_PUT_TARGETS:
+        return False
+    indices = [i if i is None else i.meta["val"] for i in node.args[1]]
+    given = {dim: index for dim, index in enumerate(indices) if index is not None}
+    if len(given) == len(indices):
+        return len(indices) > 1 and any(
+            index.dtype == torch.bool
+            or not statically_known_true(sym_eq(index.shape, indices[0].shape))
+            for index in indices
+        )
+    accumulate = (
+        node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False)
+    )
+    y = node.args[0].meta["val"]
+    if (
+        accumulate
+        or (y.dtype == torch.bool and not ios18)
+        or any(
+            index.dtype == torch.bool
+            or index.dim() > 1
+            or not statically_known_true(index.numel() == 1)
+            for index in given.values()
+        )
+    ):
+        return True
+    selected = [1 if dim in given else size for dim, size in enumerate(y.shape)]
+    values = node.args[2].meta["val"].shape
+    reorders = len(given) > 1 or any(index.dim() == 0 for index in given.values())
+    after = len(selected) - 1 - max(given)
+    return (
+        any(isinstance(size, torch.SymInt) for size in selected)
+        or len(values) > len(selected)
+        or not all(
+            statically_known_true(v == 1) or statically_known_true(v == s)
+            for v, s in zip(reversed(values), reversed(selected))
+        )
+        or (
+            reorders
+            and not all(
+                statically_known_true(v == 1)
+                for v in values[: max(0, len(values) - after)]
+            )
+        )
+    )
+
+
 def _makes_tensor(
     node: torch.fx.Node, dtypes: Optional[Collection[torch.dtype]] = None
 ) -> bool:
@@ -419,12 +490,14 @@ def __init__(
         skip_ops_for_coreml_delegation: Optional[List[str]] = None,
         lower_full_graph: bool = False,
         log: bool = False,
+        ios18: bool = False,
     ) -> None:
         if skip_ops_for_coreml_delegation is None:
             skip_ops_for_coreml_delegation = []
         super().__init__()
         self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
         self.lower_full_graph = lower_full_graph
+        self.ios18 = ios18
         self._logged_msgs = set()
         self._log = log
 
@@ -514,6 +587,13 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
+        if _is_index_put_coremltools_cannot_convert(node, self.ios18):
+            self.log_once(
+                f"{node.target} with these indices is not supported by CoreML.  "
+                "Overriding op support."
+            )
+            return True
+
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
         # # in the placeholders due to partitioning, which CoreML does not support
@@ -665,9 +745,10 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 self.skip_ops_for_coreml_delegation,
                 self.lower_full_graph,
                 log=True,
+                ios18=self._targets_ios18(),
             ),
             kept_inside,
-            self._allows_partitions_without_inputs(),
+            self._targets_ios18(),
         )
         if left_out:
             inputs = sorted(
@@ -776,8 +857,9 @@ def _check_state_is_used_in_one_place(self, ep: ExportedProgram) -> None:
                     f"outside of CoreML.  Each place would use its own copy.{hint}"
                 )
 
-    def _allows_partitions_without_inputs(self) -> bool:
-        """Whether Core ML accepts a partition with no inputs, which needs iOS 18."""
+    def _targets_ios18(self) -> bool:
+        """Whether the deployment target is iOS 18 or later, which Core ML needs
+        for a partition with no inputs."""
         target = CoreMLBackend.min_deployment_target_from_compile_specs(
             self.delegation_spec.compile_specs
         )
@@ -818,6 +900,7 @@ def ops_to_not_decompose(
             self.skip_ops_for_coreml_delegation,
             self.lower_full_graph,
             log=False,
+            ios18=self._targets_ios18(),
         )
 
         # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace
@@ -907,7 +990,7 @@ def _drop_ops_left_outside(
                         and not (skip and _edge_op_names(node.target) & skip)
                     ),
                     kept_inside,
-                    self._allows_partitions_without_inputs(),
+                    self._targets_ios18(),
                 )
             kept = set(keep)
             late = {
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 84c13e4c2f1..ece732cc0d3 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -592,14 +592,16 @@ def forward(self, x):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
-    def test_int_and_bool_fills_get_their_dtype(self):
+    def test_ops_coremltools_gets_wrong_are_refused(self):
         """
         coremltools makes a filled int or bool tensor with the wrong dtype in
         some cases, such as ``full`` without a dtype or ``zeros`` with one.
         Index math with such a tensor then fails to lower, or loses integers
-        above 2048 to float16 precision.  The partitioner must leave those
-        fills to the portable backend, or let lowering decompose them into a
-        ``full`` with the dtype given.
+        above 2048 to float16 precision.  It also fails to lower some forms of
+        ``index_put``, such as ``y[:, idx] = v`` with two entries in ``idx``.
+        The partitioner must refuse those ops, so they run on the portable
+        backend, or lowering decomposes a fill into a ``full`` with the dtype
+        given.
         """
 
         class Positions(torch.nn.Module):
@@ -620,13 +622,53 @@ class Flags(torch.nn.Module):
             def forward(self, mask):
                 return mask & torch.full_like(mask, 1)
 
+        class Scatter(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("cols", torch.tensor([0, 2]))
+                self.register_buffer("col", torch.tensor([1]))
+                self.register_buffer("last", torch.tensor(3))
+                self.register_buffer("rows", torch.tensor([[1], [2]]))
+
+            def forward(self, x, v):
+                y = x.clone()
+                y[:, self.cols] = v
+                y[self.rows, self.cols] = v[:2]
+                torch.ops.aten.index_put_(y, [None, self.col], v[:, :1], True)
+                torch.ops.aten.index_put_(y, [None, self.last], v[:, 0])
+                mask = y > 0.5
+                mask[:, self.col] = True
+                return y * 2 + mask
+
+        class Scatter3d(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("last", torch.tensor(2))
+
+            def forward(self, x, v):
+                y = x.clone()
+                torch.ops.aten.index_put_(y, [None, None, self.last], v)
+                return y * 2
+
         mask = torch.arange(16)[None] % 3 != 0
         cases = [
-            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default"),
-            (Counts(), (torch.arange(6).reshape(2, 3),), None),
-            (Flags(), (mask,), "aten.full_like.default"),
+            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 1),
+            (Counts(), (torch.arange(6).reshape(2, 3),), "aten.full.default", 0),
+            (Flags(), (mask,), "aten.full_like.default", 1),
+            (
+                Scatter(),
+                (torch.rand(3, 4), torch.rand(3, 2)),
+                "aten.index_put.default",
+                5,
+            ),
+            (
+                Scatter3d(),
+                (torch.rand(2, 2, 3), torch.rand(2, 1)),
+                "aten.index_put.default",
+                1,
+            ),
         ]
-        for model, inputs, cpu_op in cases:
+        for model, inputs, cpu_op, cpu_calls in cases:
             with self.subTest(model=type(model).__name__):
                 ep = torch.export.export(model.eval(), inputs, strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
@@ -638,8 +680,7 @@ def forward(self, mask):
                     if n.op == "call_function"
                 ]
                 self.assertIn("executorch_call_delegate", op_names)
-                if cpu_op:
-                    self.assertIn(cpu_op, op_names)
+                self.assertEqual(op_names.count(cpu_op), cpu_calls)
                 if _TEST_RUNTIME:
                     program = Runtime.get().load_program(edge.to_executorch().buffer)
                     out = program.load_method("forward").execute(inputs)[0]

From 1380f93a0fa5fc38c665d9c8dc1f7577ba76288a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:23:28 -0700
Subject: [PATCH 32/44] Check tanh gelu, RMSNorm eps and int zeros without the
 Core ML runtime

The tanh gelu and RMSNorm tests and the int zeros case of the fill test
only failed without their fixes when the Core ML runtime ran the
program, so they passed on runners without it. The gelu test now also
checks that the converted program uses the tanh mode, the RMSNorm test
checks the eps in the converted program, and the fill test checks that
the int zeros is not kept whole, so lowering decomposes it.

The RMSNorm override now also says why it has to stay once coremltools
accepts a missing eps: coremltools would use 1e-5, not the torch eps.

The RMSNorm test is skipped in fbcode, whose coremltools has no rms_norm
converter to check.
---
 backends/apple/coreml/compiler/torch_ops.py      |  7 ++++---
 .../apple/coreml/test/test_coreml_partitioner.py |  5 +++++
 backends/apple/coreml/test/test_torch_ops.py     | 16 ++++++++++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index 6da56142faa..7d222149d54 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -88,9 +88,10 @@ def _to_dim_order_copy_TMP_EXECUTORCH_ALIAS_HACK(context, node):
     def _empty_dim_order_TMP_EXECUTORCH_ALIAS_HACK(context, node):
         _empty_dim_order(context, node)
 
-    # coremltools 9.0 requires the eps that torch.nn.RMSNorm(eps=None) leaves out,
-    # and https://github.com/apple/coremltools/pull/2743 would use 1e-5 for it.
-    # Torch computes in float32 with its machine epsilon, also for float16 inputs.
+    # coremltools 9.0 requires the eps that torch.nn.RMSNorm(eps=None) leaves out.
+    # Torch then uses the float32 machine epsilon, also for float16 inputs, while
+    # https://github.com/apple/coremltools/pull/2743 would use 1e-5, so keep this
+    # override after that change lands.
     @register_torch_op(override=True)
     def rms_norm(context, node):
         inputs = list(node.inputs) + [None] * (4 - len(node.inputs))
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index ece732cc0d3..4e4dca37742 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -687,6 +687,11 @@ def forward(self, x, v):
                     ref = model(*inputs)
                     self.assertEqual(out.dtype, ref.dtype)
                     self.assertTrue(torch.allclose(out, ref, atol=1e-2))
+        counts = torch.export.export(Counts(), cases[1][1], strict=True)
+        self.assertNotIn(
+            torch.ops.aten.zeros.default,
+            CoreMLPartitioner().ops_to_not_decompose(counts)[0],
+        )
 
     def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
         """
diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py
index aed29d86cce..7968a6d391e 100644
--- a/backends/apple/coreml/test/test_torch_ops.py
+++ b/backends/apple/coreml/test/test_torch_ops.py
@@ -14,6 +14,7 @@
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 from executorch.backends.apple.coreml.test.test_coreml_utils import (
+    is_fbcode,
     IS_VALID_TEST_RUNTIME,
 )
 from executorch.exir.backend.utils import format_delegated_graph
@@ -349,6 +350,11 @@ def test_gelu_tanh(self):
             compute_precision=ct.precision.FLOAT32
         )
         ep = torch.export.export(model, example_inputs)
+        program = ct.convert(ep.run_decompositions({}), convert_to="milinternal")
+        self.assertEqual(
+            [op.mode.val for op in program.find_ops(op_type="gelu")],
+            ["TANH_APPROXIMATION"],
+        )
         delegated_program = executorch.exir.to_edge_transform_and_lower(
             ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
         )
@@ -360,6 +366,7 @@ def test_gelu_tanh(self):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs, atol=1e-5, rtol=1e-5)
 
+    @unittest.skipIf(is_fbcode(), "The coremltools in fbcode has no rms_norm.")
     def test_rms_norm_without_eps(self):
         """
         torch.nn.RMSNorm(eps=None) leaves eps out of rms_norm, which
@@ -377,6 +384,15 @@ def test_rms_norm_without_eps(self):
         ):
             with self.subTest(elementwise_affine=model.elementwise_affine):
                 ep = torch.export.export(model.eval(), (x,))
+                program = ct.convert(
+                    ep.run_decompositions({}),
+                    convert_to="milinternal",
+                    compute_precision=ct.precision.FLOAT32,
+                )
+                self.assertEqual(
+                    [op.y.val for op in program.find_ops(op_type="add")],
+                    [torch.finfo(torch.float32).eps],
+                )
                 delegated_program = executorch.exir.to_edge_transform_and_lower(
                     ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
                 )

From 8e55fe4936d72ea4d01e187e215b615bff2eb3a6 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 20:24:29 -0700
Subject: [PATCH 33/44] Pick the Core ML ops to keep once per program

to_edge_transform_and_lower asks the partitioner twice in a row for the
ops to keep whole, with the same program. When the program has an op
Core ML rejects that has a decomposition, such as an int zeros or a
bfloat16 linear, each answer decomposes the whole program in a dry run,
which is slow for a large model.

The Core ML partitioner now keeps its last answer and returns it again
for the same program, as the MLX partitioner does. It holds the program
through a weak reference, so it never keeps a program alive.
---
 .../apple/coreml/partition/coreml_partitioner.py     | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 98ba0b4079b..1f40ea6f1d7 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -5,6 +5,7 @@
 import inspect
 import logging
 import operator
+import weakref
 from typing import Any, Callable, Collection, List, Optional, Set, Tuple
 
 import coremltools as ct
@@ -689,6 +690,7 @@ def __init__(
         self.lower_full_graph = lower_full_graph
         self.take_over_constant_data = take_over_constant_data
         self._logged_msgs = set()
+        self._ops_to_not_decompose_cache = None
 
         if self.lower_full_graph:
             assert (
@@ -895,6 +897,11 @@ def _state_kept_inside(self, ep: ExportedProgram) -> Set[str]:
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
+        # to_edge_transform_and_lower asks twice in a row for the same program,
+        # and the dry run in _drop_ops_left_outside can decompose all of it.
+        cached = self._ops_to_not_decompose_cache
+        if cached is not None and cached[0]() is ep:
+            return list(cached[1]), None
         do_not_decompose = []
         op_support = _OperatorsSupportedForCoreMLBackend(
             self.skip_ops_for_coreml_delegation,
@@ -941,7 +948,10 @@ def is_supported(submodules, node: torch.fx.Node) -> bool:
         # Without a per-node filter, a preserved op stays whole at every call, so
         # one rejected call would be left with no kernel.
         keep = [op for op in do_not_decompose if op not in rejected]
-        return self._drop_ops_left_outside(ep, keep, rejected, is_supported), None
+        keep = self._drop_ops_left_outside(ep, keep, rejected, is_supported)
+        # A weak reference, so the partitioner does not keep the program alive.
+        self._ops_to_not_decompose_cache = (weakref.ref(ep), keep)
+        return list(keep), None
 
     def _drop_ops_left_outside(
         self,

From 0ee707dbaa7e8d6b2e7702f10443f47568b39d17 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 22:25:36 -0700
Subject: [PATCH 34/44] Check only the Core ML state that Core ML takes over

The partitioner raised for every state buffer whose readers and update
were in different places. But tag_mutated_buffer only moves a buffer
into Core ML when an op in a partition reads it. When only CPU ops read
it, for example an op Core ML rejects or skips, the buffer stays a
program buffer and the delegate returns its new value, which worked
before the check. Such a model, like

    y = torch.acosh(self.state)
    self.state.copy_(x * 2)

failed to lower with "used in more than one place".

The check now skips a buffer that no partition reads.
---
 backends/apple/coreml/partition/coreml_partitioner.py |  3 +++
 backends/apple/coreml/test/test_coreml_partitioner.py | 10 +++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 1f40ea6f1d7..ce40a565faf 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -845,6 +845,9 @@ def _check_state_is_used_in_one_place(self, ep: ExportedProgram) -> None:
         }
         for name in self._state_kept_inside(ep):
             users = [user for user in nodes[name].users if user.op != "output"]
+            # tag_mutated_buffer leaves a buffer no partition reads outside of CoreML.
+            if all(user.meta.get("delegation_tag") is None for user in users):
+                continue
             users.append(updates[signature.inputs_to_buffers[name]])
             if len({user.meta.get("delegation_tag") for user in users}) > 1:
                 hint = (
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 4e4dca37742..9e95b6d71c5 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -911,7 +911,8 @@ def test_uint8_input_with_state(self):
         own copy of the state.
         Lowering refuses that, at any deployment target, unless the state stays
         outside of Core ML.  State that only a partition with no inputs uses
-        moves to the CPU with it below iOS 18.
+        moves to the CPU with it below iOS 18, and state that only CPU ops read
+        stays there, even when Core ML makes its new value.
         """
 
         class Split(torch.nn.Module):
@@ -943,6 +944,12 @@ def forward(self, image):
                 self.count.add_(1)
                 return (image + self.count.to(torch.uint8)).float()
 
+        class ReadOnCpu(CopyIn):
+            def forward(self, image):
+                y = torch.asinh(self.state)
+                self.state.copy_(image.float() + 1)
+                return y
+
         x = torch.randint(0, 3, (1, 3, 2, 2), dtype=torch.uint8)
         ios18 = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18
@@ -963,6 +970,7 @@ def forward(self, image):
         lowered = [
             (Split, CoreMLPartitioner(take_over_mutable_buffer=False)),
             (Counter, CoreMLPartitioner()),
+            (ReadOnCpu, CoreMLPartitioner()),
         ]
         for model_class, partitioner in lowered:
             with self.subTest(model=model_class.__name__):

From 8f3f6531cbbb5dfe2c8abe3f3af758a59a12b425 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 00:51:10 -0700
Subject: [PATCH 35/44] Convert Core ML fills with the dtype torch gives them

coremltools makes full without a dtype, zeros, new_zeros and, from iOS
16, ones_like as float, zeros_like with the dtype of its input, and
full_like, new_full, fill and scalar_tensor with the type of their fill
value, instead of with the dtype torch gives the result. So an int or
bool tensor came out as float: index math on it failed to lower, or lost
integers above 2048 at float16 precision. A float tensor could come out
as int or bool too: the zeros_like(mask, dtype=float) that
nn.MultiheadAttention makes from a bool mask stayed bool, and attention
came out wrong. The partitioner refused the int and bool cases, which
moved them to the CPU, and with lower_full_graph=True it failed models
that worked before, such as x + torch.full((1, 16), 0).

The backend now converts these ops with the dtype in the node's
metadata, and the partitioner no longer refuses them. A TorchScript
node, which has no such metadata, and a fill value known only at run
time keep the coremltools converter.
---
 backends/apple/coreml/compiler/torch_ops.py   | 89 +++++++++++++++++++
 .../coreml/partition/coreml_partitioner.py    | 63 -------------
 .../coreml/test/test_coreml_partitioner.py    | 48 ++++++----
 3 files changed, 120 insertions(+), 80 deletions(-)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index 7d222149d54..2ddd8e3e6b9 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -13,9 +13,11 @@
 from coremltools import _logger
 from coremltools.converters.mil import Builder as mb
 from coremltools.converters.mil.frontend import _utils
+from coremltools.converters.mil.frontend.torch import ops as _ct_ops
 from coremltools.converters.mil.frontend.torch.ops import (
     _get_inputs,
     _get_kwinputs,
+    _make_fill_op,
     noop,
     NUM_TO_NUMPY_DTYPE,
     NUM_TO_TORCH_DTYPE,
@@ -65,6 +67,93 @@ def gelu(context, node):
     context.add(mb.gelu(x=inputs[0], mode=mode, name=node.name))
 
 
+def _torch_dtype(node):
+    """The dtype torch gives the output of node, if node carries it, as a node of
+    an ExportedProgram does."""
+    val = (node.meta or {}).get("val")
+    return val.dtype if isinstance(val, _torch.Tensor) else None
+
+
+def _core_ml_type(node):
+    """The numpy type Core ML uses for the dtype torch gives the output of node."""
+    dtype = _torch_dtype(node)
+    if dtype == _torch.bool:
+        return np.bool_
+    return np.float32 if dtype.is_floating_point else np.int32
+
+
+def _add_fill(context, node, shape, value):
+    context.add(_make_fill_op(shape, _core_ml_type(node)(value), node.name))
+
+
+# coremltools gives these fills the type of their fill value or of their input,
+# or makes them float, instead of the dtype torch gives them, so an int or bool
+# fill could come out as float. Not fixed in coremltools yet, except for
+# zeros_like in https://github.com/apple/coremltools/pull/2865. A TorchScript
+# node, which carries no dtype, and a fill value known only at run time keep
+# the coremltools converter.
+@register_torch_op(override=True)
+def full(context, node):
+    inputs = _get_inputs(context, node, min_expected=2)
+    if _torch_dtype(node) is None or inputs[1].val is None:
+        return _ct_ops.full(context, node)
+    _add_fill(context, node, inputs[0], inputs[1].val)
+
+
+@register_torch_op(override=True)
+def new_full(context, node):
+    inputs = _get_inputs(context, node, min_expected=3)
+    if _torch_dtype(node) is None or inputs[2].val is None:
+        return _ct_ops.new_full(context, node)
+    _add_fill(context, node, inputs[1], inputs[2].val)
+
+
+@register_torch_op(torch_alias=["fill"], override=True)
+def full_like(context, node):
+    inputs = _get_inputs(context, node, min_expected=2)
+    if _torch_dtype(node) is None or inputs[1].val is None:
+        return _ct_ops.full_like(context, node)
+    _add_fill(context, node, mb.shape(x=inputs[0]), inputs[1].val)
+
+
+@register_torch_op(torch_alias=["empty", "empty.memory_format"], override=True)
+def zeros(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.zeros(context, node)
+    _add_fill(context, node, _get_inputs(context, node, min_expected=1)[0], 0)
+
+
+@register_torch_op(torch_alias=["new_empty"], override=True)
+def new_zeros(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.new_zeros(context, node)
+    _add_fill(context, node, _get_inputs(context, node, min_expected=2)[1], 0)
+
+
+@register_torch_op(torch_alias=["empty_like"], override=True)
+def zeros_like(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.zeros_like(context, node)
+    x = _get_inputs(context, node, min_expected=1)[0]
+    _add_fill(context, node, mb.shape(x=x), 0)
+
+
+@register_torch_op(override=True)
+def ones_like(context, node):
+    if _torch_dtype(node) is None:
+        return _ct_ops.ones_like(context, node)
+    x = _get_inputs(context, node, min_expected=1)[0]
+    _add_fill(context, node, mb.shape(x=x), 1)
+
+
+@register_torch_op(override=True)
+def scalar_tensor(context, node):
+    x = _get_inputs(context, node, expected=[1, 5])[0]
+    if _torch_dtype(node) is None or x.val is None:
+        return _ct_ops.scalar_tensor(context, node)
+    context.add(mb.const(val=_core_ml_type(node)(x.val), name=node.name))
+
+
 def is_fbcode():
     return not hasattr(_torch.version, "git_version")
 
diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index ce40a565faf..f13a64e11a9 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -205,62 +205,6 @@ def _is_min_max_of_bool(node: torch.fx.Node) -> bool:
     return isinstance(val, torch.Tensor) and val.dtype == torch.bool
 
 
-_FULL_TARGETS = (torch.ops.aten.full.default, exir_ops.edge.aten.full.default)
-
-
-# Where the fill value is, for ops that coremltools types by it.
-_FILL_VALUE_INDEX = {
-    torch.ops.aten.full_like.default: 1,
-    torch.ops.aten.fill.Scalar: 1,
-    torch.ops.aten.scalar_tensor.default: 0,
-    exir_ops.edge.aten.full_like.default: 1,
-    exir_ops.edge.aten.fill.Scalar: 1,
-    exir_ops.edge.aten.scalar_tensor.default: 0,
-}
-
-
-_DECOMPOSED_FILL_TARGETS = (
-    torch.ops.aten.zeros.default,
-    torch.ops.aten.new_zeros.default,
-    torch.ops.aten.zeros_like.default,
-    torch.ops.aten.ones_like.default,
-    torch.ops.aten.new_full.default,
-    exir_ops.edge.aten.zeros.default,
-    exir_ops.edge.aten.new_zeros.default,
-    exir_ops.edge.aten.zeros_like.default,
-    exir_ops.edge.aten.ones_like.default,
-    exir_ops.edge.aten.new_full.default,
-)
-
-
-def _fills_int_or_bool_tensor_with_wrong_type(node: torch.fx.Node) -> bool:
-    """coremltools 9.0 often ignores the dtype of an op that fills a tensor.
-
-    It makes ``full`` without a dtype, ``zeros``, ``new_zeros`` and, from iOS
-    16, ``ones_like`` as float, ``zeros_like`` with the dtype of its input, and
-    ``full_like``, ``new_full``, ``fill`` and ``scalar_tensor`` with the type
-    of their fill value. An int or bool tensor made as float silently loses
-    values above 2048 at float16 precision, and fails to lower in an op that
-    does not promote types, such as ``index_put``. Ops in
-    ``_DECOMPOSED_FILL_TARGETS`` that make one are refused even when they would
-    work, so lowering decomposes them into ``full`` or ``full_like`` with the
-    dtype given.
-    """
-    val = node.meta.get("val")
-    if not isinstance(val, torch.Tensor) or val.dtype.is_floating_point:
-        return False
-    if node.target in _DECOMPOSED_FILL_TARGETS:
-        return True
-    if node.target in _FULL_TARGETS:
-        return node.kwargs.get("dtype") is None
-    if node.target in _FILL_VALUE_INDEX:
-        value = node.args[_FILL_VALUE_INDEX[node.target]]
-        return isinstance(value, (bool, int, float)) and type(value) is not (
-            bool if val.dtype == torch.bool else int
-        )
-    return False
-
-
 _INDEX_PUT_TARGETS = (
     torch.ops.aten.index_put.default,
     exir_ops.edge.aten.index_put.default,
@@ -581,13 +525,6 @@ def should_override_support(self, node) -> bool:
             )
             return True
 
-        if _fills_int_or_bool_tensor_with_wrong_type(node):
-            self.log_once(
-                f"{node.target} making an int or bool tensor gets the wrong dtype "
-                "in CoreML.  Overriding op support."
-            )
-            return True
-
         if _is_index_put_coremltools_cannot_convert(node, self.ios18):
             self.log_once(
                 f"{node.target} with these indices is not supported by CoreML.  "
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 9e95b6d71c5..ae803022073 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -14,6 +14,7 @@
 
 import torch
 import torchvision
+from coremltools.converters.mil.mil import types
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 from executorch.backends.apple.coreml.partition import CoreMLPartitioner
@@ -592,16 +593,16 @@ def forward(self, x):
             out = program.load_method("forward").execute((x,))[0]
             self.assertTrue(torch.allclose(out, model(x), atol=1e-2))
 
-    def test_ops_coremltools_gets_wrong_are_refused(self):
+    def test_ops_coremltools_gets_wrong(self):
         """
-        coremltools makes a filled int or bool tensor with the wrong dtype in
-        some cases, such as ``full`` without a dtype or ``zeros`` with one.
-        Index math with such a tensor then fails to lower, or loses integers
-        above 2048 to float16 precision.  It also fails to lower some forms of
-        ``index_put``, such as ``y[:, idx] = v`` with two entries in ``idx``.
-        The partitioner must refuse those ops, so they run on the portable
-        backend, or lowering decomposes a fill into a ``full`` with the dtype
-        given.
+        coremltools makes a filled int or bool tensor as float in some cases,
+        such as ``full`` without a dtype or ``zeros`` with one.  Index math
+        with such a tensor then fails to lower, or loses integers above 2048
+        to float16 precision, so the backend gives fills the dtype torch gives
+        them.  coremltools also fails some forms of ``index_put``, such as
+        ``y[:, idx] = v`` with two entries in ``idx``, or writes them to the
+        wrong place, so the partitioner refuses those, and they run on the
+        portable backend.
         """
 
         class Positions(torch.nn.Module):
@@ -616,12 +617,21 @@ def forward(self, x, mask):
 
         class Counts(torch.nn.Module):
             def forward(self, ids):
-                return (ids + torch.zeros((2, 3), dtype=torch.long)) * 1001
+                ones = torch.ones_like(ids) + ids.new_zeros((2, 3))
+                twos = ids.new_full((2, 3), 2.0) + torch.empty_like(ids).fill_(0.0)
+                zeros = torch.zeros((2, 3), dtype=torch.long) + torch.zeros_like(ids)
+                count = torch.scalar_tensor(1001.0, dtype=torch.long)
+                return (ids + ones + twos + zeros) * count
 
         class Flags(torch.nn.Module):
             def forward(self, mask):
                 return mask & torch.full_like(mask, 1)
 
+        # Turns a bool mask into a float one with zeros_like(mask, dtype=float).
+        class Attention(torch.nn.MultiheadAttention):
+            def forward(self, x, mask):
+                return super().forward(x, x, x, attn_mask=mask, need_weights=False)[0]
+
         class Scatter(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -652,9 +662,15 @@ def forward(self, x, v):
 
         mask = torch.arange(16)[None] % 3 != 0
         cases = [
-            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 1),
-            (Counts(), (torch.arange(6).reshape(2, 3),), "aten.full.default", 0),
-            (Flags(), (mask,), "aten.full_like.default", 1),
+            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 0),
+            (Counts(), (torch.arange(6).reshape(2, 3),), "aten.zeros.default", 0),
+            (Flags(), (mask,), "aten.full_like.default", 0),
+            (
+                Attention(16, 2, batch_first=True),
+                (torch.randn(1, 6, 16), torch.ones(6, 6, dtype=torch.bool).triu(1)),
+                "aten.zeros_like.default",
+                0,
+            ),
             (
                 Scatter(),
                 (torch.rand(3, 4), torch.rand(3, 2)),
@@ -688,10 +704,8 @@ def forward(self, x, v):
                     self.assertEqual(out.dtype, ref.dtype)
                     self.assertTrue(torch.allclose(out, ref, atol=1e-2))
         counts = torch.export.export(Counts(), cases[1][1], strict=True)
-        self.assertNotIn(
-            torch.ops.aten.zeros.default,
-            CoreMLPartitioner().ops_to_not_decompose(counts)[0],
-        )
+        program = ct.convert(counts.run_decompositions({}), convert_to="milinternal")
+        self.assertEqual(program.functions["main"].outputs[0].dtype, types.int32)
 
     def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
         """

From 8ac3daec7f3caa37d0f8d3e265f9e00d4081f76b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sat, 31 Oct 2026 22:26:46 -0700
Subject: [PATCH 36/44] Say why Core ML refuses an op when
 lower_full_graph=True

The partitioner refuses some ops that coremltools converts, such as an
index_put form coremltools gets wrong, or an op on a bfloat16 tensor.
With lower_full_graph=True, lowering then failed with "CoreML does not
support the op index_put" and advice to ask coremltools for the op,
although coremltools has it. The real reason only went to the info log.

Lowering now fails with the reason, for example "aten.index_put.default
with these indices is not supported by CoreML. Set
lower_full_graph=False in the CoreMLPartitioner to run it outside of
CoreML." The info log keeps the same reasons, now with op names instead
of their schemas.
---
 .../coreml/partition/coreml_partitioner.py    | 54 +++++++------------
 .../coreml/test/test_coreml_partitioner.py    |  8 ++-
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index f13a64e11a9..661efb69a3a 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -464,7 +464,9 @@ def should_skip_op_for_delegation(self, node_target_name: str) -> bool:
             return True
         return False
 
-    def should_override_support(self, node) -> bool:
+    def override_reason(self, node) -> Optional[str]:
+        """Why CoreML cannot take node, even when coremltools converts its op."""
+        name = getattr(node.target, "__name__", node.target)
         # https://github.com/apple/coremltools/issues/2573
         if (
             node.target
@@ -477,20 +479,14 @@ def should_override_support(self, node) -> bool:
             and "alpha" in node.kwargs
             and node.kwargs["alpha"] != 1
         ):
-            self.log_once(
-                "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML.  Overriding support."
-            )
-            return True
+            return "torch.ops.aten.{sub, add}.Tensor with alpha != 1 is not supported by CoreML."
 
         # Ops that are unsupported by CoreML purely on the basis of their
         # target — no per-arg conditions to check.  Grouped by upstream issue
         # so the comment trail still points at the underlying coremltools /
         # executorch bug for each entry.
         if node.target in _UNSUPPORTED_OP_TARGETS:
-            self.log_once(
-                f"{node.target} is not supported by CoreML.  Overriding op support."
-            )
-            return True
+            return f"{name} is not supported by CoreML."
 
         if _touches_unsupported_high_rank_tensor(
             node,
@@ -498,39 +494,22 @@ def should_override_support(self, node) -> bool:
                 getattr(n.target, "__name__", "").lower()
             ),
         ):
-            self.log_once(
-                f"{node.target} on a tensor above rank 5 is not supported by "
-                "CoreML.  Overriding op support."
-            )
-            return True
+            return f"{name} on a tensor above rank 5 is not supported by CoreML."
 
         if _is_arg_min_max_over_flattened_input(node):
-            self.log_once(
+            return (
                 "torch.ops.aten.{argmax, argmin}.default with dim=None is "
-                "not supported by CoreML.  Overriding op support."
+                "not supported by CoreML."
             )
-            return True
 
         if _reads_or_produces_bfloat16(node):
-            self.log_once(
-                f"{node.target} on a bfloat16 tensor is not supported by CoreML.  "
-                "Overriding op support."
-            )
-            return True
+            return f"{name} on a bfloat16 tensor is not supported by CoreML."
 
         if _is_min_max_of_bool(node):
-            self.log_once(
-                f"{node.target} over a bool tensor is not supported by CoreML.  "
-                "Overriding op support."
-            )
-            return True
+            return f"{name} over a bool tensor is not supported by CoreML."
 
         if _is_index_put_coremltools_cannot_convert(node, self.ios18):
-            self.log_once(
-                f"{node.target} with these indices is not supported by CoreML.  "
-                "Overriding op support."
-            )
-            return True
+            return f"{name} with these indices is not supported by CoreML."
 
         # TODO: enable this after bugs in ExecuTorch's partitioner are fixed
         # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args
@@ -549,7 +528,7 @@ def should_override_support(self, node) -> bool:
         #     )
         #     return True
 
-        return False
+        return None
 
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # get_attr node can always be supported on any backend
@@ -567,10 +546,17 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             is_supported = ct.converters.mil.frontend.torch.is_torch_fx_node_supported(
                 node
             )
-            if self.should_override_support(node):
+            reason = self.override_reason(node)
+            if reason is not None:
+                self.log_once(f"{reason}  Overriding op support.")
                 is_supported = False
 
             if not is_supported:
+                if self.lower_full_graph and reason is not None:
+                    raise NotImplementedError(
+                        f"{reason}  Set lower_full_graph=False in the "
+                        "CoreMLPartitioner to run it outside of CoreML."
+                    )
                 if self.lower_full_graph:
                     raise NotImplementedError(
                         f"""CoreML does not support the op {node_target_name}, but you have set lower_full_graph=True in the CoreMLPartitioner.
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index ae803022073..e4d049a20d3 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -602,7 +602,8 @@ def test_ops_coremltools_gets_wrong(self):
         them.  coremltools also fails some forms of ``index_put``, such as
         ``y[:, idx] = v`` with two entries in ``idx``, or writes them to the
         wrong place, so the partitioner refuses those, and they run on the
-        portable backend.
+        portable backend.  With lower_full_graph=True, lowering fails and says
+        why.
         """
 
         class Positions(torch.nn.Module):
@@ -706,6 +707,11 @@ def forward(self, x, v):
         counts = torch.export.export(Counts(), cases[1][1], strict=True)
         program = ct.convert(counts.run_decompositions({}), convert_to="milinternal")
         self.assertEqual(program.functions["main"].outputs[0].dtype, types.int32)
+        with self.assertRaisesRegex(NotImplementedError, "with these indices"):
+            executorch.exir.to_edge_transform_and_lower(
+                torch.export.export(Scatter3d(), cases[-1][1], strict=True),
+                partitioner=[CoreMLPartitioner(lower_full_graph=True)],
+            )
 
     def test_int8_made_for_the_cpu_runs_on_the_cpu(self):
         """

From 4d81fbf8627abe6a87a0244907de1085e11e39c2 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 01:03:42 -0700
Subject: [PATCH 37/44] Stop the Core ML dry run when a drop cannot change the
 program

The dry run in ops_to_not_decompose drops kept ops that end up on the
CPU, and runs again, since a dropped op gets decomposed and can change
the program. A dropped op with no decomposition, such as the cast of a
uint8 input, changes nothing, so the next round repeated the same
decomposition and search only to find nothing more to drop. The dry run
now stops there, and builds the decomposition table once.

Its docstring also says that it does not see transform passes, or the
ops an earlier partitioner takes.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 661efb69a3a..f07cb8627d2 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -895,13 +895,16 @@ def _drop_ops_left_outside(
         not kept. That can make new calls of a kept op that CoreML rejects, such
         as the matmul of a bfloat16 linear, or change which ops read a uint8 or
         int8 tensor from outside of CoreML. So this decomposes the program the
-        same way and runs the same search, again after each drop.
+        same way and runs the same search, again after each drop that changes
+        the decomposition. It does not see transform passes, or the ops an
+        earlier partitioner takes.
         """
         kept_inside = self._placeholders_kept_inside(ep)
         skip = set(self.skip_ops_for_coreml_delegation)
         not_kept = set(rejected)
+        decompositions = _default_decomposition_table()
         while keep:
-            table = _default_decomposition_table()
+            table = dict(decompositions)
             for op in keep:
                 table.pop(op, None)
             # Decomposing only views and blocklisted ops changes nothing here.
@@ -944,4 +947,6 @@ def _drop_ops_left_outside(
                 )
             not_kept |= late
             keep = [op for op in keep if op not in late]
+            if not any(op in decompositions for op in late):
+                break
         return keep

From c3153234aa94fe2c579e91bcc3ed31d442e0b02c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 02:20:09 -0800
Subject: [PATCH 38/44] Name the ops that split Core ML state

When ops in more than one place use a buffer that Core ML takes over,
lowering fails, since each place would use its own copy. The error said
only that the buffer is used in more than one place. For a model whose
cache update runs on the CPU, such as an index_put Core ML cannot
convert, it did not say which op that is.

The error now lists each op that uses the buffer and where it runs, and
the docs for take_over_mutable_buffer mention the error.
---
 backends/apple/coreml/partition/coreml_partitioner.py | 11 ++++++++---
 backends/apple/coreml/test/test_coreml_partitioner.py |  4 +++-
 docs/source/backends/coreml/coreml-partitioner.md     |  2 +-
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index f07cb8627d2..635f37c37a2 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -772,17 +772,22 @@ def _check_state_is_used_in_one_place(self, ep: ExportedProgram) -> None:
             if all(user.meta.get("delegation_tag") is None for user in users):
                 continue
             users.append(updates[signature.inputs_to_buffers[name]])
-            if len({user.meta.get("delegation_tag") for user in users}) > 1:
+            tags = {user: user.meta.get("delegation_tag") for user in users}
+            if len(set(tags.values())) > 1:
                 hint = (
                     ""
                     if self.lower_full_graph
                     else "  Set take_over_mutable_buffer=False in the "
                     "CoreMLPartitioner to keep the state outside of CoreML."
                 )
+                places = ", ".join(
+                    f"{user.name} ({'CoreML ' + tag if tag else 'CPU'})"
+                    for user, tag in tags.items()
+                )
                 raise NotImplementedError(
                     f"CoreML takes over {name} as state, but it is used in more "
-                    "than one place: in two CoreML partitions, or both in and "
-                    f"outside of CoreML.  Each place would use its own copy.{hint}"
+                    f"than one place: {places}.  Each place would use its own "
+                    f"copy.{hint}"
                 )
 
     def _targets_ios18(self) -> bool:
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index e4d049a20d3..7f03c37eca9 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -983,7 +983,9 @@ def forward(self, image):
                 model=model_class.__name__, ios18=compile_specs is not None
             ):
                 ep = torch.export.export(model_class().eval(), (x,), strict=True)
-                with self.assertRaisesRegex(NotImplementedError, "b_state as state"):
+                with self.assertRaisesRegex(
+                    NotImplementedError, "b_state as state, but .* place: \\w+ \\("
+                ):
                     executorch.exir.to_edge_transform_and_lower(
                         ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
                     )
diff --git a/docs/source/backends/coreml/coreml-partitioner.md b/docs/source/backends/coreml/coreml-partitioner.md
index 7c249d539fa..87f947a3d90 100644
--- a/docs/source/backends/coreml/coreml-partitioner.md
+++ b/docs/source/backends/coreml/coreml-partitioner.md
@@ -5,7 +5,7 @@ The Core ML partitioner API allows for configuration of the model delegation to
 
  - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by Core ML.  By default, all ops that Core ML supports will be delegated.  See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation.
 - `compile_specs`: A list of `CompileSpec`s for the Core ML backend.  These control low-level details of Core ML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32).  These are discussed more below.
-- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate).  If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood.  Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0.
+- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate).  If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood.  Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0.  If ops in more than one place, such as two Core ML partitions, or Core ML and the CPU, use a buffer, lowering fails and names them, since each place would use its own copy; set `take_over_mutable_buffer` to False for such a model.
 - `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate.  If set to False, constant data is passed to the Core ML delegate as inputs.  By default, take_over_constant_data=True.  A uint8 buffer that the model casts, such as a lookup table, needs take_over_constant_data=False: coremltools cannot cast a uint8 constant, while with False the cast runs on the CPU.
 - `lower_full_graph`: A boolean that indicates whether the entire graph must be lowered to Core ML.  If set to True and Core ML does not support an op, an error is raised during lowering.  If set to False and Core ML does not support an op, the op is executed on the CPU by ExecuTorch.  Although setting `lower_full_graph`=False can allow a model to lower where it would otherwise fail, it can introduce performance overhead in the model when there are unsupported ops.  You will see warnings about unsupported ops during lowering if there are any.  By default, `lower_full_graph`=False.
 

From 1d1953372e667d8e40e54ab7abc4e52aae1a309c Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 09:59:02 -0800
Subject: [PATCH 39/44] Leave the shared constant copy out of the Core ML
 change

This reverts commit 6ba764d4e1, which made duplicate_constant_node also
replace a constant that an op reads from a list.

duplicate_constant_node also copies a mutated buffer that a delegate
owns, for a CPU op that reads it. A list reader such as cat used to keep
the original buffer, so lowering failed. With the change it reads a copy
that never sees the delegate's updates. That affects every backend that
shares state, so the change needs its own review, apart from Core ML.
---
 .../duplicate_constant_node_pass.py           | 21 ++++++++++++--
 exir/backend/test/test_passes.py              | 28 -------------------
 2 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
index 1d76e541a04..07c3eb5e57a 100644
--- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
+++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py
@@ -96,9 +96,26 @@ def duplicate_constant_node(
                     copied_constant_or_attribute_node.meta["val"] = (
                         constant_or_attribute_node.meta["val"]
                     )
-                    users[ith].replace_input_with(
-                        constant_or_attribute_node, copied_constant_or_attribute_node
+                    new_args = tuple(
+                        [
+                            (
+                                arg
+                                if arg != constant_or_attribute_node
+                                else copied_constant_or_attribute_node
+                            )
+                            for arg in users[ith].args
+                        ]
                     )
+                    new_kwargs = {
+                        key: (
+                            value
+                            if value != constant_or_attribute_node
+                            else copied_constant_or_attribute_node
+                        )
+                        for key, value in users[ith].kwargs.items()
+                    }
+                    users[ith].args = new_args
+                    users[ith].kwargs = new_kwargs
                     # The copy keeps the original target, so it reads the
                     # original tensor and needs no data of its own.
                     new_input_specs.append(
diff --git a/exir/backend/test/test_passes.py b/exir/backend/test/test_passes.py
index 3ded45e7ce2..815a5b5bbed 100644
--- a/exir/backend/test/test_passes.py
+++ b/exir/backend/test/test_passes.py
@@ -86,31 +86,3 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         self.assertEqual(len(copied_nodes), 1)
         self.assertNotEqual(user.kwargs["other"], buffer_node)
-
-    def test_duplicate_constant_node_in_a_list_argument(self) -> None:
-        """
-        An op such as cat or index_put reads its tensors from a list, and the
-        copy has to replace the constant there too.
-        """
-
-        class CatConst(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.register_buffer("const", torch.ones(2))
-
-            def forward(self, x):
-                return torch.cat([x, self.const]), torch.cat([self.const, x])
-
-        edge = exir.to_edge(export(CatConst(), (torch.ones(2),), strict=True))
-        program = edge.exported_program()
-        name = next(
-            node.name
-            for node in program.graph.nodes
-            if node.op == "placeholder" and is_buffer(program, node)
-        )
-
-        copied_nodes = duplicate_constant_node(program, name)
-
-        users = {node.name: len(node.users) for node in program.graph.nodes}
-        self.assertEqual(users[name], 1)
-        self.assertEqual([users[copy] for copy in copied_nodes], [1])

From cdd72bc066d050dc8a6a54895c7f669c39807ae0 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 07:09:59 -0800
Subject: [PATCH 40/44] Keep float16 fills float16 in Core ML

The Core ML fill converters gave every floating fill a float32 value, so
a float16 zeros_like, ones_like or empty_like, which coremltools made
float16 before, came out float32. The ops that read it then promote to
float32, which adds casts, and with FLOAT32 compute precision makes
float32 results where torch has float16.

A float16 fill is now made as float16.
---
 backends/apple/coreml/compiler/torch_ops.py  |  2 ++
 backends/apple/coreml/test/test_torch_ops.py | 17 +++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index 2ddd8e3e6b9..6d28551fb2d 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -79,6 +79,8 @@ def _core_ml_type(node):
     dtype = _torch_dtype(node)
     if dtype == _torch.bool:
         return np.bool_
+    if dtype == _torch.float16:
+        return np.float16
     return np.float32 if dtype.is_floating_point else np.int32
 
 
diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py
index 7968a6d391e..a7c33410cd9 100644
--- a/backends/apple/coreml/test/test_torch_ops.py
+++ b/backends/apple/coreml/test/test_torch_ops.py
@@ -9,6 +9,7 @@
 import executorch.exir
 
 import torch
+from coremltools.converters.mil.mil import types
 from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
@@ -366,6 +367,22 @@ def test_gelu_tanh(self):
         et_prog = delegated_program.to_executorch()
         self._compare_outputs(et_prog, model, example_inputs, atol=1e-5, rtol=1e-5)
 
+    def test_float16_fill(self):
+        """A float16 fill stays float16, so the ops that read it stay float16."""
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return x * torch.ones_like(x) + torch.zeros_like(x)
+
+        ep = torch.export.export(Model(), (torch.rand(2, 3, dtype=torch.float16),))
+        program = ct.convert(
+            ep.run_decompositions({}),
+            convert_to="milinternal",
+            compute_precision=ct.precision.FLOAT32,
+            minimum_deployment_target=ct.target.iOS16,
+        )
+        self.assertEqual(program.functions["main"].outputs[0].dtype, types.fp16)
+
     @unittest.skipIf(is_fbcode(), "The coremltools in fbcode has no rms_norm.")
     def test_rms_norm_without_eps(self):
         """

From d5f92a4413540868bb38cc5da909d0d9d55bbdbb Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 08:32:05 -0800
Subject: [PATCH 41/44] Test an index_put form that Core ML keeps through the
 reordering rule

The index_put check keeps a whole-slice form with a 0-d index in Core ML
when the values vary only in the dims after the index. The tests only
had forms that the check refuses, so a rule that refused too much, or
let a wrong form through, went unnoticed. The 3-D case now also writes
such values, which must stay in Core ML and match eager.
---
 backends/apple/coreml/test/test_coreml_partitioner.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 7f03c37eca9..62bcb7bf54f 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -655,10 +655,13 @@ class Scatter3d(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.register_buffer("last", torch.tensor(2))
+                self.register_buffer("mid", torch.tensor(1))
 
             def forward(self, x, v):
                 y = x.clone()
                 torch.ops.aten.index_put_(y, [None, None, self.last], v)
+                # Values that vary only after the index stay in Core ML.
+                torch.ops.aten.index_put_(y, [None, self.mid], v[0] * torch.arange(3.0))
                 return y * 2
 
         mask = torch.arange(16)[None] % 3 != 0

From bcef1277c10f4a3aafa7a178e6bd80a5f40a98da Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 08:32:05 -0800
Subject: [PATCH 42/44] Clarify the Core ML docs for uint8 tables and shared
 state

The Core ML docs pointed a model with a uint8 lookup table that it casts
to take_over_constant_data=False. That setting passes every weight to
Core ML as an input, and runs the dequantize of int8 or uint8 weights on
the CPU. The docs now suggest storing such a table as float16 or int32
first, and say what the setting costs.

The note on take_over_mutable_buffer said lowering fails whenever Core
ML and the CPU both use a buffer. It fails only for a buffer that Core
ML takes over, which a buffer that only CPU ops read is not.
---
 docs/source/backends/coreml/coreml-partitioner.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/source/backends/coreml/coreml-partitioner.md b/docs/source/backends/coreml/coreml-partitioner.md
index 87f947a3d90..cf009e78f32 100644
--- a/docs/source/backends/coreml/coreml-partitioner.md
+++ b/docs/source/backends/coreml/coreml-partitioner.md
@@ -5,8 +5,8 @@ The Core ML partitioner API allows for configuration of the model delegation to
 
  - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by Core ML.  By default, all ops that Core ML supports will be delegated.  See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation.
 - `compile_specs`: A list of `CompileSpec`s for the Core ML backend.  These control low-level details of Core ML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32).  These are discussed more below.
-- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate).  If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood.  Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0.  If ops in more than one place, such as two Core ML partitions, or Core ML and the CPU, use a buffer, lowering fails and names them, since each place would use its own copy; set `take_over_mutable_buffer` to False for such a model.
-- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate.  If set to False, constant data is passed to the Core ML delegate as inputs.  By default, take_over_constant_data=True.  A uint8 buffer that the model casts, such as a lookup table, needs take_over_constant_data=False: coremltools cannot cast a uint8 constant, while with False the cast runs on the CPU.
+- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate).  If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood.  Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0.  If a buffer that Core ML takes over is also used on the CPU or in another Core ML partition, lowering fails and names the ops, since each place would use its own copy; set `take_over_mutable_buffer` to False for such a model.
+- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate.  If set to False, constant data is passed to the Core ML delegate as inputs.  By default, take_over_constant_data=True.  A uint8 buffer that the model casts, such as a lookup table, fails to lower with the default settings, since coremltools cannot cast a uint8 constant.  Storing such a table as float16 or int32 avoids that.  take_over_constant_data=False also works, since the cast then runs on the CPU, but it passes every weight to Core ML as an input, and runs the dequantize of int8 or uint8 weights on the CPU.
 - `lower_full_graph`: A boolean that indicates whether the entire graph must be lowered to Core ML.  If set to True and Core ML does not support an op, an error is raised during lowering.  If set to False and Core ML does not support an op, the op is executed on the CPU by ExecuTorch.  Although setting `lower_full_graph`=False can allow a model to lower where it would otherwise fail, it can introduce performance overhead in the model when there are unsupported ops.  You will see warnings about unsupported ops during lowering if there are any.  By default, `lower_full_graph`=False.
 
 

From 0cb79180d6f84793ea89716b5a3c5490d5d217af Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 11:40:12 -0800
Subject: [PATCH 43/44] Pass a shared constant to Core ML as an input when an
 op reads it from a list

From iOS 18, `mask[:, col] = True` stays in Core ML while an accumulate
write at the same `col` runs on the CPU, so both read the `col` buffer
from their index lists. Lowering then failed with "constant data node
(b_col) is tagged with (tag1) but has user (aten_index_put_default_2)
which has tag (None)": to_backend gives each place its own copy of a
shared constant, but swaps in the copy only where an op takes the
constant directly.

The partitioner now leaves such a constant untagged, so Core ML takes it
as an input, as it does for every constant with
take_over_constant_data=False. Mutable buffers are not touched.

The fills and index_put table test also lowers Scatter at iOS 18, which
fails without this change.
---
 .../coreml/partition/coreml_partitioner.py    | 28 ++++++++++++++++++
 .../coreml/test/test_coreml_partitioner.py    | 29 +++++++++++++++----
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py
index 635f37c37a2..6d8b4f9e104 100644
--- a/backends/apple/coreml/partition/coreml_partitioner.py
+++ b/backends/apple/coreml/partition/coreml_partitioner.py
@@ -410,6 +410,33 @@ def _propose_partitions(
         left_out |= new
 
 
+def _reads_from_a_list(user: torch.fx.Node, node: torch.fx.Node) -> bool:
+    return any(
+        isinstance(arg, (list, tuple)) and any(item is node for item in arg)
+        for arg in (*user.args, *user.kwargs.values())
+    )
+
+
+def _untag_constants_shared_through_lists(ep: ExportedProgram) -> None:
+    """Untags each constant that ops in more than one place, CoreML partitions
+    or the CPU, read, when one of them reads it from a list, such as the
+    indices of an ``index_put``.
+
+    to_backend gives each place its own copy of a shared constant, but it only
+    swaps in the copy where an op takes the constant directly, so a constant
+    read from a list fails to lower. An untagged constant is passed to each
+    CoreML partition as an input instead.
+    """
+    for node in ep.graph.nodes:
+        if node.op != "placeholder" or "delegation_tag" not in node.meta:
+            continue
+        places = {user.meta.get("delegation_tag") for user in node.users}
+        if len(places) > 1 and any(
+            _reads_from_a_list(user, node) for user in node.users
+        ):
+            del node.meta["delegation_tag"]
+
+
 def _edge_op_names(target: Any) -> Set[str]:
     """The lowercase names target can have once lowered to the edge dialect,
     which partition() matches skip names against: an ATen op gains its
@@ -735,6 +762,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
         if self.take_over_constant_data:
             tag_constant_data(exported_program)
+            _untag_constants_shared_through_lists(exported_program)
         if self.take_over_mutable_buffer:
             logger.info(
                 "Core ML partitioner will take over torch mutable buffer as Core ML state, "
diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py
index 62bcb7bf54f..729c88a4907 100644
--- a/backends/apple/coreml/test/test_coreml_partitioner.py
+++ b/backends/apple/coreml/test/test_coreml_partitioner.py
@@ -665,34 +665,51 @@ def forward(self, x, v):
                 return y * 2
 
         mask = torch.arange(16)[None] % 3 != 0
+        ios18 = CoreMLBackend.generate_compile_specs(
+            minimum_deployment_target=ct.target.iOS18
+        )
         cases = [
-            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 0),
-            (Counts(), (torch.arange(6).reshape(2, 3),), "aten.zeros.default", 0),
-            (Flags(), (mask,), "aten.full_like.default", 0),
+            (Positions(), (torch.rand(1, 16, 4), mask), "aten.full.default", 0, None),
+            (Counts(), (torch.arange(6).reshape(2, 3),), "aten.zeros.default", 0, None),
+            (Flags(), (mask,), "aten.full_like.default", 0, None),
             (
                 Attention(16, 2, batch_first=True),
                 (torch.randn(1, 6, 16), torch.ones(6, 6, dtype=torch.bool).triu(1)),
                 "aten.zeros_like.default",
                 0,
+                None,
             ),
             (
                 Scatter(),
                 (torch.rand(3, 4), torch.rand(3, 2)),
                 "aten.index_put.default",
                 5,
+                None,
+            ),
+            # From iOS 18 the bool write stays in Core ML, and reads its index
+            # from the same buffer as the accumulate write on the CPU.
+            (
+                Scatter(),
+                (torch.rand(3, 4), torch.rand(3, 2)),
+                "aten.index_put.default",
+                4,
+                ios18,
             ),
             (
                 Scatter3d(),
                 (torch.rand(2, 2, 3), torch.rand(2, 1)),
                 "aten.index_put.default",
                 1,
+                None,
             ),
         ]
-        for model, inputs, cpu_op, cpu_calls in cases:
-            with self.subTest(model=type(model).__name__):
+        for model, inputs, cpu_op, cpu_calls, compile_specs in cases:
+            with self.subTest(
+                model=type(model).__name__, ios18=compile_specs is not None
+            ):
                 ep = torch.export.export(model.eval(), inputs, strict=True)
                 edge = executorch.exir.to_edge_transform_and_lower(
-                    ep, partitioner=[CoreMLPartitioner()]
+                    ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)]
                 )
                 op_names = [
                     n.target.__name__

From 3efab3af9b947eaeba9e8dd0f6b1cc21b7922a93 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 11:40:14 -0800
Subject: [PATCH 44/44] Say when the Core ML empty aliases matter

Lowering turns empty, new_empty and empty_like into _empty_dim_order
unless dim order ops are skipped, and coremltools converts that op with
its dtype. Say so where the fill converters take these names, so the
aliases do not look unused.
---
 backends/apple/coreml/compiler/torch_ops.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py
index 6d28551fb2d..f1f87d17e1a 100644
--- a/backends/apple/coreml/compiler/torch_ops.py
+++ b/backends/apple/coreml/compiler/torch_ops.py
@@ -118,6 +118,9 @@ def full_like(context, node):
     _add_fill(context, node, mb.shape(x=inputs[0]), inputs[1].val)
 
 
+# Lowering turns empty, new_empty and empty_like into _empty_dim_order, which
+# coremltools converts with its dtype, so these aliases matter only when dim
+# order ops are skipped.
 @register_torch_op(torch_alias=["empty", "empty.memory_format"], override=True)
 def zeros(context, node):
     if _torch_dtype(node) is None: