Fix torch var/std converter ignoring correction > 1 under tracing

[adityasingh2400]

Bug

When a model is converted via torch.jit.trace, torch.var/torch.std called with the correction keyword lower to the same aten::var/aten::std signature as the older unbiased keyword. Both keywords end up in the same positional slot (aten::var(self, dim, <slot>, keepdim)), and the slot's constant is just an int.

The converter always read that slot as the boolean unbiased. So any correction >= 2 was truthy and silently treated as unbiased=True, i.e. correction=1. correction=0 and correction=1 happened to come out right only because they coincide with unbiased=False/True.

Repro (torch 2.12, coremltools from main):

import torch, coremltools as ct, numpy as np
torch.manual_seed(0)
x = torch.randn(3, 5)

class M(torch.nn.Module):
    def forward(self, x):
        return torch.var(x, dim=-1, correction=2, keepdim=True)

ts = torch.jit.trace(M().eval(), x)
m = ct.convert(ts, inputs=[ct.TensorType(name="x", shape=x.shape, dtype=np.float32)],
               compute_units=ct.ComputeUnit.CPU_ONLY, minimum_deployment_target=ct.target.iOS17)

print("torch ", M()(x).flatten()[:2].tolist())
print("coreml", np.asarray(list(m.predict({"x": x.numpy()}).values())[0]).flatten()[:2].tolist())

torch  [2.762..., 1.057...]   # divides by N - 2 = 3
coreml [2.072..., 0.793...]   # divides by N - 1 = 4  (wrong, == correction=1)

correction=3, 4, ... are all collapsed to the correction=1 result the same way. std has the same problem since it reuses this logic.

Fix

Read the argument as correction rather than unbiased. correction is a strict generalization: unbiased=False is correction=0 and unbiased=True is correction=1, so existing unbiased-traced graphs keep their current behavior, and correction >= 2 now divides by N - correction as PyTorch does. The value is routed through the existing _var(correction=...) path. The unbiased= kwarg (export) is still accepted and mapped onto correction.

Test

test_var_std_with_correction only parametrized correction over [0, 1], which is exactly the range the old code got right by accident, so the bug was invisible. Extended it to include correction=2. With this fix the full matrix (var/std x correction[0,1,2] x dim[[0,2],[1],[2]] x keepdim) passes; on main the correction=2 cases fail with a numerical mismatch.

[adityasingh2400]

Still reproduces on main at 375ce90, so this has not been overtaken by anything since I opened it.

In the TORCHSCRIPT branch of var, the second or third positional argument is read into unbiased (lines 7824 and 7827 today). Under torch.jit.trace that slot can carry a correction value instead, and _var only tests if unbiased: at line 7789, so correction=2 is truthy and silently yields the correction=1 result rather than raising.

Happy to rebase or split this if that makes it easier to take.

[TobyRoseman]

Yes, please rebase this change on top of latest main.

Also, please make all of the comments more concise and clear. All comments should be in your own words.

Once those two things are done, I will kick off a CI run.

[LeSingh1]

Not trying to step on this — I hit the same bug independently and found your PR while checking for duplicates, so I dropped mine. One thing from my digging that might be worth folding in, take it or leave it.

test_var_std_with_correction in test_torch_ops.py parametrizes correction over [0, 1]. Those are exactly the two values where the correction and the boolean unbiased flag agree, which is why the bug never tripped a test. Widening that list to [0, 1, 2] gives you a regression test for free — no new test function needed. On current main that turns into 12 failures across the var/std TorchScript cases, and they go green with the fix.

Numbers I measured on main @ be45795f, in case they're useful for the description:

var(dim=1, correction=2)   torch [2.5, 10.0]    coreml [1.666, 6.664]
var(dim=1, correction=3)   torch [5.0, 20.0]    coreml [1.666, 6.664]
std(correction=2)          torch 2.5            coreml 2.3145

The Core ML result is the same for correction=2 and 3, which is the tell — both collapse to the unbiased n-1 path.

[adityasingh2400]

@TobyRoseman both done.

Rebased onto main, so the branch is now zero commits behind. I also cut the two comments down. They now read:

# torch.jit.trace does not distinguish by name `var` and `var.dim`,
# instead by nargs 2 or 4. It also merges the `unbiased` and
# `correction` overloads, so this argument may be either one.
# Read it as `correction`, which covers both.

# A bool `unbiased` is just correction 1 (True) or 0 (False).
# When absent, torch defaults to the unbiased estimator.

@LeSingh1 thanks, and sorry you spent time on a duplicate. Your parametrize point was the right call and it is in the branch: test_var_std_with_correction now runs correction over [0, 1, 2] instead of [0, 1], so the regression is covered without adding a test function. [0, 1] were exactly the two values where correction and the boolean unbiased agree, which is why nothing caught this.

[TobyRoseman]

This change looks good now.

CI: https://gitlab.com/coremltools1/coremltools/-/pipelines/2838001215