From 737e2e33a4ef7dca8fde71e06beaf4cc7845d907 Mon Sep 17 00:00:00 2001 From: LeSingh1 Date: Sun, 9 Aug 2026 13:23:28 -0700 Subject: [PATCH] Const fold scatter_along_axis and scaled_dot_product_attention correctly Both ops had a value_inference that ignored an input which changes the result, so a program whose operands happen to be const folds to the wrong value and the op disappears from the model. scatter_along_axis: value_inference always called np.put_along_axis, i.e. it always computed mode="update". Every other mode was folded to the overwrite result, including the default mode="add". Compute the mode's reduction with the matching numpy ufunc applied through ufunc.at so that repeated indices accumulate, which is what the runtime does. scaled_dot_product_attention: value_inference only applied attn_mask when the mask had a value, and otherwise silently produced unmasked attention. This op has no @precondition, so with const query/key/value and a mask computed at runtime the whole attention collapsed to a const of the unmasked result. Return None instead, so the op stays in the graph. --- .../mil/mil/ops/defs/iOS15/scatter_gather.py | 36 +++++++- .../mil/mil/ops/defs/iOS18/transformers.py | 6 +- .../ops/tests/iOS14/test_scatter_gather.py | 88 ++++++++++++++++++ .../mil/ops/tests/iOS18/test_transformers.py | 89 +++++++++++++++++++ 4 files changed, 215 insertions(+), 4 deletions(-) diff --git a/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py b/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py index 4ac6a86d4..f9a53910a 100644 --- a/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py @@ -399,15 +399,45 @@ def default_inputs(self): mode="add", ) + # Reduction used by each ``mode`` other than ``update``. They are applied with + # ``ufunc.at``, so repeated indices accumulate, which matches the runtime. + _MODE_TO_NP_UFUNC = { + "add": np.add, + "sub": np.subtract, + "mul": np.multiply, + "div": np.divide, + "max": np.maximum, + "min": np.minimum, + } + @precondition(allow=VALUE) def value_inference(self): data = np.copy(self.data.val) indices = self.indices.val updates = self.updates.val axis = self.axis.val - np_output = data - np.put_along_axis(np_output, indices, updates, axis=axis) - return np_output + if axis < 0: + axis += self.data.rank + mode = self.mode.val + + if mode == "update": + np.put_along_axis(data, indices, updates, axis=axis) + return data + + ufunc = self._MODE_TO_NP_UFUNC.get(mode) + if ufunc is None: + # Unknown mode: do not const fold rather than fold to a wrong value. + return None + if ufunc is np.divide and not np.issubdtype(data.dtype, np.floating): + # np.divide cannot write its result back into an integer array. + return None + + # Index ``data`` along ``axis`` with ``indices``, and along every other axis with + # that axis' own coordinates, i.e. the indexing np.put_along_axis performs. + index = list(np.indices(indices.shape, sparse=True)) + index[axis] = indices + ufunc.at(data, tuple(index), updates) + return data def type_inference(self): if self.axis.val < -self.data.rank or self.axis.val >= self.data.rank: diff --git a/coremltools/converters/mil/mil/ops/defs/iOS18/transformers.py b/coremltools/converters/mil/mil/ops/defs/iOS18/transformers.py index cdc779ee3..dc9bf2042 100644 --- a/coremltools/converters/mil/mil/ops/defs/iOS18/transformers.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS18/transformers.py @@ -144,8 +144,12 @@ def value_inference(self): return None float_mask = None - if self.attn_mask is not None and self.attn_mask.val is not None: + if self.attn_mask is not None: mask = self.attn_mask.val + if mask is None: + # The mask changes the result, so without its value there is nothing + # to const fold to. + return None if mask.dtype == bool: float_mask = np.zeros(mask.shape) float_mask[np.where(np.logical_not(mask))] = -np.inf diff --git a/coremltools/converters/mil/mil/ops/tests/iOS14/test_scatter_gather.py b/coremltools/converters/mil/mil/ops/tests/iOS14/test_scatter_gather.py index fb0373c95..1bb2bf065 100644 --- a/coremltools/converters/mil/mil/ops/tests/iOS14/test_scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/tests/iOS14/test_scatter_gather.py @@ -187,6 +187,94 @@ def prog(x): rtol=1e-05, ) + # data / indices / updates shared by the ``mode`` tests below. Index 0 of the last + # axis is written twice so that the accumulating modes are distinguishable from a + # plain overwrite. + _MODE_DATA = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + _MODE_INDICES = np.array([[0, 0, 2]], dtype=np.int32) + _MODE_UPDATES = np.array([[10.0, 100.0, 30.0]], dtype=np.float32) + + @staticmethod + def _reference_scatter_along_axis(data, indices, updates, axis, mode): + """Straightforward transcription of the op definition, for axis == 1.""" + assert axis == 1 + output = np.copy(data) + for i in range(indices.shape[0]): + for j in range(indices.shape[1]): + k = indices[i, j] + if mode == "update": + output[i, k] = updates[i, j] + elif mode == "add": + output[i, k] += updates[i, j] + elif mode == "sub": + output[i, k] -= updates[i, j] + elif mode == "mul": + output[i, k] *= updates[i, j] + elif mode == "div": + output[i, k] /= updates[i, j] + elif mode == "max": + output[i, k] = max(output[i, k], updates[i, j]) + else: + assert mode == "min" + output[i, k] = min(output[i, k], updates[i, j]) + return output + + @pytest.mark.parametrize( + "backend, mode", + itertools.product(backends, ["update", "add", "sub", "mul", "div", "max", "min"]), + ) + def test_builder_eval_mode(self, backend, mode): + """The const folded value must honor ``mode``, not silently overwrite.""" + data, indices, updates = self._MODE_DATA, self._MODE_INDICES, self._MODE_UPDATES + + @mb.program( + input_specs=[mb.TensorSpec(shape=(1,), dtype=types.fp32)], + opset_version=backend.opset_version, + ) + def prog(x): + return mb.scatter_along_axis( + data=data, indices=indices, updates=updates, axis=1, mode=mode + ) + + scatter_op = prog.functions["main"].find_ops(op_type="scatter_along_axis")[0] + np.testing.assert_allclose( + self._reference_scatter_along_axis(data, indices, updates, 1, mode), + scatter_op.outputs[0].val, + atol=1e-04, + rtol=1e-05, + ) + + @pytest.mark.parametrize( + "compute_unit, backend, mode", + itertools.product( + compute_units, backends, ["update", "add", "sub", "mul", "div", "max", "min"] + ), + ) + def test_builder_to_backend_mode(self, compute_unit, backend, mode): + """The runtime result for each ``mode``, which the const folding must reproduce.""" + data, indices, updates = self._MODE_DATA, self._MODE_INDICES, self._MODE_UPDATES + + input_placeholders = { + "data": mb.placeholder(shape=data.shape), + "updates": mb.placeholder(shape=updates.shape), + } + input_values = {"data": data, "updates": updates} + + def build(data, updates): + return mb.scatter_along_axis( + data=data, indices=indices, updates=updates, axis=1, mode=mode + ) + + run_compare_builder( + build, + input_placeholders, + input_values, + (1, 3, types.fp32), + self._reference_scatter_along_axis(data, indices, updates, 1, mode), + compute_unit=compute_unit, + backend=backend, + ) + @staticmethod def _test_builder_to_backend_programmatic( compute_unit, backend, rank_axis, force_non_negative_indices diff --git a/coremltools/converters/mil/mil/ops/tests/iOS18/test_transformers.py b/coremltools/converters/mil/mil/ops/tests/iOS18/test_transformers.py index 475ea127c..be3a9c9c5 100644 --- a/coremltools/converters/mil/mil/ops/tests/iOS18/test_transformers.py +++ b/coremltools/converters/mil/mil/ops/tests/iOS18/test_transformers.py @@ -100,6 +100,95 @@ def test_builder_eval_stress(self, batches, float_dtype, mask_dtype): rtol=1e-6 if float_dtype == np.float32 else 1e-3, ) + @staticmethod + def _const_query_key_value_with_mask_placeholder(mask_dtype): + """ + Const query / key / value, plus a mask that attends to the first key position + only. Every key position carries a very different value, so ignoring the mask + gives a result far away from honoring it. + """ + S, L, E, EV = 5, 7, 16, 32 + np.random.seed(0) + query = np.random.rand(2, L, E).astype(np.float32) + key = np.random.rand(2, S, E).astype(np.float32) + value = np.zeros((2, S, EV), dtype=np.float32) + value[:, 0, :] = 100.0 + + mask = np.zeros((1, 1, S), dtype=mask_dtype) + if mask_dtype is bool: + mask[:, :, 0] = True + else: + mask[:, :, 1:] = -np.inf + return query, key, value, mask + + @pytest.mark.parametrize( + "mask_dtype", + (bool, np.float32), + ) + def test_builder_eval_non_const_mask(self, mask_dtype): + """ + ``attn_mask`` changes the result, so when only ``attn_mask`` is non const there is + nothing to const fold to. + """ + query, key, value, mask = self._const_query_key_value_with_mask_placeholder(mask_dtype) + + @mb.program( + input_specs=[ + mb.TensorSpec( + shape=mask.shape, dtype=types.numpy_type_to_builtin_type(mask_dtype) + ) + ], + opset_version=ct.target.iOS18, + ) + def prog(mask): + return mb.scaled_dot_product_attention( + query=query, + key=key, + value=value, + attn_mask=mask, + ) + + attention = prog.functions["main"].find_ops(op_type="scaled_dot_product_attention")[0] + assert attention.outputs[0].val is None + + @pytest.mark.parametrize( + "compute_unit, backend, mask_dtype", + itertools.product(compute_units, backends, (bool, np.float32)), + ) + def test_builder_to_backend_non_const_mask(self, compute_unit, backend, mask_dtype): + """ + End to end counterpart of ``test_builder_eval_non_const_mask``. The attention + feeds another op so that a const folded attention would be propagated into the + model by ``const_elimination``. + """ + query, key, value, mask = self._const_query_key_value_with_mask_placeholder(mask_dtype) + + def build(mask): + attention = mb.scaled_dot_product_attention( + query=query, + key=key, + value=value, + attn_mask=mask, + ) + return mb.mul(x=attention, y=np.float32(2.0)) + + attention_torch = 2.0 * self._torch_scaled_dot_product_attention(query, key, value, mask) + run_compare_builder( + build, + { + "mask": mb.placeholder( + shape=mask.shape, dtype=types.numpy_type_to_builtin_type(mask_dtype) + ) + }, + {"mask": mask}, + expected_output_types=[attention_torch.shape + (types.fp32,)], + expected_outputs=[attention_torch], + compute_unit=compute_unit, + backend=backend, + atol=1e-6 if backend.precision == "fp32" else 1e-3, + rtol=1e-6 if backend.precision == "fp32" else 1e-3, + ) + @pytest.mark.parametrize( "compute_unit, backend, batches, float_dtype, mask_dtype", itertools.product(