fix(autojac): Support tensors of different dtypes in jac_to_grad

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
 

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
 
 

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:
     """