refactor(autojac): Optimize jac_to_grad

[ValerianRey]

Avoid concatenation of the jacobians when the aggregator is gramian-based. Also use a deque to free each jacobian as soon as it is used.

• Add gramian-based jac_to_grad
• Update changelog

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/torchjd/_linalg/_gramian.py	100.00% <100.00%> (ø)
src/torchjd/autojac/_jac_to_grad.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[claude]

Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

[ValerianRey]

See #524 (comment) for (a bit outdated) performance comparison.

[ValerianRey]

With this PR, we rely a lot more on compute_gramian being competitively fast, because we call it hundreds of times instead of just once per iteration. So I optimized it for the case contracted_dims=-1 (the most usual case). This has a non-negligible impact on performance (for instance, reducing time of jac_to_grad from 38 ms to 32 ms for InstanceNormMobileNetV2()-bs2

[PierreQuinton]

You can accept all commits together if you want, I prefer to use TypeGuards, it is prettier and clearer.

[ValerianRey]

You can accept all commits together if you want, I prefer to use TypeGuards, it is prettier and clearer.

I think you missed one of my commit messages:

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.

=> I asked claude if it was ok to use typeguard here, but it said that it should really only be used if the function only checks that the type of the aggregator is GramianWeightedAggergator. So i'd rather use cast just to be safe.

[PierreQuinton]

I disagree with Claude here. The logic of typeguards is "if the check pass, then the type is A", not "the check passes if and only if the type is A". So here is is correct. Otherwise every function returning a typeguard would be exactly:

def check_B(a: A) -> TypeGuard[B]:
    return isinstance(a, B)

But of course if this was the only intended use of TypeGuards, then they would have provided exactly this function instead.

[ValerianRey]

I disagree with Claude here. The logic of typeguards is "if the check pass, then the type is A", not "the check passes if and only if the type is A". So here is is correct. Otherwise every function returning a typeguard would be exactly:

def check_B(a: A) -> TypeGuard[B]:
    return isinstance(a, B)

But of course if this was the only intended use of TypeGuards, then they would have provided exactly this function instead.

You're right. Please read this, it's short and very interesting: https://docs.python.org/3/library/typing.html#typing.TypeGuard

Apparently there's also a TypeIs, which has a few differences with TypeGuard. Here, we wouldn't be allowed to use TypeIs, because returning False with a TypeIs would tell the type checker that the type is not what is in the brackets of the TypeIs. But with TypeGuard, it doesn't say anything to the type checker when returning False.

So I think that:

• This can safely be a TypeGuard
• TypeGuard can be changed to TypeIs in is_matrix, is_psd_tensor and is_psd_matrix

Should we fire @claude?

[claude]

Claude finished @ValerianRey's task in 52s —— View job

---

Analysis of TypeGuard vs TypeIs

You're absolutely right, and I apologize for my earlier incorrect guidance! Let me clarify the semantics:

Todo List

• Read the current implementation in _jac_to_grad.py
• Read other files with type guard functions
• Analyze TypeGuard vs TypeIs semantics
• Provide recommendations

---

TypeGuard vs TypeIs Semantics

TypeGuard:

• When the function returns True, the type checker narrows the type to the guarded type
• When the function returns False, the type checker learns nothing about the type
• This is appropriate when the check is sufficient but not necessary for the type

TypeIs (Python 3.13+):

• When the function returns True, the type checker narrows the type to the guarded type
• When the function returns False, the type checker narrows the type to exclude the guarded type
• This is appropriate when the check is both necessary and sufficient for the type (i.e., if and only if)

Analysis of Your Functions

1. _can_skip_jacobian_combination (src/torchjd/autojac/_jac_to_grad.py:85)

Current implementation:

def _can_skip_jacobian_combination(aggregator: Aggregator) -> bool:
    return isinstance(aggregator, GramianWeightedAggregator) and not _has_forward_hook(aggregator)

Recommendation: Use TypeGuard[GramianWeightedAggregator] ✅

This is the correct choice because:

• When True: The aggregator is indeed a GramianWeightedAggregator (and has no hooks, but that's an additional runtime constraint)
• When False: The aggregator might still be a GramianWeightedAggregator (if it has hooks), so we can't exclude that type

Should NOT use TypeIs because returning False doesn't mean the aggregator is NOT a GramianWeightedAggregator - it just means we can't use the optimization.

2. is_matrix, is_psd_tensor, is_psd_matrix (src/torchjd/_linalg/_matrix.py:26-38)

Current implementation:

def is_matrix(t: Tensor) -> TypeGuard[Matrix]:
    return t.ndim == 2

def is_psd_tensor(t: Tensor) -> TypeGuard[PSDTensor]:
    half_dim = t.ndim // 2
    return t.ndim % 2 == 0 and t.shape[:half_dim] == t.shape[: half_dim - 1 : -1]

def is_psd_matrix(t: Tensor) -> TypeGuard[PSDMatrix]:
    return t.ndim == 2 and t.shape[0] == t.shape[1]

Recommendation: Use TypeIs for is_matrix ✅

For is_matrix:

• When True: The tensor is a 2D tensor (Matrix)
• When False: The tensor is definitively NOT a 2D tensor (not a Matrix)
• This is an if-and-only-if condition, so TypeIs is semantically correct

Recommendation: Keep TypeGuard for is_psd_tensor and is_psd_matrix ✅

For is_psd_tensor and is_psd_matrix:

• The checks verify shape requirements but NOT the actual PSD property (as noted in comments)
• When False: A tensor could still conceptually be a PSDTensor/PSDMatrix even if it fails the shape check (e.g., if we had a different representation)
• More importantly, these functions don't check the actual positive semi-definiteness, which is part of the semantic contract of these types
• So these are sufficient but not strictly necessary conditions

However, there's an argument for TypeIs here too: if we consider the type definitions to be purely about shape (since PSD checking is expensive and done separately in tests), then TypeIs would be appropriate. But TypeGuard is safer and more conservative.

Summary

1. _can_skip_jacobian_combination: Should use TypeGuard[GramianWeightedAggregator] ✅ (PierreQuinton is correct)
2. is_matrix: Could use TypeIs[Matrix] (it's an exact if-and-only-if condition)
3. is_psd_tensor and is_psd_matrix: Should keep TypeGuard (they don't check the full semantic meaning of PSD)

And no, you shouldn't fire me - I was being overly conservative before! 😄 The key insight is that TypeGuard is appropriate when you're checking sufficient conditions that don't rule out the type when they fail.

---

[ValerianRey]

Look at him trying to keep his 22$/month job. How pathetic.

In all this noise he just said, he said something really important: TypeIs is only available since python 3.13. Let's use that in 3 years. (Technically, we could use the typing-extensions backport to have TypeIs in python 3.10, but it's an extra dependency for almost no value).

[ValerianRey]

Up-to-date time comparison between branch main (a55afca) and this (fc9bbcf) on CPU. Note that there is a bit of randomness in those results. Overall, I think it's better when jac_to_grad takes

Model	Batch Size	Time before (jac_to_grad)	Time after (jac_to_grad)	Difference (jac_to_grad)	Time before (autojac_forward_backward)	Time after (autojac_forward_backward)	Difference (autojac_forward_backward)
AlexNet	4	301 ms	251 ms	-51 ms (-17%)	812 ms	667 ms	-144 ms (-18%)
Cifar10Model	64	47 ms	51 ms	+4 ms (+8%)	1715 ms	1618 ms	-96 ms (-6%)
GroupNormMobileNetV3Small	8	27 ms	29 ms	+1 ms (+5%)	711 ms	722 ms	+11 ms (+1%)
InstanceNormMobileNetV2	2	17 ms	22 ms	+5 ms (+27%)	353 ms	359 ms	+6 ms (+2%)
InstanceNormResNet18	4	62 ms	48 ms	-14 ms (-23%)	812 ms	779 ms	-32 ms (-4%)
SqueezeNet	4	8 ms	11 ms	+3 ms (+41%)	452 ms	434 ms	-18 ms (-4%)
WithTransformerLarge	4	220 ms	155 ms	-65 ms (-29%)	656 ms	544 ms	-112 ms (-17%)

I think that the most significant results are those of AlexNet and WithTransformerLarge, where jac_to_grad takes a large proportion of the total time of the autojac_forward_backward. So basically, when it matters, this optimization is great.

We have been quite lucky with the values obtained for the forward and backward for this PR, but we really could expect a -50 ms (out of 812 ms, so that's -6%) on AlexNet and a -65 ms (out of 656 ms, so that's -10%) on WithTransformerLarge.

Gaining 10% iteration speed (including forward, backward and jac_to_grad) just from this little change seems like a huge improvement IMO.

Let's hope the results are similar on cuda (I'm scared that it could launch way too many cuda kernels and thus slow down things). I had a problem when profiling on cuda, so I need to investigate that.

[ValerianRey]

I just made a homemade memory visualizer because torch doesn't provide any for CPU and the one for GPU is broken.
Memory usage over time on main (blue) vs this PR (orange) for WithTransformerLarge, on CPU:

=> It seems that we do half peak memory usage. Really nice.

On the blue line, the jump from ~0.7 to ~1.41 corresponds to the memory allocation of torch.concat when concatenating the jacobians. Then, the jump from ~1.41 to ~1.59 corresponds to the memory allocation for the multiplication of the combined jacobian matrix with the extracted weights (aggregator.combine). I think this jump should be avoided even in the _jacobian_based approach of this PR because we now properly free the jacobians as early as we can (thanks to aaf2544)

[ValerianRey]

All plots:

AlexNet (24% peak memory reduction):

Cifar10Model (no peak memory reduction: we flatten a peak that is a local maximum):

GroupNormMobileNetV3Small (no peak memory reduction: we flatten a peak that is a local maximum):

InstanceNormMobileNetV2 (no peak memory reduction: we flatten a peak that is a local maximum):

InstanceNormResNet18 (no peak memory reduction: we flatten a peak that is a local maximum):

SqueezeNet (no peak memory reduction: we flatten a peak that is a local maximum):

WithTransformerLarge: (54% peak memory reduction - more than half):

[PierreQuinton]

The TypeIs provides an intersection and a negation operator thus making the set of types a Boolean algebra. We do not have intersection type, but we can have a variable typed checked to the intersection type by mypy with this. Still seems tedious, it would be nice to have intersection type.

The plots are extremely cool, we should put one in the release notes eventually.

[ValerianRey]

CUDA memory timelines:

[PierreQuinton]

Following our discussion where we considered making the Generics Weighting[PSDMatrix] and Weighting[Matrix] into abstract classes which we can then use to differentiate the method used for jac_to_grad, we could in principle avoid testing instances and rather have methods in each specifying how we handle a list of Jacobians. This is good in terms of not have cases, but is bad in terms of having a less clear cut algebraic meaning to the aggregation package.

[ValerianRey]

CUDA Traces Comparison

Model	Batch Size	Time before (jac_to_grad)	Time after (jac_to_grad)	Difference (jac_to_grad)	Time before (autojac_forward_backward)	Time after (autojac_forward_backward)	Difference (autojac_forward_backward)
AlexNet	4	98 ms	98 ms	0 ms (+1%)	132 ms	129 ms	-2 ms (-2%)
Cifar10Model	64	121 ms	163 ms	+42 ms (+35%)	184 ms	210 ms	+26 ms (+14%)
GroupNormMobileNetV3Small	8	27 ms	39 ms	+12 ms (+44%)	87 ms	99 ms	+12 ms (+14%)
InstanceNormMobileNetV2	2	6 ms	19 ms	+13 ms (+201%)	93 ms	105 ms	+12 ms (+13%)
InstanceNormResNet18	4	41 ms	41 ms	0 ms (0%)	84 ms	85 ms	+2 ms (+2%)
SqueezeNet	4	18 ms	22 ms	+4 ms (+21%)	40 ms	43 ms	+3 ms (+8%)
WithTransformerLarge	4	76 ms	78 ms	+2 ms (+3%)	175 ms	184 ms	+9 ms (+5%)

=> This seems too slow on cuda. I think we can't merge this until we have a way to parallelize the cuda kernel launches or whetever it is that takes this extra time.