[BugFix] Add TensorDictModuleBase repr

tensordict/nn/common.py

@@ -793,8 +793,17 @@ def __name__(self):
         # This is necessary to make compiled vmap over TDModule happy
         return self.__class__.__name__
 
-    def __repr__(self):
-        return f"{self.__class__.__name__}()"
+    def __repr__(self) -> str:
+        entries = [f"{name}={module}" for name, module in self.named_children()]
+        # a wrapped callable that is not an nn.Module (a plain function, a lambda, ...)
+        # is stored as a regular attribute, hence missing from the children
+        module = self.__dict__.get("module")
+        if module is not None:
+            entries.insert(0, f"module={module}")
+        entries.append(f"in_keys={self.in_keys}")
+        entries.append(f"out_keys={self.out_keys}")
+        fields = indent(",\n".join(entries), 4 * " ")
+        return f"{type(self).__name__}(\n{fields})"
 
 
 class TensorDictModule(TensorDictModuleBase):
@@ -1225,17 +1234,6 @@ def device(self) -> torch.device:
             return p.device
         return torch.device("cpu")
 
-    def __repr__(self) -> str:
-        fields = indent(
-            f"module={self.module},\n"
-            f"device={self.device},\n"
-            f"in_keys={self.in_keys},\n"
-            f"out_keys={self.out_keys}",
-            4 * " ",
-        )
-
-        return f"{type(self).__name__}(\n{fields})"
-
     def __getattr__(self, name: str) -> Any:
         if not is_compiling():
             __dict__ = self.__dict__

test/nn/test_nn.py

@@ -360,6 +360,40 @@ def my_func(self, tensor: torch.Tensor, *, an_integer: int):
 
         assert td["b"] == 4
 
+    def test_repr(self):
+        class MyModule(TensorDictModuleBase):
+            in_keys = ["a"]
+            out_keys = ["b"]
+
+            def __init__(self):
+                super().__init__()
+                self.net = nn.Linear(3, 4)
+
+            def forward(self, tensordict):
+                tensordict.set("b", self.net(tensordict.get("a")))
+                return tensordict
+
+        string = repr(MyModule())
+        assert "MyModule(" in string
+        assert "net=Linear(in_features=3, out_features=4" in string
+        assert "device=" not in string
+        assert "in_keys=['a']" in string
+        assert "out_keys=['b']" in string
+
+        # an nn.Module payload is a child: listed once, and not through __dict__
+        mod = TensorDictModule(nn.Linear(3, 4), in_keys=["a"], out_keys=["b"])
+        assert "module" not in mod.__dict__
+        string = repr(mod)
+        assert string.count("module=") == 1
+        assert "module=Linear(in_features=3, out_features=4" in string
+
+        # a callable that is not an nn.Module is not a child: it comes from __dict__
+        mod = TensorDictModule(lambda x: x + 1, in_keys=["a"], out_keys=["b"])
+        assert dict(mod.named_children()) == {}
+        string = repr(mod)
+        assert string.count("module=") == 1
+        assert "module=<function" in string
+
     def test_mutable_sequence(self):
         in_keys = self.MyMutableSequence(["a", "b", "c"])
         out_keys = self.MyMutableSequence(["d", "e", "f"])