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

[SajalDevX]

Problem

jac_to_grad fails with a RuntimeError when the tensors do not all have the same dtype, for example a model with some parameters in float32 and others in float64. torch.autograd.backward and autojac.backward both handle this fine (each .jac field is in the dtype of its tensor), but jac_to_grad cannot then turn these Jacobians into gradients:

• With optimize_gramian_computation=False, _unite_jacobians concatenates the Jacobians into a single matrix, which promotes them to the highest dtype. The aggregated gradient vector is therefore in that dtype, and assigning a slice of it to the .grad field of a tensor of lower precision is rejected by torch.
• With optimize_gramian_computation=True, the weights are extracted from the sum of the Gramians (again in the promoted dtype), and torch.tensordot(weights, jacobian, dims=1) fails for every Jacobian of a different dtype.

Repro (CPU):

import torch
from torch.nn import Linear, Sequential

from torchjd.aggregation import UPGrad
from torchjd.autojac import backward, jac_to_grad

torch.manual_seed(0)
model = Sequential(Linear(3, 2), Linear(2, 2, dtype=torch.float64))
x = torch.randn(4, 3)
losses = model[1](model[0](x).double()).pow(2).mean(dim=0)  # one loss per output

backward(losses)
print([p.jac.dtype for p in model.parameters()])
jac_to_grad(model.parameters(), UPGrad())
print([p.grad.dtype for p in model.parameters()])

On main:

[torch.float32, torch.float32, torch.float64, torch.float64]
RuntimeError: attempting to assign a gradient with dtype 'double' to a tensor with grad_dtype 'Float'. The gradient must match the tensor's grad_dtype (defaults to the tensor's dtype). [...]

and with jac_to_grad(model.parameters(), UPGrad(), optimize_gramian_computation=True):

RuntimeError: both inputs should have same dtype

Solution

Cast each gradient to the dtype of its tensor, which is what torch.autograd does:

• in _disunite_gradient, each reshaped slice of the gradient vector is cast with .to(dtype=t.dtype);
• in _gramian_based, the weights are cast to the dtype of each Jacobian before the tensordot.

When all tensors share the same dtype (the usual case), .to(dtype=...) is a no-op and returns the same tensor, so nothing changes. With this change the repro prints [torch.float32, torch.float32, torch.float64, torch.float64] for the .grad fields in both modes, and the gradients obtained with Mean() match torch.autograd.grad(losses.mean(), model.parameters()).

Tests

• Added test_mixed_dtypes in tests/unit/autojac/test_jac_to_grad.py, parametrized over optimize_gramian_computation. It gives a float32 tensor and a float64 tensor .jac fields of their own dtype, calls jac_to_grad and checks that each .grad has the dtype of its tensor and the expected value. Both parametrizations fail on main (with the two errors above) and pass with the fix. The dtypes are set explicitly, so the test checks the same thing under float32 and under PYTEST_TORCH_DTYPE=float64.
• pytest tests/unit/autojac -W error: 396 passed, for both float32 and PYTEST_TORCH_DTYPE=float64 (CPU only). pytest tests/unit/aggregation tests/unit/test_deprecations.py -W error: 1883 passed, 2 skipped.
• ruff check, ruff format --check and ty check pass on the changed files.

Added an entry under [Unreleased] / Fixed in the changelog.

[ValerianRey]

Very good catch. This was a major bug IMO.