[BugFix] Preserve uniform metadata during non-tensor concatenation

Author: vmoens

tensordict/tensorclass.py

@@ -4663,9 +4663,15 @@ def _cat_non_tensor(cls, tensors, dim=0, out=None):
                     f"dimension, got {batch_size} and {shape}."
                 )
             values.extend(tensor.unbind(dim))
-        result = (
-            cls._stack_non_tensor(values, dim=dim) if values else tensors[0].clone()
-        )
+        if values and all(isinstance(tensor, cls) for tensor in tensors):
+            # Concatenating uniform data preserves scalar metadata, including
+            # callable payloads, independently of the stacking capture setting.
+            with set_capture_non_tensor_stack(True):
+                result = cls._stack_non_tensor(values, dim=dim)
+        else:
+            result = (
+                cls._stack_non_tensor(values, dim=dim) if values else tensors[0].clone()
+            )
         if out is not None:
             if out.batch_size != result.batch_size:
                 raise RuntimeError("out.batch_size and cat batch size must match.")

test/tensordict/test_nontensor.py

@@ -240,12 +240,17 @@ def test_cat_preserves_non_tensor_values(self, capture, dim, stacked, with_out):
         assert result.tolist() == expected.tolist()
 
     @pytest.mark.parametrize("capture", [False, True])
-    def test_cat_non_tensor_data_out(self, capture):
+    @pytest.mark.parametrize("with_out", [False, True])
+    def test_cat_uniform_non_tensor_data(self, capture, with_out):
         items = [NonTensorData("value", batch_size=[2])] * 2
         out = NonTensorData("old", batch_size=[4])
         with set_capture_non_tensor_stack(capture):
-            assert torch.cat(items, out=out) is out
-        assert out.tolist() == ["value"] * 4
+            result = torch.cat(items, out=out if with_out else None)
+        if with_out:
+            assert result is out
+        assert isinstance(result, NonTensorData)
+        assert result.data == "value"
+        assert result.tolist() == ["value"] * 4
 
     @pytest.mark.parametrize("capture", [False, True])
     def test_cat_pads_nested_non_tensor_values(self, capture):