From 1a06d50c12bed0a2d77908475fefb607dc4001c0 Mon Sep 17 00:00:00 2001
From: Ricardo Vieira <ricardo.vieira1994@gmail.com>
Date: Thu, 27 Jul 2023 13:39:50 +0200
Subject: [PATCH] Raise NotImplementedError for multivariate CustomDists

---
 pymc/distributions/distribution.py       |  5 +++++
 tests/distributions/test_distribution.py | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
index 252268a728f..8056fbd6bb3 100644
--- a/pymc/distributions/distribution.py
+++ b/pymc/distributions/distribution.py
@@ -488,6 +488,11 @@ def dist(
         class_name: str = "CustomDist",
         **kwargs,
     ):
+        if ndim_supp > 0:
+            raise NotImplementedError(
+                "CustomDist with ndim_supp > 0 and without a `dist` function are not supported."
+            )
+
         dist_params = [as_tensor_variable(param) for param in dist_params]
 
         # Assume scalar ndims_params
diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py
index 99a89dbace6..9bdfdb2053d 100644
--- a/tests/distributions/test_distribution.py
+++ b/tests/distributions/test_distribution.py
@@ -217,6 +217,10 @@ def test_custom_dist_without_random(self):
         with pytest.raises(NotImplementedError):
             pm.sample_posterior_predictive(idata, model=model)
 
+    @pytest.mark.xfail(
+        NotImplementedError,
+        reason="Support shape of multivariate CustomDist cannot be inferred. See https://github.com/pymc-devs/pytensor/pull/388",
+    )
     @pytest.mark.parametrize("size", [(), (3,), (3, 2)], ids=str)
     def test_custom_dist_with_random_multivariate(self, size):
         supp_shape = 5
@@ -264,6 +268,10 @@ def test_custom_dist_old_api_error(self):
             ):
                 CustomDist("a", lambda x: x)
 
+    @pytest.mark.xfail(
+        NotImplementedError,
+        reason="Support shape of multivariate CustomDist cannot be inferred. See https://github.com/pymc-devs/pytensor/pull/388",
+    )
     @pytest.mark.parametrize("size", [None, (), (2,)], ids=str)
     def test_custom_dist_multivariate_logp(self, size):
         supp_shape = 5
@@ -314,6 +322,10 @@ def density_moment(rv, size, mu):
         assert evaled_moment.shape == to_tuple(size)
         assert np.all(evaled_moment == mu_val)
 
+    @pytest.mark.xfail(
+        NotImplementedError,
+        reason="Support shape of multivariate CustomDist cannot be inferred. See https://github.com/pymc-devs/pytensor/pull/388",
+    )
     @pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str)
     def test_custom_dist_custom_moment_multivariate(self, size):
         def density_moment(rv, size, mu):
@@ -328,6 +340,10 @@ def density_moment(rv, size, mu):
         assert evaled_moment.shape == to_tuple(size) + (5,)
         assert np.all(evaled_moment == mu_val)
 
+    @pytest.mark.xfail(
+        NotImplementedError,
+        reason="Support shape of multivariate CustomDist cannot be inferred. See https://github.com/pymc-devs/pytensor/pull/388",
+    )
     @pytest.mark.parametrize(
         "with_random, size",
         [