From 30c520df15c131c7ca8680d3937c0c9a04b61a3b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:59:03 -0800 Subject: [PATCH 1/4] 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 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. --- .../coreml/partition/coreml_partitioner.py | 50 +++++++++++++++++++ .../coreml/test/test_coreml_partitioner.py | 36 +++++++++++++ 2 files changed, 86 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 4453818f956..15252f69f34 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -162,6 +162,49 @@ def _is_view_op(op: torch._ops.OpOverload) -> bool: return (alias_info is not None) and (not alias_info.is_write) +_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, +) + + +# 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, @@ -249,6 +292,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 0129eb1d33f..dc9d2c01cf3 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -521,6 +521,42 @@ 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 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) + 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) + floats = x.amax(-1, keepdim=True) + torch.maximum(x > 0.5, x) + return x * 2 + flags + both + floats + + 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.assertEqual(op_names.count("aten.amax.default"), 1) + self.assertEqual(op_names.count("aten.maximum.default"), 1) + 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_tied_embedding_is_quantized_with_its_linear(self): """A weight that is both an embedding table and an output projection lowers. From 619cd8e999d2d4fb90d4350fa6a66161073556d1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:59:03 -0800 Subject: [PATCH 2/4] 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. 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, and on macOS 15 compares the output with eager. --- .../coreml/partition/coreml_partitioner.py | 96 ++++++++++++++++ .../coreml/test/test_coreml_partitioner.py | 106 ++++++++++++++++++ 2 files changed, 202 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 15252f69f34..99514a1eda4 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 @@ -205,18 +206,98 @@ 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. 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: index for dim, index in enumerate(indices) if index is not None} + 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 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 with a minimum_deployment_target below iOS 18" + 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 @@ -299,6 +380,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 @@ -393,6 +482,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: @@ -449,6 +543,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, ) @@ -488,6 +583,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 dc9d2c01cf3..7dcb0330542 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 @@ -557,6 +560,109 @@ 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])) + + 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 + 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 + z = a * 2 + z[:, self.yes] = c[:, None] + return y + mask + z + + 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 + + ios18 = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + scatter_inputs = ( + torch.rand(3, 4), + torch.rand(3, 1), + torch.rand(2, 2), + torch.rand(3), + ) + scatter3d_inputs = (torch.rand(2, 2, 3), torch.rand(2, 1)) + cases = [ + (Scatter(), scatter_inputs, 8, None), + # From iOS 18 the bool write stays in Core ML. + (Scatter(), scatter_inputs, 7, ios18), + (Scatter3d(), scatter3d_inputs, 1, 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 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"), 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 eb2deca23289234fe7a445ed9e658ad3675c462e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:59:04 -0800 Subject: [PATCH 3/4] 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. The error now gives the reason, such as "aten.index_put.default on a bool tensor with a minimum_deployment_target below iOS 18 is not supported by CoreML. Set lower_full_graph=False in the CoreMLPartitioner to run it outside of CoreML." 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 | 54 +++++++------------ .../coreml/test/test_coreml_partitioner.py | 27 +++++++++- 2 files changed, 46 insertions(+), 35 deletions(-) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 99514a1eda4..a7e59c914eb 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -319,7 +319,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 @@ -332,20 +334,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, @@ -353,40 +349,23 @@ 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} 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 @@ -405,7 +384,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 @@ -423,10 +402,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 7dcb0330542..64f1308786d 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -565,7 +565,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): @@ -644,6 +645,30 @@ 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): + super().__init__() + self.register_buffer("index", index) + + def forward(self, x, a): + return torch.ops.aten.index_put(x > 0.5, [None, self.index], 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], + "minimum_deployment_target below iOS 18", + ), + # 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"), + ): + with self.assertRaisesRegex(NotImplementedError, reason): + executorch.exir.to_edge_transform_and_lower( + torch.export.export(model, inputs, strict=True), + partitioner=[CoreMLPartitioner(lower_full_graph=True)], + ) # 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 98b4f25d014ca6695289952e59bcdede4c924d4a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 00:12:32 -0800 Subject: [PATCH 4/4] Keep the Core ML int4 per-group recipe test from failing at random test_int4_weight_only_per_group used random weights and a random input, and it failed for some of them on macOS, on main too. - Core ML runs this recipe in float16, so its output differs from the eager quantized model by up to about 2e-3, more than the test's 1e-3 tolerance. The other int4 and int8 recipe tests use 1e-2. - The SQNR against the float model 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 test now feeds 16 rows, uses the 1e-2 tolerance, and fixes the seed. Over seeds 0 to 99 every case passes, with the SQNR at 17 dB or more. --- backends/apple/coreml/test/test_coreml_recipes.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/backends/apple/coreml/test/test_coreml_recipes.py b/backends/apple/coreml/test/test_coreml_recipes.py index 98d240d74b5..47b21207c2a 100644 --- a/backends/apple/coreml/test/test_coreml_recipes.py +++ b/backends/apple/coreml/test/test_coreml_recipes.py @@ -136,6 +136,7 @@ def test_int4_weight_only_per_channel(self): 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 +150,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 +166,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