Swap expand_dims of Periodic and WrappedPeriodic kernel full method

Author: lucianopaz

pymc/gp/cov.py

@@ -779,8 +779,8 @@ def full(self, X: TensorLike, Xs: Optional[TensorLike] = None) -> TensorVariable
         X, Xs = self._slice(X, Xs)
         if Xs is None:
             Xs = X
-        f1 = pt.expand_dims(X, axis=(0,))
-        f2 = pt.expand_dims(Xs, axis=(1,))
+        f1 = pt.expand_dims(X, axis=(1,))
+        f2 = pt.expand_dims(Xs, axis=(0,))
         r = np.pi * (f1 - f2) / self.period
         r2 = pt.sum(pt.square(pt.sin(r) / self.ls), 2)
         return self.full_from_distance(r2, squared=True)
@@ -946,8 +946,8 @@ def full(self, X: TensorLike, Xs: Optional[TensorLike] = None) -> TensorVariable
         X, Xs = self._slice(X, Xs)
         if Xs is None:
             Xs = X
-        f1 = pt.expand_dims(X, axis=(0,))
-        f2 = pt.expand_dims(Xs, axis=(1,))
+        f1 = pt.expand_dims(X, axis=(1,))
+        f2 = pt.expand_dims(Xs, axis=(0,))
         r = np.pi * (f1 - f2) / self.period
         r2 = pt.sum(pt.square(pt.sin(r) / self.cov_func.ls), 2)
         return self.cov_func.full_from_distance(r2, squared=True)

tests/gp/test_cov.py

@@ -872,3 +872,29 @@ def test_raises3(self):
         with pm.Model() as model:
             with pytest.raises(ValueError):
                 B = pm.gp.cov.Coregion(1)
+
+
+@pytest.mark.parametrize(
+    ["kernel", "args"],
+    [
+        ["Constant", (1.0,)],
+        ["WhiteNoise", (1.0,)],
+        ["ExpQuad", (1, 1.0)],
+        ["RatQuad", (1, 1.0, 1.0)],
+        ["Exponential", (1, 1.0)],
+        ["Matern12", (1, 1.0)],
+        ["Matern32", (1, 1.0)],
+        ["Matern52", (1, 1.0)],
+        ["Periodic", (1, 1.0, 1.0)],
+        ["Circular", (1, 1.0)],
+        ["Linear", (1, 1.0)],
+        ["Cosine", (1, 1.0)],
+        ["Polynomial", (1, 1.0, 1.0, 1.0)],
+        ["WrappedPeriodic", (pm.gp.cov.ExpQuad(1, 1.0), 1.0)],
+        ["Gibbs", (1, lambda x: pt.ones(x.shape))],
+    ],
+)
+def test_full_shape(kernel, args):
+    X = np.arange(10)[:, None]
+    Xs = np.arange(5)[:, None]
+    assert tuple(getattr(pm.gp.cov, kernel)(*args).full(X, Xs).shape.eval()) == (len(X), len(Xs))