fix(aggregation): Cast ConstantWeighting weights to matrix dtype

[SajalDevX]

Problem

ConstantWeighting.forward returns the stored weights tensor as-is, without considering the dtype and device of the input matrix. The aggregators that combine the rows of the Jacobian with these weights directly fail with a RuntimeError as soon as the weights and the Jacobian differ in dtype or device:

• Constant (weights @ matrix in WeightedAggregator._combine);
• ConFIG with a pref_vector (torch.linalg.pinv(units) @ weights);
• AlignedMTL with a pref_vector (B @ weights in AlignedMTLWeighting).

A typical case is a pref_vector created with torch.tensor([...]) (so float32, on CPU) used with a model in float64, or on CUDA. UPGrad and DualProj only work by accident, because their projector goes through numpy and casts its result back to the dtype and device of the Gramian. The other weightings that hold a user-provided vector already cast it at use (SDMGradWeighting._resolve_w_tilde and PCDWeighting._get_taus), so ConstantWeighting is the only one left behaving differently.

Repro (CPU):

import torch

from torchjd.aggregation import AlignedMTL, ConFIG, Constant

pref_vector = torch.tensor([1.0, 2.0])
matrix = torch.tensor([[-4.0, 1.0, 1.0], [6.0, 1.0, 1.0]], dtype=torch.float64)

for aggregator in [Constant(pref_vector), ConFIG(pref_vector=pref_vector), AlignedMTL(pref_vector=pref_vector)]:
    try:
        print(f"{aggregator}: {aggregator(matrix)}")
    except RuntimeError as e:
        print(f"{aggregator}: RuntimeError: {e}")

On main:

Constant([1., 2.]): RuntimeError: expected m1 and m2 to have the same dtype, but got: float != double
ConFIG([1., 2.]): RuntimeError: addmv input tensors must have the same dtype, but got Float, Double, and Float
AlignedMTL([1., 2.]): RuntimeError: addmv input tensors must have the same dtype, but got Float, Double, and Float

The same happens with a float64 pref_vector and a float32 matrix, and with a CPU pref_vector and a CUDA matrix (Expected all tensors to be on the same device).

Solution

In ConstantWeighting.forward, return self.weights.to(dtype=matrix.dtype, device=matrix.device) instead of self.weights. When the weights already match the matrix (the usual case), .to returns the same tensor, so nothing changes; the stored weights attribute is left untouched. With this change the repro prints:

Constant([1., 2.]): tensor([8., 3., 3.], dtype=torch.float64)
ConFIG([1., 2.]): tensor([0.4684, 2.1656, 2.1656], dtype=torch.float64)
AlignedMTL([1., 2.]): tensor([2.0079, 2.7512, 2.7512], dtype=torch.float64)

Tests

• Added test_weights_dtype_differs_from_matrix_dtype in tests/unit/aggregation/test_constant.py and test_pref_vector_dtype_differs_from_matrix_dtype in tests/unit/aggregation/test_config.py, both parametrized over (float32, float64) and (float64, float32) for the weights and the matrix. The weights are created explicitly on CPU, so on a CUDA run (PYTEST_TORCH_DEVICE=cuda:0) the tests also cover the device mismatch. All 4 parametrizations fail on main with the errors above and pass with the fix.
• pytest tests/unit/aggregation -W error: 2281 passed, 2 skipped (CPU, float32); the two modified modules also pass under PYTEST_TORCH_DTYPE=float64. pytest tests/unit/autojac/test_jac_to_grad.py -W error passes.
• ruff check, ruff format --check and ty check are unchanged.

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

[PierreQuinton]

Good catch, however I feel a bit unconfortable casting the weights everytime. Should we copy the weights and cast/return the copy instead?

[SajalDevX]

Thanks! .to(dtype=…, device=…) returns self.weights itself when the dtype and device already match, so in the usual case nothing is cast or copied and the behaviour is the same as before. It only allocates a converted tensor when they differ, which is the copy you describe.

If you'd rather not pay that conversion on every call for a mismatched matrix, I can cache the converted tensor per (dtype, device) in the weighting and return the cached copy. Happy to do that if you prefer it over the one-liner.

[PierreQuinton]

I'm not completely convinced about the cast to be honest. It feels like if this happens then the user most likely did something wrong, and hiding the mistake might be the best things. Note that as Weightings are modules, we could try to make the weights a parameter so that the to method of modules actually works. If we do this, then we will have an error if the jacobian and weights are not on the same device/dtype, and the way to fix it is to cast the weighting itself. We could also add an informative error message in the forward method of WeightedAggregator.

I can't really decide what is a good way to handle this, @ValerianRey @KhusPatel4450 any opinion?

[ValerianRey]

I think I agree with Pierre. This is more of a user mistake than a real issue IMO. Maybe we could improve error messages though.

[SajalDevX]

Fair point, a silent cast would hide the user's mistake.

One thing I noticed though: Constant(torch.tensor([1., 2.])).to(torch.float64) leaves the weights in float32, since weights is a plain attribute and not a buffer. So there is currently no clean way for the user to fix the mismatch.

I can rework this PR to register the weights as a buffer and raise an informative error in WeightedAggregator.forward instead. Or I can close it if you'd rather handle it yourselves. Let me know which you prefer.

[ValerianRey]

Fair point, a silent cast would hide the user's mistake.

One thing I noticed though: Constant(torch.tensor([1., 2.])).to(torch.float64) leaves the weights in float32, since weights is a plain attribute and not a buffer. So there is currently no clean way for the user to fix the mismatch.

I can rework this PR to register the weights as a buffer and raise an informative error in WeightedAggregator.forward instead. Or I can close it if you'd rather handle it yourselves. Let me know which you prefer.

I agree that weights should be registered as buffers (not parameters since they're not trainable). Maybe even as buffers not affecting the state dict:

self.register_buffer('foo', foo, persistent=False)

I'm not entirely sure about that persistent=False though.

Also, it would be nice to do that not only for weights, but for every tensor that we need to move with a call to .to() on the aggregator / weighting, that is not currently registered as a buffer or parameter.

Feel free to do that in a new PR.
Please make a different PR if you want to raise an error when the dtype of device mismatches.

thanks in advance! :)