Add _can_skip_jacobian_combination helper function

Tbh I don't like it very much (because it's an extra function + some cast is required) but it's the only way to easily test that the correct aggregators use the optimized _gramian_based method. I also tried using return type hint of TypeGuard[GramianWeightedAggergator] instead of bool for _can_skip_jacobian_combination, but it's not really correct since we also check that the aggregator has no forward hook, so that TypeGuard would be really weird. So in the end we have to use this cast.

Author: ValerianRey

src/torchjd/autojac/_jac_to_grad.py

@@ -75,15 +75,17 @@ def jac_to_grad(
     if not retain_jac:
         _free_jacs(tensors_)
 
-    if isinstance(aggregator, GramianWeightedAggregator) and not _has_forward_hook(aggregator):
-        # When it's possible, avoid the concatenation of the jacobians that can be very costly in
-        # memory.
-        gradients = _gramian_based(aggregator, jacobians, tensors_)
+    if _can_skip_jacobian_combination(aggregator):
+        gradients = _gramian_based(cast(GramianWeightedAggregator, aggregator), jacobians, tensors_)
     else:
         gradients = _jacobian_based(aggregator, jacobians, tensors_)
     accumulate_grads(tensors_, gradients)
 
 
+def _can_skip_jacobian_combination(aggregator: Aggregator) -> bool:
+    return isinstance(aggregator, GramianWeightedAggregator) and not _has_forward_hook(aggregator)
+
+
 def _has_forward_hook(module: nn.Module) -> bool:
     """Return whether the module has any forward hook registered."""
     return len(module._forward_hooks) > 0 or len(module._forward_pre_hooks) > 0