//  [python]
def predict(self,
                x: Tensor,
                full_cov: bool,
                sample_idxs=None) -> Tuple[Tensor, Tensor]:
        """
        Parameters
        ----------
        x : Tensor (single kernel) or Tensor list (product kernels)
            test input tensor(s) with dimensions (n_b x n_samples x d x m)
        full_cov : bool
            returns full covariance if true otherwise returns the diagonal

        Returns
        -------
        mu : Tensor 
            mean of predictive density at test inputs [ s ]
        v : Tensor 
            variance/covariance of predictive density at test inputs [ s ]
            if full_cov is true returns full covariance, otherwise
            returns diagonal variance

        Notes
        -----
        """

        q_mu, q_sqrt, z = self.prms
        kernel = self.kernel
        q_mu = q_mu[..., None]

        assert (q_mu.shape[0] == q_sqrt.shape[0])
        if (not self.tied_samples) and sample_idxs is not None:
            q_mu = q_mu[sample_idxs]
            q_sqrt = q_sqrt[sample_idxs]

        # see ELBO for explanation of _expand
        z = self._expand_z(z)
        x = self._expand_x(x)
        kzz = kernel(z, z)  # dims: (1 x n x n_z x n_z)
        kzx = kernel(z, x)  # dims: (n_mc x n_samples x n x n_inducing x m)
        e = torch.eye(self.n_inducing,
                      dtype=torch.get_default_dtype()).to(kzz.device)

        # [ l ] has dims: (1 x n x n_inducing x n_inducing)
        l = torch.linalg.cholesky(kzz + (jitter * e))
        # [ alpha ] has dims: (n_b x n_samples x n x n_inducing x m)
        alpha = torch.triangular_solve(kzx, l, upper=False)[0]
        alphat = alpha.transpose(-1, -2)

        if self.whiten:
            # [ mu ] has dims : (n_b x n_samples x n x m x 1)
            mu = torch.matmul(alphat, q_mu)
        else:
            # [ beta ] has dims : (n_b x n_samples x n x n_inducing x m)
            beta = torch.triangular_solve(alpha,
                                          l.transpose(-1, -2),
                                          upper=True)[0]
            # [ betat ] has dims : (n_b x n_samples x n x m x n_inducing)
            betat = beta.transpose(-1, -2)
            mu = torch.matmul(betat, q_mu)

        if full_cov:
            # [ tmp1 ] has dims : (n_b x n_samples, n x m x n_inducing)
            if self.whiten:
                tmp1 = torch.matmul(alphat, q_sqrt)
            else:
                tmp1 = torch.matmul(betat, q_sqrt)
            # [ v1 ] has dims : (n_b x n_samples x n x m x m)
            v1 = torch.matmul(tmp1, tmp1.transpose(-1, -2))
            # [ v2 ] has dims : (n_b x n_samples x n x m x m)
            v2 = torch.matmul(alphat, alphat)
            # [ kxx ] has dims : (n_b x n_samples x n x m x m)
            kxx = kernel(x, x)
            v = kxx + v1 - v2
        else:
            # [ kxx ] has dims : (n_b x n_samples x n x m)
            kxx = kernel.diagK(x)
            # [ tmp1 ] has dims : (n_b x n_samples x n x m x n_inducing)
            if self.whiten:
                tmp1 = torch.matmul(alphat, q_sqrt)
            else:
                tmp1 = torch.matmul(betat, q_sqrt)
            # [ v1 ] has dims : (n_b x n_samples x n x m)
            v1 = torch.square(tmp1).sum(-1)
            # [ v2 ] has dims : (n_b x n_samples x n x m)
            v2 = torch.square(alpha).sum(-2)
            v = kxx + v1 - v2

        return mu.squeeze(-1), v