[BUG] pad_sequence raises TypeError on tensordicts that hold only non-tensor data #2490

tensordict.pad_sequence of tensordicts with a batch size whose entries are all non-tensor data, such as strings, raises TypeError from torch.Size. The same tensordicts with one tensor entry added pad fine, and the strings are padded with None. Without a batch size, only-non-tensor tensordicts also work.

from tensordict import NonTensorData, TensorDict, pad_sequence

tds = [TensorDict(s=NonTensorData("x", batch_size=[n]), batch_size=[n]) for n in (2, 3)]
pad_sequence(tds)

TypeError: torch.Size() takes an iterable of 'int' (item 1 is 'float')

Expected: a tensordict of batch size [2, 3] whose s is [["x", "x", None], ["x", "x", "x"]], as pad_sequence returns when each tensordict also holds a=torch.zeros(n). Torch's pad_sequence pads to the longest sequence, here 3.

Cause: pad_sequence (tensordict/functional.py:474) takes the padded length as the maximum of the tensor leaves' sizes along the pad dim, starting from max_seq_length = float("-inf") (tensordict/functional.py:564). Non-tensor leaves are skipped (tensordict/functional.py:576), so with no tensor leaf the length stays -inf, is written into the batch size (tensordict/functional.py:610), and torch.Size(shape) raises (tensordict/functional.py:618).

Possible fix: when the pad dim is a batch dim, also take each tensordict's batch_size[pad_dim] into the maximum. A tensor leaf has the same size as the batch along a batch dim, so only calls that raise today would change.

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