Multioutput gp.Latent

pymc/gp/gp.py

@@ -408,7 +408,9 @@ def _build_marginal_likelihood(self, X, noise, jitter):
         cov = Kxx + Knx
         return mu, stabilize(cov, jitter)
 
-    def marginal_likelihood(self, name, X, y, noise, jitter=0.0, is_observed=True, **kwargs):
+    def marginal_likelihood(
+        self, name, X, y, noise, n_outputs=None, jitter=0.0, is_observed=True, **kwargs
+    ):
         R"""
         Returns the marginal likelihood distribution, given the input
         locations `X` and the data `y`.
@@ -428,10 +430,14 @@ def marginal_likelihood(self, name, X, y, noise, jitter=0.0, is_observed=True, *
             vector with shape `(n, 1)`.
         y: array-like
             Data that is the sum of the function with the GP prior and Gaussian
-            noise.  Must have shape `(n, )`.
+            noise.  Must have shape `(n, n_outputs)`.
         noise: scalar, Variable, or Covariance
             Standard deviation of the Gaussian noise.  Can also be a Covariance for
             non-white noise.
+        n_outputs: int
+            The number of observations, or outputs of the GP.  If you have a matrix
+            instead of vector `y`, this is the number of columns.  This argument only
+            needs to be supplied if `y` is a tensor whose shape isn't known at runtime.
         jitter: scalar
             A small correction added to the diagonal of positive semi-definite
             covariance matrices to ensure numerical stability.  Default value is 0.0.
@@ -446,6 +452,7 @@ def marginal_likelihood(self, name, X, y, noise, jitter=0.0, is_observed=True, *
         self.X = X
         self.y = y
         self.noise = noise
+        self.n_outputs = shape(y, kwargs.pop("n_outputs", None))
         if is_observed:
             return pm.MvNormal(name, mu=mu, cov=cov, observed=y, **kwargs)
         else:
@@ -483,8 +490,8 @@ def _build_conditional(
         rxx = y - mean_total(X)
         L = cholesky(stabilize(Kxx, jitter) + Knx)
         A = solve_lower(L, Kxs)
-        v = solve_lower(L, rxx)
-        mu = self.mean_func(Xnew) + at.dot(at.transpose(A), v)
+        v = solve_lower(L, rxx.T)
+        mu = self.mean_func(Xnew) + at.dot(A.T, v).T
         if diag:
             Kss = self.cov_func(Xnew, diag=True)
             var = Kss - at.sum(at.square(A), 0)
@@ -493,12 +500,14 @@ def _build_conditional(
             return mu, var
         else:
             Kss = self.cov_func(Xnew)
-            cov = Kss - at.dot(at.transpose(A), A)
+            cov = Kss - at.dot(A.T, A)
             if pred_noise:
                 cov += noise(Xnew)
             return mu, cov if pred_noise else stabilize(cov, jitter)
 
-    def conditional(self, name, Xnew, pred_noise=False, given=None, jitter=0.0, **kwargs):
+    def conditional(
+        self, name, Xnew, pred_noise=False, given=None, n_outputs=None, jitter=0.0, **kwargs
+    ):
         R"""
         Returns the conditional distribution evaluated over new input
         locations `Xnew`.
@@ -526,6 +535,10 @@ def conditional(self, name, Xnew, pred_noise=False, given=None, jitter=0.0, **kw
             Can optionally take as key value pairs: `X`, `y`, `noise`,
             and `gp`.  See the section in the documentation on additive GP
             models in PyMC for more information.
+        n_outputs: int
+            The number of observations, or outputs of the GP.  If you have a matrix
+            instead of vector `y`, this is the number of columns.  This argument only
+            needs to be supplied if `y` is a tensor whose shape isn't known at runtime.
         jitter: scalar
             A small correction added to the diagonal of positive semi-definite
             covariance matrices to ensure numerical stability.  For conditionals
@@ -537,7 +550,19 @@ def conditional(self, name, Xnew, pred_noise=False, given=None, jitter=0.0, **kw
 
         givens = self._get_given_vals(given)
         mu, cov = self._build_conditional(Xnew, pred_noise, False, *givens, jitter)
-        return pm.MvNormal(name, mu=mu, cov=cov, **kwargs)
+
+        try:
+            shape = infer_shape(Xnew, kwargs.pop("shape", None))
+            return pm.MvNormal(name, mu=mu, cov=cov, shape=shape, **kwargs)
+        except TypeError as e:
+            n_points = infer_shape(Xnew, kwargs.pop("shape", None))
+            return pm.MatrixNormal(
+                name,
+                mu=mu,
+                rowcov=at.eye(self.n_outputs),
+                colcov=cov,
+                shape=(self.n_outputs, n_points),
+            )
 
     def predict(
         self, Xnew, point=None, diag=False, pred_noise=False, given=None, jitter=0.0, model=None

pymc/gp/util.py

@@ -105,9 +105,9 @@ def infer_size(X, n_points=None):
     """
     if n_points is None:
         try:
-            n_points = int(X.shape[0])
-        except TypeError:
-            raise TypeError("Cannot infer 'shape', provide as an argument")
+            n_points = np.int(X.shape[0])
+        except TypeError as e:
+            raise TypeError("Cannot infer 'shape', it must be provided as an argument") from e
     return n_points