AIProgram.optimize() removes broadcasting-significant axis moves and silently miscompiles N×N distance expressions

[dkomoroske]

Summary

AIProgram.optimize() silently changes the semantics of broadcast arithmetic used by the standard expanded squared-distance expression:

D[i,j] = ||x_i||² - 2·x_i·y_j + ||y_j||²

For equal-length inputs, the two norm tensors have these shapes:

x norms: (1,N,1)
y norms: (1,1,N)

The second shape is commonly produced by either:

torch.sum(y ** 2, dim=-1).unsqueeze(-2)

or:

torch.sum(y ** 2, dim=-1, keepdim=True).transpose(-1, -2)

Although moving (1,N,1) to (1,1,N) does not reorder its element values, the shape change is essential to broadcasting. After optimize(), that axis move is absent. The final addition consumes the unrearranged (1,N,1) reduction and broadcasts ||y_i||² where ||y_j||² belongs.

The resulting expression is:

||x_i||² - 2·x_i·y_j + ||y_i||²

Because the operands have equal lengths, the invalid rewrite still produces an output with the expected (1,N,N) shape. No diagnostic is emitted.

Environment

• macOS 27.0 on Apple Silicon
  • originally reproduced on build 26A5378j
  • reproduced again on 2026-07-22 on build 26A5388g
• coreai-torch 0.4.1
• coreai-core 1.0.0b2
• torch 2.11.0
• Python 3.12.13

These were the newest PyPI releases, including prereleases, when retested on 2026-07-22.

Steps to reproduce

Run the attached repro_coreai_optimize_axismove.py. It creates a fresh exported program for every optimized and unoptimized arm, saves each Core AI asset, executes it, and compares its output with eager PyTorch.

The minimal failing form uses z as a graph input, so a matmul is not required to trigger the problem:

s1 = torch.sum(x ** 2, dim=-1).unsqueeze(-1)  # (1,N,1)
s2 = torch.sum(y ** 2, dim=-1).unsqueeze(-2)  # (1,1,N)
out = (s1 - 2 * z + s2).clamp(min=0.0)

Expected result

AIProgram.optimize() preserves the program’s output. All arms should match eager PyTorch to normal float32 noise.

Actual result

Chain           optimize=False: max|d| = 1.907e-06  OK
Chain           optimize=True : max|d| = 1.022e+01  MISCOMPILED
ChainKeepdim    optimize=False: max|d| = 1.907e-06  OK
ChainKeepdim    optimize=True : max|d| = 1.022e+01  MISCOMPILED
ChainReordered  optimize=False: max|d| = 3.815e-06  OK
ChainReordered  optimize=True : max|d| = 3.815e-06  OK

The optimized result matches an independent NumPy implementation of the wrong-axis expression to 1.907e-06.

Isolation and controls

The following controls were tested independently:

Control	Result
Eager PyTorch versus NumPy	Exact
exported.module() versus eager	Exact
Conversion without optimize()	Correct
SpecializationOptions.cpu_only()	Same miscompile
Expression without the final clamp	Miscompiled
Real x @ y.transpose(-1,-2), distinct inputs of equal length	Miscompiled
Real self-distance expression using the same tensor	Miscompiled
Reordered (s1 + s2) - 2*z expression	Correct
s1 + s2 alone	Correct
Unequal input lengths, producing a 17×23 result	Correct

For the distinct-input case, the optimized graph retains separate reductions for x and y but removes the axis move from the y reduction. This is therefore not a valid common-subexpression merge.

The identical CPU-only result also indicates that this is not caused by GPU placement or reduced-precision execution.

IR evidence

Before optimization, the y reduction is moved to (1,1,32), and the final addition consumes that shape:

%y_norm = coreai.reduce_sum ... -> tensor<1x32x1xf32>
%y_norm_moved = coreai.expand_dims ... -> tensor<1x1x32xf32>
%tmp = ...broadcasting_sub ...
    : (tensor<1x32x1xf32>, tensor<1x32x32xf32>)
    -> tensor<1x32x32xf32>
%out = ...broadcasting_add %tmp, %y_norm_moved
    : (tensor<1x32x32xf32>, tensor<1x1x32xf32>)
    -> tensor<1x32x32xf32>

After optimization, both reductions remain, but the axis move is gone:

%x_norm = coreai.reduce_sum ... %arg0 ... -> tensor<1x32x1xf32>
%y_norm = coreai.reduce_sum ... %arg1 ... -> tensor<1x32x1xf32>
%tmp = ...broadcasting_sub ...
    : (tensor<1x32x1xf32>, tensor<1x32x32xf32>)
    -> tensor<1x32x32xf32>
%out = ...broadcasting_add %tmp, %y_norm
    : (tensor<1x32x32xf32>, tensor<1x32x1xf32>)
    -> tensor<1x32x32xf32>

In the keepdim=True variant, the pre-optimization graph contains one transpose operation and the optimized graph contains none.

Scope and impact

The verified failure requires shape compatibility that allows the unrearranged operand to broadcast successfully. Equal-length and square/self-distance matrices are affected; the tested unequal-length 17×23 case remained correct.

Models lowered through this expanded squared-distance formulation may include point-cloud registration and segmentation networks, geometric attention, and k-nearest-neighbor graph construction. Fused distance operators and differently ordered formulations have not been shown to be affected.

The failure is particularly hazardous because the output has the expected shape and contains plausible values. In a larger GeoTransformer conversion, this appeared as approximately 17 dB PSNR versus eager PyTorch and scrambled nearest-neighbor relationships. Disabling optimize() restored approximately 78–85 dB parity.

Workarounds

Both of these workarounds were verified:

1. Do not call AIProgram.optimize(). Conversion, save_asset, specialization, loading, and inference work correctly without it.

2. Reorder the expression so the two norm operands are added first:

   (||x_i||² + ||y_j||²) - 2·x_i·y_j
   

   The reordered graph remains correct after optimization.

Related report

Possibly related to #9, which also reports a silent semantics-changing simplification reached through prog.optimize(), but the affected operation and rewrite pattern are different.

This issue was also submitted through Feedback Assistant as FB23695952, with the original sysdiagnose attached there.

repro_coreai_optimize_axismove.py