refactor(autojac): Remove reshaping to/from matrix in Jac

[ValerianRey]

• Fixes No need to reshape to and from matrix in Jac #419

• Stop reshaping Jacobians in Jac
• Factorize _get_vjp between Grad and Jac
• Add changelog entry

Since the first commit also made _get_vjp be the same between Grad and Jac, I also factorized them.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/torchjd/autojac/_transform/_differentiate.py	100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/_grad.py	100.00% <100.00%> (ø)
src/torchjd/autojac/_transform/_jac.py	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

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

[PierreQuinton]

I don't remember why we have all this things with chunks rather than using chunk_size of vmap (For the chunk_size=1 I get it). I guess there was a good reason to have that so LGTM.

[ValerianRey]

I don't remember why we have all this things with chunks rather than using chunk_size of vmap (For the chunk_size=1 I get it). I guess there was a good reason to have that so LGTM.

2 reasons IIRC:

• Avoid using vmap when chunk_size=1, because some modules like RNNs on some version of cuda are not compatible with vmap
• Be able to not retain graph for the last chunk

[PierreQuinton]

Right, this is quite nasty. Good thing we have one memory for two. I think this can be merged, unless we should test some more?

[ValerianRey]

Yes ready to merge.