Reduce runtime checking

[ValerianRey]

In order to make torchjd as efficient as possible, I think we need to limit ourselves to only checking the inputs provided by the user, at the interface (e.g. are the losses provided to mtl_backward correct?), or things that are unknown before runtime (e.g. are the values of the Jacobian finite?). A few examples of checks that we either have (or used to have - in which case the box is ticked), that reduce performance while never failing, are:

• Key checks when composing transforms => This always passed in publicly available transforms (i.e. those created by backward and mtl_backward). Fixed in Isolate key checks in transforms #279.
• _AggregateMatrices._disunite checks that the united_gradient_vector has the expected length => This always passes. Fix proposed in Remove _disunite check #377.
• TensorDict checks => These checks also always pass when the transforms are correctly implemented (which a test should check anyway). Fix proposed in Move tensor dict checks out of constructor #380.
• The aggregators check that their inputs are finite (i.e. no nan, inf or -inf values allowed). Technically, we can't really know ahead of time if the Jacobian will be finite, this check is not useless. However, it can be very time or memory-consuming on huge matrices (and the Jacobian can be huge). It takes 1.62 sec on average on my cpu when testing a matrix of shape [10, 10^8]. On cuda, it's even worse: it requires the instantiation of another [10, 10^8] matrix (to contain the boolean values), which is too much memory for my gpu to handle. So I think we need to address this. Fix proposed in Remove _check_is_finite #381.

A few extra cases that I think are not worth changing (so I tick them already):

• The checks of get_leaf_tensors (e.g. if any([tensor.grad_fn is None for tensor in tensors]):) could technically be moved from this function to backward and mtl_backward). However, it does not bring any performance and it makes the code less readable IMO.
• The aggregators check that their input is a matrix in Aggregator._check_is_matrix. In autojac, it will always be the case, but aggregators are also publicly available so users can use them on their own tensors (they're thus another interface to torchjd). In this case, the tensors might not be matrices, and it's important to raise a clear error. Also, performance-wise, this results in 1 extra trivial check per backward pass.
• _check_expects_grad could technically be moved from the Accumulate transform to the backward and mtl_backward function (to be at the interface). In case the tensors with respect to which we differentiate are not provided, our get_leaf_tensors function should already provide tensors that expect a grad, so the check could be skipped. Otherwise, we could simply check that all the tensors provided by the user do expect a grad. This would only lead to a (negligible) performance improvement in the first case (as otherwise, we would still run the same number of checks, just at a different place), but since it's the default, it's very likely to happen. Also, it goes in the direction of having checks only at the interface. However, I think we decided that the requires_grad, is_leaf and retains_grad fields might change between the creation of the Transform and the time it is applied, so the current version is safer in that respect. For that reason (and the fact the performance improvement would be negligible), I think it's fine to leave Accumulate as it is.

[ValerianRey]

Here are the results of my investigations (c.f. 7c95297) about what can happen if we remove the call to _check_is_finite in aggregators. I tested 6 different ways of making the matrix not fully finite:

• 10x100 matrix filled with 1.0, but matrix[0, 0] replaced by one of {nan, inf, -inf}
• 10x100 matrix filled with one of {nan, inf, -inf}

I then tested all 16 of our aggregators against all 6 of these cases, and the aggregators could be grouped in 8 groups:

• The result is full of nan in all 6 cases: UPGrad, DualProj and MGDA
• Just 1 non-finite (nan, inf or -inf) value is trimmed, otherwise the value is kept in the result: TrimmedMean
• Just the result of the column with a non-finite value becomes the value: Sum, Mean, Random and Constant.
• Almost the same, except that for some reason when matrix[0, 0] is -inf, result[0] becomes nan and the rest of the result becomes +inf: PCGrad (might be up to the random seed).
• Just the result of the affected column becomes nan: Krum, GradDrop
• Always raises a ValueError from cvxpy: NashMTL
• Always raises a LinalgError from torch: CAGrad, IMTLG and AlignedMTL
• Full nan result when full non-finite input matrix, result full of the value (inf, -inf or nan) when only matrix[0, 0] has this non-finite value: ConFIG

To summarize, we either have:

• Result that clearly makes sense: TrimmedMean, Sum, Mean, Random, Constant
• Result that seems to make sense (as it's not really possible to be more precise than nan in those cases): UPGrad, DualProj, MGDA, Krum, GradDrop
• Result for which it's hard to tell: ConFIG, PCGrad
• Exception raised: NashMTL, CAGrad, IMTLG, AlignedMTL

When the result at least seems to make sense, we're good.
When it's hard to tell, I don't think it's the role of the aggregator to be correct in those cases. If the Jacobian contains a non-finite value, there's already a bigger issue. And we can't really fix the aggregators for those cases.
When the aggregator raises an exception, we could arguably try to catch it and re-raise it with a clearer error message. However, we cannot guarantee that a ValueError from cvxpy will always be due to a non-finite input, and same for the LinalgError. And again, there's obviously something that went wrong prior to that, which the user should realize themselves. For those aggregators, the only change will be that the error message will be less clear without the _check_is_finite call, but it will make them faster and more memory-efficient.

So my opinion on this is that we should:

• Remove _check_is_finite
• Add a changelog entry
• Archive the 16x6 = 96 new tests I just added, as they're not very important but we might still want them one day.