[BUG] reshape, flatten and unflatten with inplace=True keep dim names that no longer match the batch dims #2492

Anyone who names the batch dims of a tensordict and then calls reshape, flatten or unflatten with inplace=True ends up with names that describe the old shape. The names list can even have a different length from batch_size, so later named-dim lookups (refine_names, rename, indexing by name) work on the wrong dims or fail. The same calls without inplace=True give the right names.

import torch
from tensordict import TensorDict

def make():
    return TensorDict({"x": torch.zeros(2, 3, 4)}, batch_size=[2, 3], names=["a", "b"])

print(make().reshape(6).names, make().flatten(0, 1).names)
for op in (
    lambda td: td.reshape(6, inplace=True),
    lambda td: td.flatten(0, 1, inplace=True),
    lambda td: td.unflatten(1, (3, 1), inplace=True),
):
    td = make()
    op(td)
    print(td.batch_size, td.names)

[None] [None]
torch.Size([6]) ['a', 'b']
torch.Size([6]) ['a', 'b']
torch.Size([2, 3, 1]) ['a', 'b']

Expected: the in-place calls should give the same names as the out-of-place ones: [None] after reshape(6) and flatten(0, 1), and the names of unflatten(1, (3, 1)) without inplace after unflatten. masked_select(mask, inplace=True) already does this: it erases the names once the masked dims merge into one.

Cause: the in-place branches of reshape (tensordict/_base/shape.py:531), flatten (tensordict/_base/shape.py:2323) and unflatten (tensordict/_base/shape.py:2394) return self._inplace_rebind_leaves(...), which changes the batch size through _change_batch_size (tensordict/base.py:7338). That only sets _batch_size and leaves _td_dim_names alone. flatten computes the new names before the branch (tensordict/_base/shape.py:2302) but passes them only to the out-of-place _fast_apply. masked_select erases the names after its rebind (tensordict/_base/shape.py:2611-2613).

Possible fix: after the rebind, set the names that the out-of-place path would give (self._erase_names() for reshape, the computed names for flatten and unflatten), on the root and, through nested_fn, on nested tensordicts. Only tensordicts with names change; calls without names are unaffected.

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