reduce_sum over a short last axis after mul is wrong past 65536 elements #2906

reduce_sum(mul(x, y), axes=[-1]) comes back wrong once the tensor gets past 65536 elements and the last axis is short . mul on its own and reduce_sum on its own are both exact .

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

rng = np.random.default_rng(0)
for shape in ((96, 168, 4), (98, 168, 4)):
    x = rng.random(shape, dtype=np.float32)
    y = rng.random(shape, dtype=np.float32)

    @mb.program(
        input_specs=[mb.TensorSpec(shape=shape), mb.TensorSpec(shape=shape)],
        opset_version=ct.target.macOS15,
    )
    def prog(x, y):
        return mb.reduce_sum(x=mb.mul(x=x, y=y), axes=[-1])

    m = ct.convert(
        prog,
        convert_to="mlprogram",
        pass_pipeline=ct.PassPipeline.EMPTY,
        minimum_deployment_target=ct.target.macOS15,
        compute_units=ct.ComputeUnit.CPU_ONLY,
    )
    out = list(m.predict({"x": x, "y": y}).values())[0]
    print(shape, np.abs(out - (x * y).sum(-1)).max())

(96, 168, 4) 2.3841858e-07
(98, 168, 4) 2.7332225

96 x 168 x 4 is 64512 elements and 98 x 168 x 4 is 65856 . the last axis matters : with 4 to 16 it goes wrong , with 2 or 3 it does not , and a last axis of 768 at 98304 elements is fine . add in place of mul goes wrong the same way , and the result is the same for every compute_units setting .

it reaches plain torch code too . (x * y).sum(-1) on a (1, 3, 128, 168, 4) tensor , converted with fp32 precision and run on the cpu , is off by 3.2 . with fp16 on all compute units it is fine .

macOS 26.5.2 (25F84) , apple m5 , coremltools main at a207b1a .

this is what broke the separable upsample lowering at image sizes . the pull request i am opening for the antialias resize works around it with a matmul .