Composite specs can create named tensors with 'zero' and 'rand'

Author: Louis Faury

test/test_specs.py

@@ -4585,6 +4585,26 @@ def test_names_repr(self):
         assert "Composite" in repr_str
         assert "obs" in repr_str
 
+    def test_zero_create_names(self):
+        """Test that creating tensors with 'zero' propagates names."""
+        spec = Composite(
+            {"obs": Bounded(low=-1, high=1, shape=(10, 3, 4))},
+            shape=(10,),
+            names=["batch"],
+        )
+        td = spec.zero()
+        td.names = ["batch"]
+
+    def test_rand_create_names(self):
+        """Test that creating tensors with 'rand' propagates names."""
+        spec = Composite(
+            {"obs": Bounded(low=-1, high=1, shape=(10, 3, 4))},
+            shape=(10,),
+            names=["batch"],
+        )
+        td = spec.rand()
+        td.names = ["batch"]
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

torchrl/data/tensor_specs.py

@@ -5750,6 +5750,7 @@ def rand(self, shape: torch.Size = None) -> TensorDictBase:
             _dict,
             batch_size=_size([*shape, *_remove_neg_shapes(self.shape)]),
             device=self.device,
+            names=self.names,
         )
 
     def keys(
@@ -6030,6 +6031,7 @@ def zero(self, shape: torch.Size = None) -> TensorDictBase:
             },
             batch_size=_size([*shape, *self._safe_shape]),
             device=device,
+            names=self.names,
         )
 
     def __eq__(self, other: object) -> bool: