[BUG] LazyStackedTensorDict.empty() drops all dim names #2500

empty() on a lazy stack with named dims returns a stack without any names: the stack dim and the dims of the members both lose theirs. The same call on a TensorDict keeps them, so code that builds an output buffer with td.empty() and then indexes or reshapes by name fails only when td is a lazy stack.

import torch
from tensordict import LazyStackedTensorDict, TensorDict

td = LazyStackedTensorDict(
    *[TensorDict(a=torch.zeros(3), batch_size=[3]) for _ in range(2)],
    stack_dim=0,
)
td.names = ["s", "c"]
print(td.names, td.empty().names)
print(td.to_tensordict().empty().names)
print(td.tensordicts[0].names, td.empty().tensordicts[0].names)

['s', 'c'] [None, None]
['s', 'c']
['c'] [None]

Expected: td.empty().names is ['s', 'c'], as for td.to_tensordict().empty(). td.empty(names=[...]) already sets the names it is given.

Cause: LazyStackedTensorDict.empty (tensordict/_lazy.py:2723-2755) defaults names to None, then passes names=None to each member's empty and stack_dim_name=None to the new stack. TensorDict.empty (tensordict/_td.py:2224-2240) uses NO_DEFAULT to mean "keep the names", and with recurse=False an explicit names=None erases them. The stack name is only read from names (_lazy.py:2738-2740), never from self._td_dim_name.

Possible fix: when no names are given, keep self._td_dim_name for the stack dim and do not pass names to the members, so they keep their own. The stack branch that falls back to TensorDict.empty for a new batch size (_lazy.py:2727-2736) passes names=None too and needs the same care. Only results that lose their names today change. Found while working on #2467.

tensordict main (eba7879), torch 2.14.1, Python 3.12.12.