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

## Summary

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

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

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

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

The second shape is commonly produced by either:

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

or:

```python
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:

```text
||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:

```python
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

```text
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:

```text
%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:

```text
%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:

```text
(||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](https://github.com/user-attachments/files/30292351/repro_coreai_optimize_axismove.py)