From 6dc5a9cf3e4dead9cc3c3bbc55b6c53bbbe6120a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 13:07:24 -0800 Subject: [PATCH 1/5] Refuse reductions of a 0-d tensor in the Core ML partitioner Core ML's reduce ops need an input of rank 1 or more, and coremltools converts a reduction of a 0-d tensor to one anyway. So x.sum().sum() lowers, but Core ML fails to load it with "Param 'x' has incorrect type for operator 'reduce_sum'". The partitioner now leaves sum, mean, amax, amin, max, min, argmax, argmin, any, all and var of a 0-d tensor to the CPU kernels, with one check for all of them. With a dim, these used to fail at lowering instead, with an IndexError, and they now run on the CPU too, except a sum. A std or vector norm of a 0-d tensor fails the same way, but stays in Core ML. ExecuTorch has no CPU kernel for it, and a refused call can decompose into an op that another call in the program keeps whole, which would then have no kernel at all. A sum over a given dim, such as x.sum().sum(0), stays in Core ML too, where coremltools fails to lower it with an IndexError, as before. ExecuTorch's optimized sum kernel aborts on a 0-d tensor with a dim, so refusing it would turn the lowering error into a crash at run time. The new test lowers such reductions of x.sum() and of x[0, 0], checks that they stay on the CPU while the rest is delegated, and on macOS compares the output with eager. Core ML computes in float16, so the input holds values float16 represents exactly, and the comparison also allows a relative error, since the output reaches about 60. The test also checks that a sum over a dim of a 0-d tensor stays in Core ML. --- .../coreml/partition/coreml_partitioner.py | 57 ++++++++++++ .../coreml/test/test_coreml_partitioner.py | 87 +++++++++++++++++++ 2 files changed, 144 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 4453818f956..1e6a1df2508 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -154,6 +154,56 @@ 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) +# No upstream issue yet. For example, Core ML fails to load x.sum().sum(): it gets +# the second sum as a reduce_sum of a 0-d tensor, "Param 'x' has incorrect type +# for operator 'reduce_sum'". A std or vector norm of a 0-d tensor fails too, but +# stays: it has no CPU kernel, and a refused call can decompose into an op that +# another call keeps whole. +_REDUCTION_TARGETS = frozenset( + target + for aten in (torch.ops.aten, exir_ops.edge.aten) + for target in ( + aten.sum.default, + aten.sum.dim_IntList, + aten.mean.default, + aten.mean.dim, + aten.amax.default, + aten.amin.default, + aten.max.default, + aten.max.dim, + aten.min.default, + aten.min.dim, + aten.argmax.default, + aten.argmin.default, + aten.any.default, + aten.any.dim, + aten.any.dims, + aten.all.default, + aten.all.dim, + aten.all.dims, + aten.var.default, + aten.var.dim, + aten.var.correction, + ) +) +_SUM_OVER_DIMS = (torch.ops.aten.sum.dim_IntList, exir_ops.edge.aten.sum.dim_IntList) + + +def _reduces_0d_tensor(node: torch.fx.Node) -> bool: + """Core ML cannot run the reductions above on a 0-d tensor: its reduce ops + need rank 1 or more. + + A sum over a given dim, such as x.sum().sum(0), stays and still fails to + lower, since ExecuTorch's optimized sum kernel aborts on it at run time. + """ + if node.target not in _REDUCTION_TARGETS or ( + node.target in _SUM_OVER_DIMS and len(node.args) > 1 and node.args[1] + ): + return False + val = node.args[0].meta.get("val") + return isinstance(val, torch.Tensor) and val.dim() == 0 + + def _is_view_op(op: torch._ops.OpOverload) -> bool: schema = op._schema if len(schema.arguments) == 0: @@ -242,6 +292,13 @@ def should_override_support(self, node) -> bool: ) return True + if _reduces_0d_tensor(node): + self.log_once( + f"{node.target} of a 0-d tensor is not supported by CoreML. " + "Overriding op support." + ) + return True + if _reads_or_produces_bfloat16(node): self.log_once( f"{node.target} on a bfloat16 tensor 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 0129eb1d33f..4a74469a210 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -521,6 +521,93 @@ 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_reduction_of_0d_tensor_is_skipped(self): + """ + Core ML cannot reduce a 0-d tensor. The partitioner must leave such + reductions to the CPU and still delegate the rest. + """ + + class Model(torch.nn.Module): + def forward(self, x): + total = x.sum() + flag = x[0, 0] > 0.5 + zero_d = [total.sum(), total.mean(), total.mean(0), total.amax()] + zero_d += [total.amin(), total.max(), total.max(0)[0], total.min()] + zero_d += [total.min(0)[0], total.argmax(0), total.argmin(0)] + zero_d += [x[0, 0].sum(), flag.any(), flag.any(0), flag.any((0,))] + zero_d += [flag.all(), flag.all(0), flag.all((0,))] + zero_d += [total.var(correction=0), total.var(unbiased=False)] + zero_d += [total.var(0, False)] + return x * 2 + sum(r.to(x.dtype) for r in zero_d) + + model = Model().eval() + # Core ML computes in float16, so use values it holds exactly, or + # x[0, 0] > 0.5 can differ from eager for x just above 0.5. + x = torch.rand(2, 3).half().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) + # One per reduction, in the order of the model. torch makes all() a + # logical_not of any() of a logical_not, which Core ML takes, and every + # var a var.correction. + reductions = [ + "sum.dim_IntList", + "mean.default", + "mean.dim", + "amax.default", + "amin.default", + "max.default", + "max.dim", + "min.default", + "min.dim", + "argmax.default", + "argmin.default", + "sum.dim_IntList", + "any.default", + "any.dim", + "any.dims", + "any.dims", + "any.dim", + "any.dims", + ] + ["var.correction"] * 3 + self.assertEqual( + sorted( + name + for name in op_names + if name not in ("executorch_call_delegate", "getitem") + ), + sorted(f"aten.{name}" for name in reductions), + ) + 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, rtol=1e-2)) + + # A sum of a 0-d tensor over a dim stays in Core ML, where lowering + # fails as before, since ExecuTorch's optimized sum kernel aborts on it. + class SumOverDim(torch.nn.Module): + def forward(self, x): + return x * x.sum().sum(0) + + ep = torch.export.export(SumOverDim(), (x,), strict=True) + result = CoreMLPartitioner().partition( + executorch.exir.to_edge(ep).exported_program() + ) + sums = [ + n + for n in result.tagged_exported_program.graph.nodes + if getattr(n.target, "__name__", "") == "aten.sum.dim_IntList" + ] + self.assertEqual(len(sums), 2) + self.assertTrue(all("delegation_tag" in n.meta for n in sums)) + def test_tied_embedding_is_quantized_with_its_linear(self): """A weight that is both an embedding table and an output projection lowers. From 92b1f400c80da5c861b37a1c3140964396e0a47f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 13:07:25 -0800 Subject: [PATCH 2/5] Refuse bool min and max in the Core ML partitioner coremltools 9.0 cannot compute min, max, amin, amax, minimum or maximum of bool tensors. It raises a ValueError instead of reporting the op as unsupported, so lowering fails, for example on the all-true check that torch.distributions.MultivariateNormal exports. A bool next to a float or int tensor converts, since the result is then float or int. The partitioner now leaves the calls with a bool result to the portable CPU kernels and still delegates the rest. The new test lowers each of these ops on bool tensors, checks which ops stay on the CPU, and that an amax of a float tensor and a maximum of a bool and a float tensor stay in Core ML, and on macOS compares the output with eager. Its inputs are values that float16 holds exactly, since Core ML compares x > 0.5 in float16. --- .../coreml/partition/coreml_partitioner.py | 44 +++++++++++++++++ .../coreml/test/test_coreml_partitioner.py | 48 +++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 1e6a1df2508..e5b752ab9d6 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -212,6 +212,43 @@ def _is_view_op(op: torch._ops.OpOverload) -> bool: return (alias_info is not None) and (not alias_info.is_write) +# The partitioner picks the ops to keep whole from torch ops. Refusing min.other +# and max.other there lets lowering turn them into minimum and maximum. The +# others stay whole, as core ATen ops or ops the partitioner keeps, so only their +# edge ops reach it. +_MIN_MAX_TARGETS = ( + torch.ops.aten.min.other, + torch.ops.aten.max.other, + exir_ops.edge.aten.min.default, + exir_ops.edge.aten.min.dim, + exir_ops.edge.aten.max.default, + exir_ops.edge.aten.max.dim, + exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.minimum.default, + exir_ops.edge.aten.maximum.default, +) + + +# No upstream issue yet. For example, converting x.amax() of a bool x raises +# 'Op "amax" (op_type: reduce_max) Input x="x" expects tensor or scalar of +# dtype from type domain ['fp16', 'fp32', 'int32'] but got tensor[3,2,bool]'. +def _is_min_max_of_bool(node: torch.fx.Node) -> bool: + """coremltools 9.0 cannot compute a min or max of bool tensors. + + It raises ``ValueError`` when every input is bool, so that the result is + bool, such as for the all-true check that torch.distributions exports, + instead of reporting the op as unsupported. A bool next to a float or int + tensor works, and so do int tensors. + """ + if node.target not in _MIN_MAX_TARGETS: + return False + val = node.meta.get("val") + if isinstance(val, (list, tuple)): + val = val[0] + return isinstance(val, torch.Tensor) and val.dtype == torch.bool + + class _OperatorsSupportedForCoreMLBackend(OperatorSupportBase): def __init__( self, @@ -306,6 +343,13 @@ def should_override_support(self, node) -> bool: ) return True + if _is_min_max_of_bool(node): + self.log_once( + f"{node.target} on bool tensors 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 4a74469a210..56f64da9fec 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -521,6 +521,54 @@ 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 every form of a min or max of bool tensors, such + as the all-true check torch.distributions exports. The partitioner + must leave those ops to the portable backend and still delegate the + rest, including a min or max with a float input. + """ + + class Model(torch.nn.Module): + def forward(self, x): + ok = (x.tril() == x).view(1, -1) + no = ok.flip(-1) + reduced = [ok.min(), ok.max(), ok.min(-1)[0], ok.max(-1)[0]] + reduced += [ok.amin(-1), ok.amax(-1)] + paired = [torch.min(ok, no), torch.max(ok, no)] + paired += [torch.minimum(ok, no), torch.maximum(ok, no)] + flags = sum(r.to(x.dtype).sum() for r in reduced) + both = sum(p.to(x.dtype) for p in paired).view(x.shape) + floats = x.amax(-1, keepdim=True) + torch.maximum(x > 0.5, x) + return x * 2 + flags + both + floats + + model = Model().eval() + # Core ML computes in float16, so use values it holds exactly, or + # x > 0.5 can differ from eager for x just above 0.5. + x = torch.rand(14, 14).half().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) + for name in ("min.default", "max.default", "min.dim", "max.dim"): + self.assertIn(f"aten.{name}", op_names) + self.assertIn("aten.amin.default", op_names) + # The float amax and maximum stay in Core ML, and lowering turns the + # min and max of two tensors into minimum and maximum. + self.assertEqual(op_names.count("aten.amax.default"), 1) + self.assertEqual(op_names.count("aten.minimum.default"), 2) + self.assertEqual(op_names.count("aten.maximum.default"), 2) + 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_reduction_of_0d_tensor_is_skipped(self): """ Core ML cannot reduce a 0-d tensor. The partitioner must leave such From ee970ae615a9aa0df5a47b221cd94a5d304fcddc Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 13:07:25 -0800 Subject: [PATCH 3/5] Refuse index_put forms that coremltools cannot convert coremltools converts y[:, idx] = v only when idx has a single entry. With idx = torch.tensor([0, 2]), lowering fails with "Input squeeze_0 has rank 1 != other inputs rank 0". accumulate=True, a bool tensor next to a slice below iOS 18, a bool index next to a slice or next to other indices, index tensors that broadcast, and a size known only at run time next to the slice fail too. With a 0-d index or several indices next to a slice, coremltools lays out the written part differently from torch, so values land in the wrong place unless they vary only in the dims after the last index. A bool tensor written without a slice comes back as int32, so ~ on the result reads it wrong. The partitioner now sends these forms to the portable CPU kernels, and logs which form it found. Forms coremltools handles, such as a single-entry index for a cache update, or a bool mask on its own, stay in Core ML. The new test counts the index_put calls left on the CPU for each form, with a static and a dynamic batch size. It includes a write with two same-shape indices and no slice, which stays in Core ML, a bool write without a slice, and a write with two indices next to a slice. On macOS it compares the output with eager, at iOS 18 only on macOS 15. The values it compares with 0.5 are ones float16 holds exactly. --- .../coreml/partition/coreml_partitioner.py | 100 +++++++++++++++ .../coreml/test/test_coreml_partitioner.py | 116 ++++++++++++++++++ 2 files changed, 216 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index e5b752ab9d6..849327df2c3 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -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 torch.export.exported_program import ExportedProgram +from torch.fx.experimental.symbolic_shapes import statically_known_true, sym_eq from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase @@ -249,18 +250,102 @@ def _is_min_max_of_bool(node: torch.fx.Node) -> bool: return isinstance(val, torch.Tensor) and val.dtype == torch.bool +_INDEX_PUT_TARGETS = ( + torch.ops.aten.index_put.default, + exir_ops.edge.aten.index_put.default, +) + + +# https://github.com/apple/coremltools/issues/2141 covers a bool index, which +# coremltools 9.0 converts only on its own. No upstream issue covers the other +# forms yet. For example, y[:, idx] = v with idx = torch.tensor([0, 2]) raises +# "Input squeeze_0 has rank 1 != other inputs rank 0". +def _index_put_refusal( + node: torch.fx.Node, minimum_deployment_target: Optional[ct.target] +) -> Optional[str]: + """The form of an ``index_put`` that coremltools 9.0 cannot convert, such + as "with accumulate=True next to a slice", or None for any other node. + + 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, and a non-bool ``y`` below iOS 18. For a + size known only at run time, it raises "Cannot add const". 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, and a bool ``y`` + comes back as int32. It raises or writes the wrong place for any other + form, such as ``idx`` with two entries, instead of reporting the op as + unsupported. + """ + if node.target not in _INDEX_PUT_TARGETS: + return None + y, values = node.args[0].meta.get("val"), node.args[2].meta.get("val") + indices = [i if i is None else i.meta.get("val") for i in node.args[1]] + given = {dim: indices[dim] for dim, i in enumerate(node.args[1]) if i is not None} + # Export always sets val, so an index_put without it is left to coremltools. + if not given or not all( + isinstance(t, torch.Tensor) for t in (y, values, *given.values()) + ): + return None + if len(given) == len(indices): + if y.dtype == torch.bool: + return "on a bool tensor without a slice" + if 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 + ): + return "with indices that differ in shape or include a bool mask" + return None + accumulate = ( + node.args[3] if len(node.args) > 3 else node.kwargs.get("accumulate", False) + ) + if accumulate: + return "with accumulate=True next to a slice" + if any( + index.dtype == torch.bool + or index.dim() > 1 + or not statically_known_true(index.numel() == 1) + for index in given.values() + ): + return ( + "with an index other than an int tensor of shape () or (1,) next to a slice" + ) + selected = [1 if dim in given else size for dim, size in enumerate(y.shape)] + if any(isinstance(size, torch.SymInt) for size in selected): + return "with a size known only at run time next to a slice" + reorders = len(given) > 1 or any(index.dim() == 0 for index in given.values()) + after = len(selected) - 1 - max(given) + if reorders and not all( + statically_known_true(v == 1) + for v in values.shape[: max(0, values.dim() - after)] + ): + return ( + "with a 0-d index or several indices, and values that vary before " + "the last index" + ) + if y.dtype == torch.bool and ( + minimum_deployment_target is None or minimum_deployment_target < ct.target.iOS18 + ): + return "on a bool tensor below minimum_deployment_target=ct.target.iOS18" + return None + + class _OperatorsSupportedForCoreMLBackend(OperatorSupportBase): def __init__( self, skip_ops_for_coreml_delegation: Optional[List[str]] = None, lower_full_graph: bool = False, log: bool = False, + minimum_deployment_target: Optional[ct.target] = None, ) -> 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.minimum_deployment_target = minimum_deployment_target self._logged_msgs = set() self._log = log @@ -350,6 +435,14 @@ def should_override_support(self, node) -> bool: ) return True + refusal = _index_put_refusal(node, self.minimum_deployment_target) + if refusal is not None: + self.log_once( + f"{node.target} {refusal} 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 @@ -444,6 +537,11 @@ def __init__( self.take_over_mutable_buffer = take_over_mutable_buffer self.lower_full_graph = lower_full_graph self.take_over_constant_data = take_over_constant_data + self._minimum_deployment_target = ( + CoreMLBackend.min_deployment_target_from_compile_specs( + self.delegation_spec.compile_specs + ) + ) self._logged_msgs = set() if self.lower_full_graph: @@ -500,6 +598,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=True, + minimum_deployment_target=self._minimum_deployment_target, ), allows_single_node_partition=True, ) @@ -539,6 +638,7 @@ def ops_to_not_decompose( self.skip_ops_for_coreml_delegation, self.lower_full_graph, log=False, + minimum_deployment_target=self._minimum_deployment_target, ) # CoreML prevents certain ops (like triu) from lowering to CoreML when put in the ExecuTorch op namespace diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 56f64da9fec..9ae748cfa79 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -17,6 +17,9 @@ 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_VALID_TEST_RUNTIME, +) from executorch.exir.backend.utils import format_delegated_graph @@ -569,6 +572,119 @@ def forward(self, x): out = program.load_method("forward").execute((x,))[0] self.assertTrue(torch.allclose(out, model(x), atol=1e-2)) + def test_index_put_forms_coremltools_gets_wrong(self): + """ + coremltools 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 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]])) + self.register_buffer("flag", torch.tensor([2])) + self.register_buffer("rows_mask", torch.tensor([True, True, True])) + self.register_buffer("trio", torch.tensor([0, 2, 3])) + self.register_buffer("grid", torch.tensor([[1]])) + self.register_buffer("yes", torch.tensor([True])) + self.register_buffer("pair", torch.tensor([0, 1])) + + def forward(self, x, a, b, c): + y = x * 2 + y[:, self.cols] = a + y[:, self.grid] = a[:, :, None] + y[self.rows, self.cols] = b + # Same-shape indices without a slice stay in Core ML, also when + # a refused write reads the same index. + y[self.cols, self.cols] = a[:2, 0] + y[self.rows_mask, self.trio] = c + torch.ops.aten.index_put_(y, [None, self.col], a, True) + torch.ops.aten.index_put_(y, [None, self.last], c) + mask = y > 0.5 + mask[:, self.flag] = a > 0.5 + # Core ML would return this bool write as int32, and ~ would + # then read it wrong. + flags = x > 0.5 + flags[self.pair, self.pair] = a[:2, 0] > 0.5 + z = a * 2 + z[:, self.yes] = c[:, None] + return y + mask + z + ~flags + + class Scatter3d(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("last", torch.tensor(2)) + self.register_buffer("mid", torch.tensor(1)) + self.register_buffer("row", torch.tensor([1])) + self.register_buffer("col", torch.tensor([2])) + + 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)) + torch.ops.aten.index_put_(y, [None, self.row, self.col], v) + return y * 2 + + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + # Values that float16 holds exactly, so that the comparisons with 0.5 + # in Core ML agree with eager. + scatter_inputs = tuple( + torch.rand(shape).half().float() for shape in ((3, 4), (3, 1), (2, 2), (3,)) + ) + scatter3d_inputs = (torch.rand(2, 2, 3), torch.rand(2, 1)) + cases = [ + (Scatter(), scatter_inputs, 9, None), + # From iOS 18 the bool write next to a slice stays in Core ML. + (Scatter(), scatter_inputs, 8, ios18), + (Scatter3d(), scatter3d_inputs, 2, None), + ] + for model, inputs, 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(compile_specs=compile_specs)] + ) + 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.assertEqual(op_names.count("aten.index_put.default"), cpu_calls) + # An iOS 18 program needs macOS 15 to run. + if _TEST_RUNTIME and (compile_specs is None or IS_VALID_TEST_RUNTIME): + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute(inputs)[0] + self.assertTrue(torch.allclose(out, model(*inputs), atol=1e-2)) + # With a batch size known only at run time, the write that stays in + # Core ML above goes to the CPU too. + batch = torch.export.Dim("batch", min=2, max=8) + ep = torch.export.export( + Scatter3d(), + scatter3d_inputs, + dynamic_shapes=({0: batch}, {0: batch}), + 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.count("aten.index_put.default"), 3) + def test_reduction_of_0d_tensor_is_skipped(self): """ Core ML cannot reduce a 0-d tensor. The partitioner must leave such From f576ae4740bf73e2d9fe96907191604629d91475 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 13:07:26 -0800 Subject: [PATCH 4/5] Say why the Core ML partitioner refuses an op With lower_full_graph=True, an op that the partitioner refuses failed with a message saying that Core ML does not support the op, even when coremltools converts the op and the partitioner refuses only one form of it. For an op that coremltools converts, the error now opens with the reason instead, such as "aten.index_put.default on a bool tensor below minimum_deployment_target=ct.target.iOS18 is not supported by CoreML", and keeps its advice to set lower_full_graph=False or rewrite the model. The index_put test checks the reason for some of its forms, including a bool write that a later deployment target would not fix, which gives its own reason. --- .../coreml/partition/coreml_partitioner.py | 62 +++++++------------ .../coreml/test/test_coreml_partitioner.py | 35 ++++++++++- 2 files changed, 56 insertions(+), 41 deletions(-) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 849327df2c3..895f47e6443 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -367,7 +367,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 @@ -380,20 +382,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, @@ -401,47 +397,26 @@ 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 _reduces_0d_tensor(node): - self.log_once( - f"{node.target} of a 0-d tensor is not supported by CoreML. " - "Overriding op support." - ) - return True + return f"{name} of a 0-d tensor is not supported by CoreML." 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} on bool tensors is not supported by CoreML. " - "Overriding op support." - ) - return True + return f"{name} on bool tensors is not supported by CoreML." refusal = _index_put_refusal(node, self.minimum_deployment_target) if refusal is not None: - self.log_once( - f"{node.target} {refusal} is not supported by CoreML. " - "Overriding op support." - ) - return True + return f"{name} {refusal} 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 @@ -460,7 +435,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 @@ -478,13 +453,20 @@ 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 is_supported else None + 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: + unsupported = ( + f"CoreML does not support the op {node_target_name}" + if reason is None + else reason.rstrip(".") + ) raise NotImplementedError( - f"""CoreML does not support the op {node_target_name}, but you have set lower_full_graph=True in the CoreMLPartitioner. + f"""{unsupported}, but you have set lower_full_graph=True in the CoreMLPartitioner. Please set lower_full_graph=False in the CoreMLPartitioner to allow running unsupported ops outside of CoreML. Note that setting lower_full_graph=False may affect performance of CoreML and the available features. As an alternative to setting lower_full_graph=False, you can try rewriting your model to avoid using this op. diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 9ae748cfa79..6066299bca3 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -577,7 +577,8 @@ def test_index_put_forms_coremltools_gets_wrong(self): coremltools 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 Scatter(torch.nn.Module): @@ -666,6 +667,38 @@ def forward(self, x, v): program = Runtime.get().load_program(edge.to_executorch().buffer) out = program.load_method("forward").execute(inputs)[0] self.assertTrue(torch.allclose(out, model(*inputs), atol=1e-2)) + + class BoolWrite(torch.nn.Module): + def __init__(self, index, sliced=True): + super().__init__() + self.register_buffer("index", index) + self.sliced = sliced + + def forward(self, x, a): + indices = [None, self.index] if self.sliced else [self.index] + return torch.ops.aten.index_put(x > 0.5, indices, a > 0.5) + + for model, inputs, reason in ( + (Scatter3d(), scatter3d_inputs, "values that vary before the last index"), + ( + BoolWrite(torch.tensor([2])), + scatter_inputs[:2], + "below minimum_deployment_target=ct.target.iOS18", + ), + # A bool write that a later target would not fix gives its own reason. + (BoolWrite(torch.tensor([0, 2])), scatter_inputs[:2], "other than an int"), + ( + BoolWrite(torch.tensor([0, 2]), sliced=False), + (scatter_inputs[0], scatter_inputs[1][:2]), + "on a bool tensor without a slice", + ), + ): + with self.assertRaisesRegex(NotImplementedError, reason) as error: + executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, inputs, strict=True), + partitioner=[CoreMLPartitioner(lower_full_graph=True)], + ) + self.assertNotIn("does not support the op", str(error.exception)) # With a batch size known only at run time, the write that stays in # Core ML above goes to the CPU too. batch = torch.export.Dim("batch", min=2, max=8) From efb3cc0249be9716c814db07d78f4c0a994446f4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 3 Nov 2026 13:07:27 -0800 Subject: [PATCH 5/5] Keep the Core ML int4 recipe tests from failing at random test_int4_weight_only_per_group and test_int4_weight_only_per_channel used random weights and a random input, and they failed for some of them on macOS, on main too. Both compare the SQNR (signal to quantization noise ratio) of the output against the float model with a threshold. - Core ML runs these recipes in float16, so the per-group output differs from the eager quantized model by up to about 2e-3, more than that test's 1e-3 tolerance. The other int4 and int8 recipe tests use 1e-2. - The per-group SQNR came from a single input row, which gives only eight outputs. For some weights it fell to 12 dB. The eager quantized model gives the same SQNR as Core ML there, so this is the int4 rounding of those weights seen through eight outputs, not a Core ML error. - The per-channel threshold of 20 dB is above what int4 per-channel rounding gives for many random weights. Over seeds 0 to 199, the eager int4 model is below 20 dB against the float model for 43 of them, and at 17.7 dB at the lowest. Both tests now fix the seed. The per-group test feeds 16 rows and uses the 1e-2 tolerance; over seeds 0 to 99 every case passes, with the SQNR at 17 dB or more. The per-channel test uses 15 dB, the threshold of the per-group test, which all 200 seeds clear. --- backends/apple/coreml/test/test_coreml_recipes.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backends/apple/coreml/test/test_coreml_recipes.py b/backends/apple/coreml/test/test_coreml_recipes.py index 98d240d74b5..17b47bc3ae9 100644 --- a/backends/apple/coreml/test/test_coreml_recipes.py +++ b/backends/apple/coreml/test/test_coreml_recipes.py @@ -120,6 +120,7 @@ def test_fp_recipes_with_custom_parameters(self): def test_int4_weight_only_per_channel(self): """Test INT4 weight-only per-channel quantization""" + torch.manual_seed(0) model = TestHelperModules.TwoLinearModule().eval() example_inputs = [(torch.randn(9, 8),)] @@ -132,10 +133,13 @@ def test_int4_weight_only_per_channel(self): ) self.check_fully_delegated(session) self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-02) - self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + self._compare_eager_unquantized_model_outputs( + session, model, example_inputs, sqnr_threshold=15 + ) def test_int4_weight_only_per_group(self): """Test INT4 weight-only per-group quantization with different group sizes""" + torch.manual_seed(0) class CustomTwoLinearModel(nn.Module): def __init__(self): @@ -149,7 +153,7 @@ def forward(self, x): return x model = CustomTwoLinearModel().eval() - example_inputs = [(torch.randn(1, 32),)] + example_inputs = [(torch.randn(16, 32),)] # Test with different group sizes for group_size in [8, 16, 32]: with self.subTest(group_size=group_size): @@ -165,7 +169,7 @@ def forward(self, x): self.check_fully_delegated(session) self._compare_eager_quantized_model_outputs( - session, example_inputs, atol=1e-3 + session, example_inputs, atol=1e-2 ) self._compare_eager_unquantized_model_outputs( session, model, example_inputs, sqnr_threshold=15