diff --git a/src/candle/_generated/variable_type.py b/src/candle/_generated/variable_type.py index 020918bb..ba246683 100644 --- a/src/candle/_generated/variable_type.py +++ b/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 diff --git a/tests/cpu/test_autograd_ops.py b/tests/cpu/test_autograd_ops.py index 304700b6..c91bfdf8 100644 --- a/tests/cpu/test_autograd_ops.py +++ b/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]]) diff --git a/tests/cpu/test_declarative_autograd.py b/tests/cpu/test_declarative_autograd.py index a7ea0c42..23d43b00 100644 --- a/tests/cpu/test_declarative_autograd.py +++ b/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: