fix str repr for KroneckerNormal and MatrixNormal

Author: Spaak

pymc3/distributions/multivariate.py

@@ -1449,6 +1449,8 @@ def logp(self, x):
             broadcast_conditions=False,
         )
 
+    def _distr_parameters_for_repr(self):
+        return ["eta", "n"]
 
 class MatrixNormal(Continuous):
     R"""
@@ -1712,6 +1714,10 @@ def logp(self, value):
         norm = -0.5 * m * n * pm.floatX(np.log(2 * np.pi))
         return norm - 0.5 * trquaddist - m * half_collogdet - n * half_rowlogdet
 
+    def _distr_parameters_for_repr(self):
+        mapping = {"tau": "tau", "cov": "cov", "chol": "chol_cov"}
+        return ["mu", "row" + mapping[self._rowcov_type], "col" + mapping[self._colcov_type]]
+
 
 class KroneckerNormal(Continuous):
     R"""
@@ -1954,3 +1960,6 @@ def logp(self, value):
         """
         quad, logdet = self._quaddist(value)
         return -(quad + logdet + self.N * tt.log(2 * np.pi)) / 2.0
+
+    def _distr_parameters_for_repr(self):
+        return ["mu"]

pymc3/tests/test_distributions.py

@@ -1782,8 +1782,18 @@ def setup_class(self):
             # add a bounded variable as well
             bound_var = Bound(Normal, lower=1.0)("bound_var", mu=0, sigma=10)
 
+            # KroneckerNormal
+            n, m = 3, 4
+            covs = [np.eye(n), np.eye(m)]
+            kron_normal = KroneckerNormal('kron_normal', mu=np.zeros(n*m), covs=covs, shape=n*m)
+
+            # MatrixNormal
+            matrix_normal = MatrixNormal('mat_normal', mu=np.random.normal(size=n), rowcov=np.eye(n),
+                colchol=np.linalg.cholesky(np.eye(n)), shape=(n, n))
+
             # Likelihood (sampling distribution) of observations
             Y_obs = Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)
+
         self.distributions = [alpha, sigma, mu, b, Z, Y_obs, bound_var]
         self.expected_latex = (
             r"$\text{alpha} \sim \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
@@ -1793,6 +1803,8 @@ def setup_class(self):
             r"$\text{Z} \sim \text{MvNormal}(\mathit{mu}=array,~\mathit{chol_cov}=array)$",
             r"$\text{Y_obs} \sim \text{Normal}(\mathit{mu}=\text{mu},~\mathit{sigma}=f(\text{sigma}))$",
             r"$\text{bound_var} \sim \text{Bound}(\mathit{lower}=1.0,~\mathit{upper}=\text{None})$ -- \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
+            r"$\text{kron_normal} \sim \text{KroneckerNormal}(\mathit{mu}=array)$",
+            r"$\text{mat_normal} \sim \text{MatrixNormal}(\mathit{mu}=array,~\mathit{rowcov}=array,~\mathit{colchol_cov}=array)$",
         )
         self.expected_str = (
             r"alpha ~ Normal(mu=0.0, sigma=10.0)",
@@ -1802,6 +1814,8 @@ def setup_class(self):
             r"Z ~ MvNormal(mu=array, chol_cov=array)",
             r"Y_obs ~ Normal(mu=mu, sigma=f(sigma))",
             r"bound_var ~ Bound(lower=1.0, upper=None)-Normal(mu=0.0, sigma=10.0)",
+            r"kron_normal ~ KroneckerNormal(mu=array)",
+            r"mat_normal ~ MatrixNormal(mu=array, rowcov=array, colchol_cov=array)",
         )
 
     def test__repr_latex_(self):