Update MatrixNormal shape_inference and remove shape and size restrictions

Author: ricardoV94

pymc/distributions/multivariate.py

@@ -1616,7 +1616,7 @@ class MatrixNormalRV(RandomVariable):
     _print_name = ("MatrixNormal", "\\operatorname{MatrixNormal}")
 
     def _infer_shape(self, size, dist_params, param_shapes=None):
-        shape = tuple(size) + tuple(dist_params[0].shape)
+        shape = tuple(size) + tuple(dist_params[0].shape[-2:])
         return shape
 
     @classmethod
@@ -1746,18 +1746,6 @@ def dist(
 
         cholesky = Cholesky(lower=True, on_error="raise")
 
-        if kwargs.get("size", None) is not None:
-            raise NotImplementedError("MatrixNormal doesn't support size argument")
-
-        if "shape" in kwargs:
-            kwargs.pop("shape")
-            warnings.warn(
-                "The shape argument in MatrixNormal is deprecated and will be ignored."
-                "MatrixNormal automatically derives the shape"
-                "from row and column matrix dimensions.",
-                FutureWarning,
-            )
-
         # Among-row matrices
         if len([i for i in [rowcov, rowchol] if i is not None]) != 1:
             raise ValueError(
@@ -1787,22 +1775,16 @@ def dist(
                 raise ValueError("colchol must be two dimensional.")
             colchol_cov = at.as_tensor_variable(colchol)
 
-        dist_shape = (rowchol_cov.shape[0], colchol_cov.shape[0])
+        dist_shape = (rowchol_cov.shape[-1], colchol_cov.shape[-1])
 
         # Broadcasting mu
         mu = at.extra_ops.broadcast_to(mu, shape=dist_shape)
-
         mu = at.as_tensor_variable(floatX(mu))
-        # mean = median = mode = mu
 
         return super().dist([mu, rowchol_cov, colchol_cov], **kwargs)
 
     def get_moment(rv, size, mu, rowchol, colchol):
-        output_shape = (rowchol.shape[0], colchol.shape[0])
-        if not rv_size_is_none(size):
-            output_shape = at.concatenate([size, output_shape])
-        moment = at.full(output_shape, mu)
-        return moment
+        return at.full_like(rv, mu)
 
     def logp(value, mu, rowchol, colchol):
         """

pymc/tests/test_distributions_moments.py

@@ -1020,11 +1020,6 @@ def test_mvstudentt_moment(nu, mu, cov, size, expected):
     assert_moment_is_expected(model, expected, check_finite_logp=x.ndim < 3)
 
 
-def check_matrixnormal_moment(mu, rowchol, colchol, size, expected):
-    with Model() as model:
-        MatrixNormal("x", mu=mu, rowchol=rowchol, colchol=colchol, size=size)
-
-
 @pytest.mark.parametrize(
     "alpha, mu, sigma, size, expected",
     [
@@ -1092,11 +1087,12 @@ def test_asymmetriclaplace_moment(b, kappa, mu, size, expected):
     ],
 )
 def test_matrixnormal_moment(mu, rowchol, colchol, size, expected):
-    if size is None:
-        check_matrixnormal_moment(mu, rowchol, colchol, size, expected)
-    else:
-        with pytest.raises(NotImplementedError):
-            check_matrixnormal_moment(mu, rowchol, colchol, size, expected)
+    with Model() as model:
+        x = MatrixNormal("x", mu=mu, rowchol=rowchol, colchol=colchol, size=size)
+
+    # MatrixNormal logp is only implemented for 2d values
+    check_logp = x.ndim == 2
+    assert_moment_is_expected(model, expected, check_finite_logp=check_logp)
 
 
 @pytest.mark.parametrize(

pymc/tests/test_distributions_random.py

@@ -1811,13 +1811,15 @@ class TestMatrixNormal(BaseTestDistributionRandom):
     mu = np.random.random((3, 3))
     row_cov = np.eye(3)
     col_cov = np.eye(3)
-    shape = None
-    size = None
     pymc_dist_params = {"mu": mu, "rowcov": row_cov, "colcov": col_cov}
     expected_rv_op_params = {"mu": mu, "rowcov": row_cov, "colcov": col_cov}
 
+    sizes_to_check = (None, (1,), (2, 4))
+    sizes_expected = [(3, 3), (1, 3, 3), (2, 4, 3, 3)]
+
     checks_to_run = [
         "check_pymc_params_match_rv_op",
+        "check_rv_size",
         "check_draws",
         "check_errors",
         "check_random_variable_prior",
@@ -1858,17 +1860,6 @@ def ref_rand(mu, rowcov, colcov):
         assert p > delta
 
     def check_errors(self):
-        msg = "MatrixNormal doesn't support size argument"
-        with pm.Model():
-            with pytest.raises(NotImplementedError, match=msg):
-                matrixnormal = pm.MatrixNormal(
-                    "matnormal",
-                    mu=np.random.random((3, 3)),
-                    rowcov=np.eye(3),
-                    colcov=np.eye(3),
-                    size=15,
-                )
-
         with pm.Model():
             matrixnormal = pm.MatrixNormal(
                 "matnormal",
@@ -1879,16 +1870,6 @@ def check_errors(self):
             with pytest.raises(ValueError):
                 logp(matrixnormal, aesara.tensor.ones((3, 3, 3)))
 
-        with pm.Model():
-            with pytest.warns(FutureWarning):
-                matrixnormal = pm.MatrixNormal(
-                    "matnormal",
-                    mu=np.random.random((3, 3)),
-                    rowcov=np.eye(3),
-                    colcov=np.eye(3),
-                    shape=15,
-                )
-
     def check_random_variable_prior(self):
         """
         This test checks for shape correctness when using MatrixNormal distribution