Const fold scatter_along_axis and scaled_dot_product_attention correctly

[LeSingh1]

Two ops have a value_inference that ignores an input which changes the result. When the operands happen to be const, the op is folded to the wrong value and then removed from the model by const_elimination, so the model is silently wrong rather than failing to convert.

scatter_along_axis ignored mode

value_inference always called np.put_along_axis, i.e. it always computed mode="update". Every other mode folded to the overwrite result — including mode="add", which is the op's own default.

import numpy as np
from coremltools.converters.mil.mil import Builder as mb

data = np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
indices = np.array([[0, 0, 2]], dtype=np.int32)      # index 0 written twice
updates = np.array([[10.0, 100.0, 30.0]], dtype=np.float32)

@mb.program(input_specs=[mb.TensorSpec(shape=(1,))])
def prog(x):
    return mb.scatter_along_axis(data=data, indices=indices, updates=updates,
                                 axis=1, mode="add")

op = prog.functions["main"].find_ops(op_type="scatter_along_axis")[0]
print(op.outputs[0].val)

mode	folded (before)	Core ML runtime
update	[[100, 2, 30]]	[[100, 2, 30]]
add	[[100, 2, 30]]	[[111, 2, 33]]
sub	[[100, 2, 30]]	[[-109, 2, -27]]
mul	[[100, 2, 30]]	[[1000, 2, 90]]
div	[[100, 2, 30]]	[[0.001, 2, 0.1]]
max	[[100, 2, 30]]	[[100, 2, 30]]
min	[[100, 2, 30]]	[[1, 2, 3]]

The runtime column was measured by running the same program with data/updates as model inputs on this machine.

The fix computes each mode with the matching numpy ufunc applied through ufunc.at, so repeated indices accumulate the way the runtime does (add above gives 1 + 10 + 100). Negative indices keep wrapping, as before. div on an integer tensor is not folded, because numpy cannot write a float quotient back into an integer array. iOS17.scatter_along_axis inherits this value_inference, so it is fixed too.

scaled_dot_product_attention dropped a non-const attn_mask

float_mask = None
if self.attn_mask is not None and self.attn_mask.val is not None:
    ...

When attn_mask is present but computed at runtime, float_mask stayed None and the op folded to unmasked attention. This op has no @precondition, so value_inference runs whenever it is asked to, and with const query/key/value the entire attention collapsed into a const of the wrong answer.

The fix returns None in that case, so the op stays in the graph and the mask is applied at runtime.

Testing

New tests in iOS14/test_scatter_gather.py and iOS18/test_transformers.py. For each op there is a builder-eval test (the folded value) and a run_compare_builder test (the runtime value), so the two are pinned to each other. In the attention backend test the attention output feeds a mul, because a bad fold on a var that is directly a block output is not propagated by const_elimination and so would not be observable end to end.

All 8 new tests fail on main and pass here. Existing TestScatterAlongAxis (iOS14/16/17) and TestScaledDotProductAttention tests still pass; they all used mode="update" or a const mask, which are the two paths whose results are unchanged.