Add batch and event shape to repr. (#1946)

numpyro/distributions/distribution.py

@@ -563,6 +563,13 @@ def icdf(self, q: ArrayLike) -> ArrayLike:
     def is_discrete(self):
         return self.support.is_discrete
 
+    def __repr__(self) -> str:
+        cls = self.__class__
+        return (
+            f"<{cls.__module__}.{cls.__name__} object at {id(self):#x} with batch "
+            f"shape {self.batch_shape} and event shape {self.event_shape}>"
+        )
+
 
 @runtime_checkable
 class DistributionLike(Protocol):

test/test_distributions.py

@@ -3681,3 +3681,9 @@ def icdf(self, q):
 
     assert my_dist.mean.shape == (2, 3, 4)
     assert my_dist.variance.shape == (2, 3, 4)
+
+
+def test_distribution_repr():
+    result = repr(dist.Wishart(7, jnp.eye(5)).expand([3, 4]).to_event(1))
+    assert "batch shape (3,)" in result
+    assert "event shape (4, 5, 5)"