feat(autojac): Add grad_tensors to mtl_backward

[ValerianRey]

• Rename losses to tensors in mtl_backward
• Add grad_tensors parameter to mtl_backward
• Update unit tests for mtl_backward and add new tests for the grad_tensors param
• Rename check_matching_shapes to check_matching_jac_shapes
• Add check_matching_grad_shapes
• Add changelog entries

[PierreQuinton]

@claude please review this.

[claude]

Code Review

• Gather context
• Read changed files
• Analyze changes
• Post review feedback

---

Overall this is a clean, well-scoped change. The addition of grad_tensors to mtl_backward is analogous to how jac_tensors was added to backward, and the implementation follows the same patterns. Below are a few issues and observations.

---

Issues

1. Wrong parameter name in error message (_mtl_backward.py:214)

The error message in _check_tensors_are_scalar says grad_outputs but the public parameter is named grad_tensors:

# current (wrong):
"When `tensors` are non-scalar, the `grad_outputs` parameter must be provided."

# should be:
"When `tensors` are non-scalar, the `grad_tensors` parameter must be provided."

Fix this →

---

2. Docstring line too long (_mtl_backward.py:43)

This line exceeds typical line-length limits:

If the ``tensors`` are non-scalar, ``mtl_backward`` requires some initial gradients in ``grad_tensors``.

PierreQuinton's review already suggests a reformatted version — worth applying.

---

3. Docstring param description for grad_tensors says grad_outputs (_mtl_backward.py:53)

# current:
parameter of ``torch.autograd.backward``. If any of the ``tensors`` is non-scalar, ``grad_tensors`` must

PierreQuinton's second suggestion is also about this line. The wording here is actually fine in the current PR but it's worth double checking the suggestion is applied.

---

Minor observations

4. Type signature inconsistency in private helper

The public mtl_backward signature declares grad_tensors: Sequence[Tensor] | None, but _create_grad_tensors_dict accepts Sequence[Tensor] | Tensor | None — accepting a single Tensor. This internal overload is never exposed from the public API, so it's dead code unless intentional. Not a bug, but worth either documenting the intent or aligning the types.

5. Missing punctuation in error message (_mtl_backward.py:110)

raise ValueError("`tensors` cannot be empty")

The features error message has a period at the end (line 107: "`features` cannot be empty."), but the tensors one doesn't. Minor but inconsistent.

Fix this →

6. test_grad_tensors_value_is_correct — tasks_params and shared_params arguments

In the new test, tasks_params=None and shared_params=None are passed explicitly (lines 713–718). This is fine but slightly inconsistent with the rest of the test suite where those are often omitted (defaults). Not a correctness issue.

7. CHANGELOG entry wording (CHANGELOG.md:9-11)

The entry says allowing to use non-scalar \losses`but the parameter was just renamed fromlossestotensors. A cleaner phrasing: allowing to use non-scalar `tensors`` (which is the new name). Worth checking for consistency.

---

Summary

The implementation is correct and the tests are thorough (scalar/non-scalar cases, shape mismatch, length mismatch, value correctness). The main actionable fix is the wrong parameter name (grad_outputs → grad_tensors) in the error message at line 214.

View job run