fix autograd var_mean backward wiring

src/candle/_generated/variable_type.py

@@ -5190,17 +5190,31 @@ def var_correction_autograd(self_, dim=None, correction=None, keepdim=False, **_
 def var_mean_correction_autograd(self_, dim=None, correction=None, keepdim=False, **_kwargs):
     active_keyset = current_dispatch_keyset()
     raw_keyset = _strip_autograd_keys(active_keyset)
-    result = redispatch("var_mean", raw_keyset, self_, dim, correction, keepdim, **_kwargs)
-    if GradMode.enabled and (self_.requires_grad):
-        grad_fn = _F.VarMeanCorrectionBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
-        annotate_node_creation(grad_fn)
-        grad_fn._save(self_=self_)
-        grad_fn._dim = dim
-        grad_fn._correction = correction
-        grad_fn._keepdim = keepdim
-        result[0].grad_fn = grad_fn
+    if "unbiased" in _kwargs:
+        unbiased = _kwargs.pop("unbiased")
+    else:
+        unbiased = correction == 1 if correction is not None else True
+    result = redispatch("var_mean", raw_keyset, self_, dim=dim, unbiased=unbiased, keepdim=keepdim, **_kwargs)
+    if GradMode.enabled and (self_.requires_grad):
+        var_grad_fn = _F.VarCorrectionBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+        annotate_node_creation(var_grad_fn)
+        var_grad_fn._save(self_=self_)
+        var_grad_fn._dim = dim
+        var_grad_fn._correction = 1 if unbiased else 0
+        var_grad_fn._keepdim = keepdim
+        result[0].grad_fn = var_grad_fn
         result[0].requires_grad = True
-        result[1].grad_fn = grad_fn
+        if dim is None:
+            mean_grad_fn = _F.MeanBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+            annotate_node_creation(mean_grad_fn)
+            mean_grad_fn._save(self_=self_)
+        else:
+            mean_grad_fn = _F.MeanDimBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+            annotate_node_creation(mean_grad_fn)
+            mean_grad_fn._save(self_=self_)
+            mean_grad_fn._dim = dim
+            mean_grad_fn._keepdim = keepdim
+        result[1].grad_fn = mean_grad_fn
         result[1].requires_grad = True
     return result
 
@@ -13285,16 +13299,30 @@ def var_correction_autograd_post(result, self_, dim=None, correction=None, keepd
 
 
 def var_mean_correction_autograd_post(result, self_, dim=None, correction=None, keepdim=False, *, raw_keyset, active_keyset, **_kwargs):
-    if GradMode.enabled and (self_.requires_grad):
-        grad_fn = _F.VarMeanCorrectionBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
-        annotate_node_creation(grad_fn)
-        grad_fn._save(self_=self_)
-        grad_fn._dim = dim
-        grad_fn._correction = correction
-        grad_fn._keepdim = keepdim
-        result[0].grad_fn = grad_fn
+    if "unbiased" in _kwargs:
+        unbiased = _kwargs["unbiased"]
+    else:
+        unbiased = correction == 1 if correction is not None else True
+    if GradMode.enabled and (self_.requires_grad):
+        var_grad_fn = _F.VarCorrectionBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+        annotate_node_creation(var_grad_fn)
+        var_grad_fn._save(self_=self_)
+        var_grad_fn._dim = dim
+        var_grad_fn._correction = 1 if unbiased else 0
+        var_grad_fn._keepdim = keepdim
+        result[0].grad_fn = var_grad_fn
         result[0].requires_grad = True
-        result[1].grad_fn = grad_fn
+        if dim is None:
+            mean_grad_fn = _F.MeanBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+            annotate_node_creation(mean_grad_fn)
+            mean_grad_fn._save(self_=self_)
+        else:
+            mean_grad_fn = _F.MeanDimBackward0((self_,), raw_keyset=raw_keyset, active_keyset=active_keyset)
+            annotate_node_creation(mean_grad_fn)
+            mean_grad_fn._save(self_=self_)
+            mean_grad_fn._dim = dim
+            mean_grad_fn._keepdim = keepdim
+        result[1].grad_fn = mean_grad_fn
         result[1].requires_grad = True
     return result
 

tests/cpu/test_autograd_ops.py

@@ -369,6 +369,32 @@ def test_training(self):
         assert x.grad.shape == x.shape
 
 
+class TestVarMeanBackward:
+    def test_tuple_outputs_backward(self):
+        x = _tensor([1.0, 2.0, 3.0])
+        var, mean = torch.var_mean(x)
+        out = var + mean
+        out.backward()
+        assert x.grad is not None
+        assert x.grad.shape == x.shape
+        _check_grad(x, [-2.0 / 3.0, 1.0 / 3.0, 4.0 / 3.0])
+
+    def test_tuple_outputs_backward_dim(self):
+        x = _tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        var, mean = torch.var_mean(x, dim=1)
+        out = var.sum() + mean.sum()
+        out.backward()
+        assert x.grad is not None
+        _check_grad(x, [[-5.0 / 6.0, 1.0 / 6.0, 7.0 / 6.0], [-5.0 / 6.0, 1.0 / 6.0, 7.0 / 6.0]])
+
+    def test_outputs_do_not_share_grad_fn(self):
+        x = _tensor([1.0, 2.0, 3.0])
+        var, mean = torch.var_mean(x)
+        assert var.grad_fn is not None
+        assert mean.grad_fn is not None
+        assert var.grad_fn is not mean.grad_fn
+
+
 class TestEmbeddingBackward:
     def test_basic(self):
         w = _tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])

tests/cpu/test_declarative_autograd.py

@@ -89,6 +89,15 @@ def test_differentiability_info_properties(self):
         assert len(mul_info.differentiable_inputs) == 1
         assert not mul_info.is_inplace
         assert not mul_info.is_multi_output
+    @_skip_no_yaml
+    def test_derivative_info_marks_multi_output_when_schema_returns_tuple(self):
+        from tools.autograd.load_derivatives import load_derivatives
+        from pathlib import Path
+        yaml_path = Path(__file__).resolve().parents[2] / "tools" / "autograd" / "derivatives.yaml"
+        infos = load_derivatives(yaml_path)
+        info = next(i for i in infos if i.name == "var_mean")
+        assert info.func_name.startswith("var_mean")
+        assert info.is_multi_output
 
 
 class TestCodegenOutput: