feat: Add autojac.jac

[ValerianRey]

The idea is to have a way to easily compute Jacobians. The interface is very similar to that of torch.autograd.grad. The output is a tuple of jacobians (not necessarily matrices) instead of a tuple of gradients.

I think it would be nice to have more tests, and also maybe some utility functions or examples to show people how to transform these jacobians into matrices or how to combine them into a single matrix.

Lastly, it would be nice to have a grad_outputs param, or maybe a jac_outputs, idk.

@PierreQuinton what do you think of this interface?

[codecov]

Codecov Report

❌ Patch coverage is 89.65517% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/torchjd/autojac/_jac.py	88.88%	3 Missing ⚠️

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

🚀 New features to boost your workflow:

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

[PierreQuinton]

Nice, thx.

I would vouch for having a jac_output parameter. Later on this will have the option of being a SLT which allows user to chose between SLT/classical jac: Somehow the user can chose what algorithm they use, and this is especially nice.

I think this change would also have to be made to backward and mtl_backward, so we can do that in another PR, it also seems that it would make most Transform used in the current implementation of jac useless (no init, no diagonalize unless we are using the default param for jac_output).