Fix proxy repr for non-tensor proxy (Lightning-AI#692)

thunder/core/proxies.py

@@ -134,7 +134,9 @@ def replace_name(self, name: str | None = None):
         return self.__class__(name=name)
 
     def __repr__(self) -> str:
-        return f'<{type(self).__name__}(name="{self.name}", dtype={self.dtype}, shape={self.shape}>'
+        # All subclasses of Proxy will have `self.name`, so this generic implementation relies on that.
+        # To have a specific repr for a subclass, override the implementation for that subclass.
+        return f'<{type(self).__name__}(name="{self.name}")>'
 
     def type_string(self) -> str:
         return "Any"
@@ -1198,6 +1200,9 @@ def true_dtype(self):
     def requires_grad(self):
         return self._requires_grad
 
+    def __repr__(self):
+        return f'<{type(self).__name__}(name="{self.name}", dtype={self.dtype}, shape={self.shape})>'
+
     def type_string(self):
         return f"FUTURE {self.device} {self.dtype.shortname()}{list(self.shape)}"
 
@@ -1293,6 +1298,9 @@ def replace_name(self, name: str):
         """Return a copy of this proxy with the given name."""
         return tensorproxy(self, name=name, history=self.history)
 
+    def __repr__(self):
+        return f'<{type(self).__name__}(name="{self.name}", dtype={self.dtype}, shape={self.shape})>'
+
     def type_string(self):
         return f"{self.device} {self.dtype.shortname()}{list(self.shape)}"
 

thunder/tests/test_core.py

@@ -1244,7 +1244,7 @@ def foo(x, y, z):
     region_bsyms = trace.bound_symbols[:3]
     region = Region(producers, consumers, region_bsyms)
     assert len(region.inputs) == 0 and sorted(str(v) for v in region.outputs) == [
-        '<TensorProxy(name="t0", dtype=thunder.dtypes.float32, shape=(1,)>'
+        '<TensorProxy(name="t0", dtype=thunder.dtypes.float32, shape=(1,))>'
     ]
 
 
@@ -2853,3 +2853,21 @@ def forward(self, x) -> torch.Tensor:
     gradcheck(jitted, (x,))
     with pytest.raises(GradcheckError):
         gradcheck(model, (x,))
+
+
+def test_proxy_repr():
+    # Verify that we can call `__repr__` on different proxy subclasses.
+    t = thunder.core.trace.TraceCtx()
+    with thunder.core.trace.tracectx(t):
+        p = thunder.core.proxies.NumberProxy("number", 1, python_type=int)
+        c = thunder.core.proxies.CollectionProxy((1, 2), name="collection")
+        t = thunder.core.proxies.TensorProxy(
+            "tensor",
+            shape=(1,),
+            dtype=thunder.core.dtypes.float16,
+            device=thunder.core.devices.Device("cpu"),
+            requires_grad=True,
+        )
+        assert p.__repr__() == '<NumberProxy(name="number")>'
+        assert t.__repr__() == '<TensorProxy(name="tensor", dtype=thunder.dtypes.float16, shape=(1,))>'
+        assert c.__repr__() == '<CollectionProxy(name="collection")>'