Const fold scatter_along_axis and scaled_dot_product_attention correctly

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:

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

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

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(