feat(autojac): Add tangents to backward and jac

[PierreQuinton]

• Add jac_tensors parameter to backward
• Add jac_outputs parameter to jac
• Add some extra usage examples using these parameters
• Add extra tests and doctests
• Add changelog entry

Different PR:

• Add grad_outputs to mtl_backward
• Make these params name-only

[ValerianRey]

@claude Review this.

[claude]

Claude finished @ValerianRey's task in 1m 14s —— View job

---

Code Review

Tasks

• Read changed files and understand modifications
• Analyze code quality and correctness
• Check adherence to project standards
• Provide review feedback

---

Summary

This PR adds a new jac_tensors parameter to the backward function, allowing users to pre-multiply the Jacobian computation with an initial Jacobian matrix. The implementation is generally solid and well-tested, but I found several issues that should be addressed.

Issues Found

1. 🐛 Bug in Example 2 Documentation (Critical)

Location: src/torchjd/autojac/_backward.py:82

The second example claims to demonstrate using jac_tensors but actually doesn't use them:

>>> J1 = torch.tensor([1.0, 0.0])
>>> J2 = torch.tensor([0.0, 1.0])
>>>
>>> backward([y1, y2])  # ❌ Should be: backward([y1, y2], jac_tensors=[J1, J2])

The variables J1 and J2 are defined but never used. This makes the example misleading and defeats its stated purpose.

Impact: Users following this example won't understand how to actually use jac_tensors.

Fix this →

2. 📝 Typo in Example 3 Documentation

Location: src/torchjd/autojac/_backward.py:98

Missing > in the docstring:

>>> from torchjd.autojac import backward
>>   # ❌ Should be: >>>

3. 🧹 Incomplete Comment

Location: src/torchjd/autojac/_backward.py:136

# correspond to the shape of the corresponding element in tensors_

This comment is incomplete. It appears to be a sentence fragment. Should probably be:

# Check that jac_tensors correspond to the shape of the corresponding element in tensors_

4. 🔍 Removed Test Could Be Valuable

Location: tests/unit/autojac/test_backward.py

The PR removes test_check_create_transform() which validated that _create_transform returns a valid Transform object. While the helper function _create_transform was removed (inlined into backward), this test verified an important invariant about the transform composition.

Recommendation: Consider whether the transform composition logic still needs validation, even if the helper is gone.

Positive Observations

✅ Excellent test coverage: Added comprehensive tests for:

• Single and multiple jac_tensors
• Different row counts
• Error cases (length mismatch, shape mismatch, inconsistent dimensions)

✅ Documentation examples are tested: All three docstring examples have corresponding tests in tests/doc/test_backward.py

✅ Type hints are correct: Proper use of Sequence[Tensor] | Tensor | None

✅ Mathematical documentation is clear: The explanation of the matrix product J_init · J is helpful

✅ Error handling: Proper validation with strict=True in zip() ensures length matching

Minor Suggestions

1. Docstring consistency: The third example description says "essentially equivalent to autograd.grad" - consider being more precise about how it differs (if at all).

2. Parameter naming: The internal variable name jac_tensors_dict is clear and descriptive.

Adherence to Project Standards

• ✅ Sphinx-style docstrings with :param format
• ✅ Type hints in function prototype
• ✅ No unnecessary :returns: key
• ✅ Tests added for new functionality
• ⚠️ Should verify with uv run ty check and uv run ruff check before merging

Recommendation

Request changes to fix the critical bug in Example 2. The other issues are minor but should be addressed for code quality.

---

[ValerianRey]

@PierreQuinton I think claude caught a few interesting things. I'll make a proper review later.

[PierreQuinton]

@claude I think we are supposed to add an entry to the changelog for this as it changes the public interface (backward is a public function), could you suggest one?

[claude]

Claude finished @PierreQuinton's task in 35s —— View job

---

Suggesting Changelog Entry

Tasks

• Read changelog file to understand format
• Review the backward function changes
• Suggest changelog entry

---

Based on the changelog format and the changes in this PR, here's my suggested entry for the Unreleased section under ### Added:

- Added a `jac_tensors` parameter to `torchjd.autojac.backward`. This allows users to compute
  pre-multiplied Jacobians by providing initial Jacobian tensors to backpropagate. When
  `jac_tensors` is provided, the function computes the matrix product `J_init · J`, where `J` is
  the Jacobian of `tensors` w.r.t `inputs`, and `J_init` is the concatenation of `jac_tensors`.
  If `None`, it defaults to the Identity matrix, resulting in the standard Jacobian computation.

Rationale:

• This follows the existing changelog style of describing what was added and its purpose
• It explains the mathematical behavior concisely (similar to how other features are described)
• It mentions the public API change (torchjd.autojac.backward)
• It clarifies the default behavior when jac_tensors is None

Alternative shorter version (if you prefer brevity):

- Added a `jac_tensors` parameter to `torchjd.autojac.backward`, allowing users to pre-multiply the
  Jacobian computation with an initial Jacobian matrix. When `None`, defaults to the standard
  Jacobian computation.

[ValerianRey]

@claude please review again.

[claude]

Claude finished @ValerianRey's task in 1m 40s —— View job

---

Code Review

Tasks

• Read all changed files
• Check if previous review issues were addressed
• Analyze new review comment context
• Provide comprehensive review feedback

---

Summary

This PR successfully adds jac_tensors and jac_outputs parameters to the backward and jac functions, enabling multi-step chain rule computations. The implementation is solid and well-tested. Most issues from the previous review have been addressed, but there are still a few documentation issues to fix.

Status of Previous Review Issues

✅ Fixed: Example 2 Bug - The second example in _backward.py now correctly uses jac_tensors=[J1, J2] (line 88). The corresponding doctest in tests/doc/test_backward.py:24-40 validates this.

✅ Fixed: Typo in Example 3 - No longer present (Example 3 doesn't exist in current version).

✅ Fixed: Incomplete Comment - The comment issue has been resolved. The validation logic is now clearer.

✅ Addressed: Test Removal - The test_check_create_transform() test has been kept and updated to work with the new _create_jac_tensors_dict helper.

Current Issues

1. 📝 Missing Documentation: Consistent First Dimension Requirement

Location: src/torchjd/autojac/_backward.py:30-35 and src/torchjd/autojac/_jac.py:37-43

Issue: The documentation for jac_tensors/jac_outputs doesn't explicitly state that all tensors must have the same first dimension (number of rows). The code enforces this via check_consistent_first_dimension() on line 143 of _backward.py and line 190 of _jac.py, but users won't know this requirement until they hit the error.

Current text (lines 30-33 in _backward.py):

:param jac_tensors: The initial Jacobians to backpropagate, analog to the `grad_tensors`
    parameter of `torch.autograd.backward`. If provided, it must have the same structure as
    ``tensors`` and each tensor in ``jac_tensors`` must match the shape of the corresponding
    tensor in ``tensors``, with an extra leading dimension representing the number of rows of
    the resulting Jacobian (e.g. the number of losses). If ``None``, defaults to the identity
    matrix. In this case, the standard Jacobian of ``tensors`` is computed, with one row for
    each value in the ``tensors``.

Suggested fix: Add a sentence clarifying the first dimension must be consistent:

:param jac_tensors: The initial Jacobians to backpropagate, analog to the `grad_tensors`
    parameter of `torch.autograd.backward`. If provided, it must have the same structure as
    ``tensors`` and each tensor in ``jac_tensors`` must match the shape of the corresponding
    tensor in ``tensors``, with an extra leading dimension representing the number of rows of
    the resulting Jacobian (e.g. the number of losses). All tensors in ``jac_tensors`` must
    have the same first dimension. If ``None``, defaults to the identity matrix. In this case,
    the standard Jacobian of ``tensors`` is computed, with one row for each value in the
    ``tensors``.

Same issue in _jac.py at lines 37-43 for jac_outputs parameter.

Fix this in backward →

Fix this in jac →

2. 🔍 Minor: Confusing Example 2 Description

Location: src/torchjd/autojac/_backward.py:73-74

Issue: The description says "we explicitly specify the identity jac_tensors" but the actual vectors used are:

J1 = torch.tensor([1.0, 0.0])
J2 = torch.tensor([0.0, 1.0])

These aren't technically "the identity" - they're the rows of an identity matrix. When used together, they form an identity transformation, but individually they're not identity matrices. The phrasing could be clearer.

Suggested improvement:

This is the same example as before, except that we explicitly specify ``jac_tensors`` as
the rows of the identity matrix (which is equivalent to using the default ``None``).

3. ✅ Validation Tests Are Excellent

The PR adds comprehensive validation with helpful error messages:

• check_matching_length - Validates len(jac_tensors) == len(tensors)
• check_matching_shapes - Validates shapes match (excluding first dimension)
• check_consistent_first_dimension - Validates all first dimensions are equal

All three have corresponding unit tests in test_backward.py:

• test_jac_tensors_length_mismatch (lines 135-147)
• test_jac_tensors_shape_mismatch (lines 150-164)
• test_jac_tensors_inconsistent_first_dimension (lines 167-190)

4. ✅ Test Coverage Is Comprehensive

The PR includes excellent test coverage:

• Unit tests: Multiple scenarios with different row counts (1, 2, 3, 5)
• Doc tests: All three docstring examples are tested in tests/doc/test_backward.py and tests/doc/test_jac.py
• Edge cases: Single vs multiple tensors, various parallel_chunk_size values

5. ✅ Changelog Entries Are Good

The changelog entries (lines 16-20) accurately describe the new functionality and correctly reference the analogous PyTorch functions.

Code Quality Observations

✅ Type hints: Properly used throughout with Sequence[Tensor] | Tensor | None

✅ Consistent API: Both backward and jac have parallel implementations with appropriate parameter names (jac_tensors vs jac_outputs)

✅ Mathematical correctness: The matrix product J_init · J is correctly implemented via the transform composition

✅ Error messages: Clear and informative, specifying parameter names and shapes

✅ Helper functions: _create_jac_tensors_dict and _create_jac_outputs_dict follow good separation of concerns

Recommendations

1. Required: Add documentation about consistent first dimension requirement (Issue Experiment with double column form sphinx design package. #1 above)
2. Suggested: Clarify Example 2 description (Issue CICD: Running pre-commit without modifying files #2 above)
3. Before merge: Run uv run ty check, uv run ruff check, and uv run ruff format per CLAUDE.md
4. Before merge: Run affected tests: uv run pytest tests/unit/autojac/test_backward.py tests/unit/autojac/test_jac.py tests/doc/test_backward.py tests/doc/test_jac.py -W error

Verdict

Request minor changes to add the missing documentation about the consistent first dimension requirement. Otherwise, this is high-quality work with excellent test coverage.

---

View job