[peterdsharpe]

Closes #3005.

A tensor with the batch shape, such as a per-sample weight, is meant to line up with the batch dims of every leaf. When a tensordict holds a lazy stack below its root, as heterogeneous data or a lazy_stack of observations often does, add, mul, maximum and the other pointwise ops lined it up with the feature dims of the stacked leaves instead. The call raised when the sizes differed, and gave wrong values with no error when they matched. With this PR the nested lazy stack gives the same values as the dense tensordict, also when a tensorclass or a TypedTensorDict wraps it, and lerp, addcmul and addcdiv, which #3002 made raise on such trees on purpose, now work too.

What was wrong

import torch
from tensordict import TensorDict, lazy_stack

w = torch.tensor([0.0, 1.0])
members = [TensorDict(a=torch.arange(2.0)), TensorDict(a=torch.arange(2.0))]
td = TensorDict(s=lazy_stack(members, 0), batch_size=[2])
dense = TensorDict(s=lazy_stack(members, 0).contiguous(), batch_size=[2])

dense.add(w)["s", "a"]    # [[0., 1.], [1., 2.]]
td.add(w)["s", "a"]       # main: [[0., 2.], [0., 2.]]
td.lerp(td + 1, w)        # main: TypeError: Multiple dispatch failed for 'torch.Tensor.lerp'

A lazy stack of tensorclasses, or a TypedTensorDict built on a lazy stack with from_tensordict, nested in a tensordict, gives the same wrong values for add on main.

The fix

For a tensor operand, the broadcast branch of _maybe_broadcast_other (tensordict/base.py) calls self._fast_apply(lambda x: op(x, expand_as_right(other, x))). A nested lazy stack runs that function on the leaves of its members (LazyStackedTensorDict._apply_nest), and those leaves lack the stack dim, so expand_as_right lines w up with the wrong dims. A lazy stack at the root never gets there: an earlier branch gives each member its slice of the operand (#2365).

When the tree holds a lazy stack (_has_lazy_stack, which looks through tensorclasses and TypedTensorDicts), the _fast_apply call now passes is_leaf=_is_leaf_or_lazy_stack, a new private helper. It accepts what _default_is_leaf accepts, plus a LazyStackedTensorDict that is not a NonTensorStack, a tensorclass and a TypedTensorDict. The function then runs on the nested stack, or on the wrapper that holds it: expand_as_right expands w to its batch size, and the op called on it reaches the root branch, which splits w along the stack dim. Deeper nesting is handled the same way at each level. The two-operand branch (lerp, addcmul, addcdiv) passes the same is_leaf to _apply_pointwise, which pairs the nested stack with the entry of the tensordict operand under the same key. That made the _has_nested_lazy_stack guard of #3002 unnecessary, so I removed it; it had no other caller.

A TypedTensorDict delegates its pointwise ops to the base class, so the op called on one that wraps a lazy stack ran _fast_apply through the stack's members as well: on main S.from_tensordict(lazy_stack(...)).add(w) gives wrong values at the root too (the in-place forms were right). The root branch now runs the op on the lazy stack that a TypedTensorDict wraps and wraps the result in the same class. Without this, the nested case above would still give wrong values for a TypedTensorDict.

Trees with no lazy stack keep main's path: is_leaf stays None, and the walk of _has_lazy_stack is the only added step (2.7 us on a tree of 30 leaves). Calling the op on every tensorclass instead cost 30% on a dense tree that holds tensorclasses, so the wider is_leaf is used only when it is needed. _fast_apply keeps its dense fast path (_dense_apply) for the new is_leaf, as it does for the default one.

What I checked:

• A grid of 11,040 calls on main and on this branch, compared leaf by leaf (values, dtype, shape and strides; NaNs count as equal), with the result, whether an in-place op returns self, and the state of self afterwards: add, add(alpha=2), sub, mul, div, pow, %, rsub, +, w + td, +=, maximum, minimum, clamp_min, clamp_max_, clamp, lerp (tensordict end, dense end, tensor end), addcmul, addcdiv and the in-place forms. The trees: dense; a lazy stack at the root; a lazy stack nested one or two levels down; a lazy stack of tensorclasses at the root and nested; a TypedTensorDict on a lazy stack at the root and nested; and a dense tensorclass and a dense TypedTensorDict nested in a tensordict. Stack dims 0 and 1, homogeneous and heterogeneous members, members that hold a nested lazy stack themselves, and operands of the batch shape, a broadcastable shape, a shape with an extra leading dim, and 0-d. Every call on a dense tree (with or without tensorclasses and TypedTensorDicts) and every call on a root lazy stack whose members hold no lazy stack gives the same result as main. No call that works on main raises here.
• The calls that changed are right. Of the 5,580 calls on lazy trees that succeed here, 4,928 were checked against the op on the densified tree, or, for heterogeneous members, against the op on each densified member with its slice of the operand. All agree except 220 in-place ops with an operand that adds a leading dim to a root lazy stack (plain, of tensorclasses, or in a TypedTensorDict), where the dense tensordict raises and the lazy stack returns a new, larger tensordict. That is main's behaviour for a root lazy stack ([BugFix] clamp writes into out, and the in-place and two-operand pointwise ops broadcast a batch-shaped tensor #3002), unchanged here. 432 calls with heterogeneous members and an operand with an extra leading dim were not checked.
• A tensorclass that holds a nested lazy stack, a locked tensordict, a tensordict with dim names, a nested lazy stack with a non-tensor leaf, and a dense tensordict with a nested-lazy end or tensor1 all give the dense values; tensorclasses stay tensorclasses and TypedTensorDicts keep their class. A TensorDictParams that holds a lazy stack raises as on main, and memory-mapped tensordicts give main's results.
• torch.compile(fullgraph=True) compiles add(w), add_(w), lerp(td + 1, w) and addcmul(w, td + 1) on a dense tree, a root lazy stack and a lazy stack nested plain, as a tensorclass or as a TypedTensorDict, with values equal to eager, when the members of the stack have a batch dim. With members of batch size [], lerp on a nested lazy stack fails to compile (Inductor raises GuardOnDataDependentSymNode), as a root lazy stack of the same members does on main; main raised on that call in eager too. td * w does not compile with fullgraph=True on main either, for dense and lazy trees, and is unchanged.

Behaviour changes

Only calls on trees that hold a lazy stack below the root (or inside the members of a root lazy stack, or inside a TypedTensorDict) change:

• add, sub, mul, div, maximum, minimum, clamp, clamp_min, clamp_max and the other ops of _maybe_broadcast_other, with a tensor that has dims, now give the dense values. On main they raised RuntimeError from expand_as_right, or gave wrong values when the leaf shapes matched the operand. This includes lazy stacks of tensorclasses, a TypedTensorDict on a lazy stack (nested, and at the root for the out-of-place ops), and nested lazy stacks seen through a view such as unsqueeze or transpose.
• The in-place forms now return self and write into its leaves. On main they returned a new tensordict when they did not raise.
• lerp, lerp_, addcmul(_) and addcdiv(_) with a tensordict operand and a batch-shaped tensor now work; they raised TypeError.
• Some calls that raised still raise, with another error. addcmul and addcdiv with members that have an empty batch size give each member a 0-d slice, and then raise main's TypeError from _foreach_addcmul for a 0-d tensor with a tensordict, as a root lazy stack of such members already does on main (the 0-d bug noted in the [BUG] lerp, addcmul and addcdiv raise when one operand is a tensordict and the other a batch-shaped tensor #2963 triage, not fixed here). In-place ops with an operand that adds a leading dim raise the same RuntimeError as the dense tensordict. lerp with a dense end and heterogeneous members raises KeyError for the keys that only some members have.

The test

TestLazyStackedTensorDict.test_arithmetic_ops_broadcast_tensor_nested in test/tensordict/test_lazy.py checks add, sub, mul, maximum, clamp_min, +, lerp, addcmul and mul_ on a lazy stack nested one and two levels down, stack dims 0 and 1, against the dense tensordict. The leaves have the shape of the batch, so main gives wrong values without an error; all 4 cases fail on main with "Tensor-likes are not close!" and pass here. TestLazyStackedTensorDict.test_arithmetic_ops_broadcast_tensor_nested_wrapped does the same for a tensorclass and a TypedTensorDict that wrap the nested lazy stack, and checks that they keep their class; its 4 cases fail on main the same way and pass here. TestPointwiseOps.test_nested_lazy_and_tensor_operands in test/tensordict/test_misc.py pinned the TypeError of the guard; it now checks that lerp, addcmul, addcdiv and their in-place forms give the dense values and that the in-place forms return self. Its 12 cases fail on main and pass here.

What I ran

Local testing was CPU-only: Python 3.12, torch 2.14.1+cpu, with CI's environment variables (TD_GET_DEFAULTS_TO_NONE=1 LIST_TO_STACK=1 TD_CHECK_INVARIANTS=1).

• test/tensordict/test_misc.py and test/tensordict/test_lazy.py: 2757 passed, 4 skipped.
• test/tensordict/test_methods.py, test/tensordict/test_nontensor.py, test/nn/test_functorch.py, test/tensorclass/test_tensorclass.py, test/tensorclass/test_typedtensordict.py: 11541 passed, 458 skipped.
• test/compile/test_compile.py -k "add or lerp or pointwise or arith or mul or broadcast or lazy": 116 passed, 8 skipped.
• torch 2.13, test_misc.py and test_lazy.py: 2743 passed, 18 skipped.
• pre-commit run on the changed files (ruff check, ruff format, torchfix, whitespace): passed. import tensordict with -W error::DeprecationWarning is silent.
• Timing on an idle 24-core machine (taskset -c 5, torch 2.14.1+cpu, Python 3.14), main against this branch, the minimum of 5 interleaved runs of 7 repeats, with a batch-shaped tensor w of shape [64]:
  • 10 leaves: add(w) 31.72 against 32.92 us (+3.8%), mul(w) +3.1%, add_(w) +4.5%.
  • 100 leaves: add(w) 280.2 against 284.3 us (+1.5%), add_(w) +1.5%.
  • 10 leaves plus a tensorclass and a nested tensordict: add(w) 60.59 against 63.19 us (+4.3%), mul(w) +4.5%, add_(w) +4.2%.
  • lerp(td, w) is 1-3% faster in all three, and td + 1.0 is unchanged.
    The cost is the walk that looks for a lazy stack below the root (about 1 us for 10 leaves, 3 us with the nested entries), paid only when an operand is a tensor.