refactor(aggregation): Register state tensors as buffers

[KhusPatel4450]

Summary

• GradVacWeighting._phi_t, CRMOGMWeighting._lambda, MoDoWeighting._lambda, and SDMGradWeighting._w were plain Python attributes initialised to None.
• Changed each to self.register_buffer(...) so that nn.Module.to(device="cuda") (or .to(dtype=...)) automatically moves the warm-started state tensor when the user migrates their model to GPU.
• _NashMTLWeighting is unchanged — its state is NumPy/cvxpy arrays, not PyTorch tensors.
• _state_key tuples are left as plain attributes (not tensors).
• Existing forward() and reset() assignments (self._w = w, self._w = None) continue to work: PyTorch's __setattr__ routes assignments to registered buffer names through _buffers automatically.

Test plan

• uv run pytest tests/unit -q — 3167 passed, 0 failed
• uv run ty check on all five files — 0 new errors (7 pre-existing NashMTL/cvxpy errors unchanged)

Claude Code was used to assist with this task

[ValerianRey]

LGTM. @PierreQuinton This is related to what you asked in #737.

[PierreQuinton]

Very cool, I guess the only problem is that they are typed as Any? But this doesn't matter too much. LGTM

[ValerianRey]

Very cool, I guess the only problem is that they are typed as Any? But this doesn't matter too much. LGTM

No, it's typed as None | Tensor, same as before. I'm not sure how ty infers that though, but it works.

[ValerianRey]

Arguably I think we could go 1 step further and also register as buffers the tensors that can only be changed through the init or through a setter.

This way, calling .to(device) will actually always work, as far as I know.
The only question is: do we want to have such tensors in the state dict of the aggregator / scalarizer? I don't think so, because these are not necessarily stateful aggregators. So in my opinion we should use persistent=False when registering those kinds of buffers. @KhusPatel4450 wanna take that one too?