feat: Use .jac fields #510

ValerianRey · 2026-01-08T00:25:39Z

This is just a draft, showing more or less how this change will impact the interface. What do you think so far @PierreQuinton ?

TODO:

Update tests of backward and mtl_backward
Add unit tests for jac_to_grad and AccumulateJac
Find solution to the dependency problem
Use utils package?
Finish solution to the jac not defined attribute
Check that the jacobians all have the same first dim in jac_to_grad
Add a changelog entry explaining that this is a breaking change and how to make the transition

PierreQuinton · 2026-01-08T08:00:37Z

src/torchjd/utils/_jac_to_grad.py

+    Aggregates the Jacobians stored in the ``.jac`` fields of ``params`` and accumulates the result
+    into their ``.grad`` fields.


Are we set on "Accumulate" vs "store" in .grad?

I think that accumulate is the way to make it more torchy, and store is a new direction related to our (my?) dislike of sums.

I think the main reason to accumulate is to be more torch-like. Also, if a user really wants to replace the .grad, they can do optimizer.zero_grad() and then call jac_to_grad(...).

The other way around (we decide that jac_to_grad replaces the .grad field and a user instead wants to accumulate) is not doable.

Pretty convinced, but let's keep this discussion for few days in case we think of other things.

src/torchjd/autojac/_backward.py

src/torchjd/utils/_accumulation.py

src/torchjd/utils/_jac_to_grad.py

src/torchjd/utils/_accumulation.py

ValerianRey · 2026-01-08T10:31:26Z

We have something a bit tricky: both jac_to_grad and the AccumulateGrad transform depend on the code from _accumulation.py.

To reuse some code in two different packages (utils and autojac), without making one depend on the other, we would need another protected package just for this code. So IMO we could:

Have a the _accumulation.py code somewhere else in torchjd (either in a standalone protected python file or in a protected accumulation module).
Duplicate some accumulation code.
Live with the fact that autojac depends on utils and utils depends on aggregation.

ValerianRey · 2026-01-08T10:51:05Z

Lastly, it's very annoying to have this Pylance error all over the place:

Cannot access attribute "jac" for class "Tensor"
  Attribute "jac" is unknown

I'm not sure how we could reliably stop showing this error, or even use a custom tensor that would properly fix this.

PierreQuinton · 2026-01-08T10:57:47Z

Lastly, it's very annoying to have this Pylance error all over the place:
Cannot access attribute "jac" for class "Tensor"
  Attribute "jac" is unknown
I'm not sure how we could reliably stop showing this error, or even use a custom tensor that would properly fix this.

I think this is just because all this is bad design. We made it more torch compatible, and this is the price. Pylance is right, there is not field jac in Tensors.

ValerianRey · 2026-01-08T11:02:38Z

I think this is just because all this is bad design. We made it more torch compatible, and this is the price. Pylance is right, there is not field jac in Tensors.

Yeah but there are still ways to make it much cleaner. I'm working on something.

PierreQuinton · 2026-01-08T11:07:09Z

Yeah but there are still ways to make it much cleaner. I'm working on something.

I think we should keep in mind that we would like the users to have a good typing experience. It would be nasty to make them use .jac fields if they get typing errors.

ValerianRey · 2026-01-08T11:11:45Z

Yeah but there are still ways to make it much cleaner. I'm working on something.

I think we should keep in mind that we would like the users to have a good typing experience. It would be nasty to make them use .jac fields if they get typing errors.

Well, they should rarely need to use these .jac fields themselves. If they want access to the jacobians, they should use autojac.jac.

PierreQuinton · 2026-01-08T11:14:18Z

right, that's very good, and if they do, they should get some amount of warning. You can delete the discussion if you want (annoying that we cannot have a thread).

src/torchjd/utils/_accumulation.py

- Remove test_aggregate.py - Update test_accumulate.py and test_interactions.py to test on AccumulateGrad instead of Accumulate - Fix tests in test_backward.py and test_mtl_backward.py to match the new interface: check the jac field instead of the .grad field. - Use _asserts.py for helper functions common to backward.py and mtl_backward.py

codecov · 2026-01-09T19:18:29Z

Codecov Report

❌ Patch coverage is 93.18182% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/torchjd/utils/_accumulation.py	87.50%	3 Missing ⚠️
src/torchjd/utils/_jac_to_grad.py	92.85%	3 Missing ⚠️

Files with missing lines	Coverage Δ
src/torchjd/autojac/_backward.py	`100.00% <100.00%> (ø)`
src/torchjd/autojac/_mtl_backward.py	`100.00% <100.00%> (ø)`
src/torchjd/autojac/_transform/__init__.py	`100.00% <100.00%> (ø)`
src/torchjd/autojac/_transform/_accumulate.py	`100.00% <100.00%> (ø)`
src/torchjd/utils/__init__.py	`100.00% <100.00%> (ø)`
src/torchjd/utils/_tensor_with_jac.py	`100.00% <100.00%> (ø)`
src/torchjd/utils/_accumulation.py	`87.50% <87.50%> (ø)`
src/torchjd/utils/_jac_to_grad.py	`92.85% <92.85%> (ø)`

🚀 New features to boost your workflow:

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

ValerianRey added the feat New feature or request label Jan 8, 2026

ValerianRey self-assigned this Jan 8, 2026

ValerianRey requested a review from PierreQuinton January 8, 2026 00:25

feat: Use .jac fields

99b4260

ValerianRey force-pushed the revamp-interface branch from 48c27b2 to 99b4260 Compare January 8, 2026 00:27

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/autojac/_backward.py Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Outdated Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_jac_to_grad.py Outdated Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_jac_to_grad.py Outdated Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Outdated Show resolved Hide resolved

Delete jac field instead of setting to None

7dda66f

ValerianRey force-pushed the revamp-interface branch from 712667c to 7dda66f Compare January 8, 2026 11:08

ValerianRey added 5 commits January 8, 2026 13:52

[WIP] Fix jac undefined errors

e08e45f

Merge branch 'main' into revamp-interface

0bcf1c0

Add check that the number of rows of the jacobians is consistant

57e5c6d

Add check of jac shape before assigning to .jac

6ea5983

Add changelog entry

b7b3a75

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Show resolved Hide resolved

PierreQuinton reviewed Jan 8, 2026

View reviewed changes

src/torchjd/utils/_accumulation.py Outdated Show resolved Hide resolved

Move check of jacobian shape outside of if/else

28cf701

ValerianRey added 3 commits January 8, 2026 16:53

Improve docstring of AccumulateJac and AccumulateGrad

9a2a0ec

Merge branch 'main' into revamp-interface

9e855e2

Simplify a test

7e95c37

		Aggregates the Jacobians stored in the ``.jac`` fields of ``params`` and accumulates the result
		into their ``.grad`` fields.

feat: Use .jac fields #510

Are you sure you want to change the base?

feat: Use .jac fields #510

Uh oh!

Conversation

ValerianRey commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

PierreQuinton Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

ValerianRey Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

PierreQuinton Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ValerianRey commented Jan 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ValerianRey commented Jan 8, 2026

Uh oh!

PierreQuinton commented Jan 8, 2026

Uh oh!

ValerianRey commented Jan 8, 2026

Uh oh!

PierreQuinton commented Jan 8, 2026

Uh oh!

ValerianRey commented Jan 8, 2026

Uh oh!

PierreQuinton commented Jan 8, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Jan 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ValerianRey commented Jan 8, 2026 •

edited

Loading

ValerianRey commented Jan 8, 2026 •

edited

Loading

codecov bot commented Jan 9, 2026 •

edited

Loading