diff --git a/CHANGELOG.md b/CHANGELOG.md index 05f30bf29..fe25e1956 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,10 @@ changelog does not include internal changes that do not affect the user. - Fixed `Engine.compute_gramian` silently ignoring the contribution of a module when one of its forward passes had no associated backward pass (e.g. because its output was detached). It now raises a `ValueError` instead. +- Fixed `jac_to_grad` failing with a `RuntimeError` when the tensors do not all have the same dtype + (e.g. a model with some parameters in `float32` and others in `float64`). The aggregated gradient + was computed in the promoted dtype and could not be assigned to the `.grad` field of the tensors + of lower precision. Each tensor now gets a `.grad` of its own dtype, like with `torch.autograd`. ## [0.17.1] - 2026-09-23 diff --git a/src/torchjd/autojac/_jac_to_grad.py b/src/torchjd/autojac/_jac_to_grad.py index e8c67f104..8ceb5641b 100644 --- a/src/torchjd/autojac/_jac_to_grad.py +++ b/src/torchjd/autojac/_jac_to_grad.py @@ -205,7 +205,7 @@ def _gramian_based( gradients = list[Tensor]() while jacobians: jacobian = jacobians.popleft() # get jacobian + dereference it to free memory asap - gradients.append(torch.tensordot(weights, jacobian, dims=1)) + gradients.append(torch.tensordot(weights.to(dtype=jacobian.dtype), jacobian, dims=1)) return gradients, weights @@ -229,7 +229,9 @@ def _disunite_gradient( tensors: list[TensorWithJac], ) -> list[Tensor]: gradient_vectors = gradient_vector.split([t.numel() for t in tensors]) - gradients = [g.reshape(t.shape) for g, t in zip(gradient_vectors, tensors, strict=True)] + gradients = [ + g.reshape(t.shape).to(dtype=t.dtype) for g, t in zip(gradient_vectors, tensors, strict=True) + ] return gradients diff --git a/tests/unit/autojac/test_jac_to_grad.py b/tests/unit/autojac/test_jac_to_grad.py index 5abe4fd24..89bde2064 100644 --- a/tests/unit/autojac/test_jac_to_grad.py +++ b/tests/unit/autojac/test_jac_to_grad.py @@ -1,5 +1,6 @@ from typing import Any +import torch from pytest import mark, raises from torch import Tensor from torch.testing import assert_close @@ -266,6 +267,29 @@ def test_noncontiguous_jac() -> None: assert_grad_close(t, g) +@mark.parametrize("optimize", [True, False]) +def test_mixed_dtypes(optimize: bool) -> None: + """ + Tests that jac_to_grad works when the tensors do not all have the same dtype. Each tensor should + get a .grad of its own dtype, as torch's backward would do. + """ + + aggregator = base_agg() + t1 = tensor_(1.0, dtype=torch.float32, requires_grad=True) + t2 = tensor_([2.0, 3.0], dtype=torch.float64, requires_grad=True) + jac = tensor_([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]], dtype=torch.float64) + t1.__setattr__("jac", jac[:, 0].to(dtype=torch.float32)) + t2.__setattr__("jac", jac[:, 1:]) + expected_grad = aggregator(jac) + + jac_to_grad([t1, t2], aggregator, optimize_gramian_computation=optimize) + + assert t1.grad is not None and t1.grad.dtype == torch.float32 + assert t2.grad is not None and t2.grad.dtype == torch.float64 + assert_grad_close(t1, expected_grad[0].to(dtype=torch.float32)) + assert_grad_close(t2, expected_grad[1:]) + + @mark.parametrize("aggregator", [base_agg(), ConFIG()]) def test_aggregator_hook_is_run(aggregator: Aggregator) -> None: """