[Feature] Add TensorDict.backward()

[vmoens]

Description

Adds TensorDictBase.backward() with semantics matching torch.Tensor.backward() across differentiable tensor leaves.

tensordict.backward(
    gradient=None,
    retain_graph=None,
    create_graph=False,
    inputs=None,
)

Behavior

• Recursively collects differentiable tensor leaves and differentiates the graph through a single torch.autograd.backward(tensors, grad_tensors, ...) call. Non-differentiable leaves are ignored; an error is raised when no leaf requires gradients.
• With gradient=None, an implicit gradient of 1 is used when every differentiable leaf is scalar, so loss_td.backward() is equivalent to loss_td.sum(reduce=True).backward(). The standard PyTorch-style error is raised if any differentiable leaf is non-scalar.
• With an explicit gradient, a matching TensorDictBase is required. Gradients are matched to leaves by nested key, with clear errors for missing keys and incompatible shapes. Heterogeneous leaf shapes, devices and dtypes are supported, which enables weighted losses:

weights = TensorDict(
    actor_loss=torch.tensor(0.5),
    critic_loss=torch.tensor(1.0),
)
loss_td.backward(weights)

• retain_graph, create_graph and inputs are forwarded consistently with torch.Tensor.backward(); inputs also accepts a TensorDictBase.
• TensorDict.sum() semantics are unchanged: plain .sum() still reduces each leaf independently and returns a TensorDict, and .sum(reduce=True) remains the explicit full reduction.
• The method is exposed on tensorclasses through the TD fallback list (with the matching .pyi stub).

Tests

TestBackward in test/tensordict/test_methods.py covers scalar and non-scalar leaves, nested keys, explicit gradient TensorDicts, weighted gradients (including parity with (loss_td * weights).sum(reduce=True).backward()), mixed differentiable/non-differentiable leaves, missing/incompatible/mis-typed gradients, no differentiable leaves, retain_graph/create_graph/inputs forwarding, and parity with the equivalent plain-PyTorch expressions. A parametrized test_backward runs against all TensorDict variants (nested, stacked, permuted, h5, params, ...), plus a tensorclass test.

Docs

Documented in the "Gradient computation" section of docs/source/reference/td.rst; the docstring ships a runnable example.

🤖 Generated with Claude Code