Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: RBF kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other.

    :return: N x dim(Z1) x dim(Z2)

def _expectation(p, kern1, feat1, kern2, feat2, nghp=None):
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: RBF kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other.

    :return: N x dim(Z1) x dim(Z2)
    """

    if kern1.on_separate_dims(kern2) and isinstance(p, DiagonalGaussian):
        # no joint expectations required
        eKxz1 = expectation(p, (kern1, feat1))
        eKxz2 = expectation(p, (kern2, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    Ka, Kb = kern1, kern2

    with params_as_tensors_for(Ka, feat1, Kb, feat2):

        # use only active dimensions
        Xcov = Ka._slice_cov(tf.matrix_diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)

        Z1, Xmu = Ka._slice(feat1.Z, p.mu)

        N = tf.shape(Xmu)[0]
        D = tf.shape(Xmu)[1]

        def get_squared_length_scales(kern):
            squared_lengthscales = kern.lengthscales ** 2. if kern.ARD \
                        else tf.zeros((D,), dtype=settings.float_type) + kern.lengthscales ** 2.
            return squared_lengthscales

        if Ka == Kb:
            La = get_squared_length_scales(Ka)
            Lb = La
            half_mean_L = La * 0.5  # average length scale
        else:
            La, Lb = map(get_squared_length_scales, (Ka, Kb))
            half_mean_L = La * Lb / (La + Lb)  # average length scale

        sqrt_det_L = tf.reduce_prod(half_mean_L) ** 0.5
        C = tf.cholesky(tf.matrix_diag(half_mean_L) + Xcov)  # [N, D, D]
        dets = sqrt_det_L / tf.exp(tf.reduce_sum(tf.log(tf.matrix_diag_part(C)), axis=1))  # N

        # for mahalanobis computation we need Zᵀ (CCᵀ)⁻¹ Z  as well as C⁻¹ Z
        # with Z = Z₁, Z₂  for two rbf kernels
        def get_cholesky_solve_terms(Z, C=C):
            C_inv_z = tf.matrix_triangular_solve(
                C, tf.tile(tf.expand_dims(tf.transpose(Z), 0),
                           [N, 1, 1]), lower=True)  # [N, D, M]
            z_CC_inv_z = tf.reduce_sum(tf.square(C_inv_z), 1)  # [N, M]

            return C_inv_z, z_CC_inv_z

        C_inv_mu = tf.matrix_triangular_solve(C, tf.expand_dims(Xmu, 2), lower=True)  # [N, D, 1]
        mu_CC_inv_mu = tf.expand_dims(tf.reduce_sum(tf.square(C_inv_mu), 1), 2)  # [N, 1, 1]

        C_inv_z1, z1_CC_inv_z1 = get_cholesky_solve_terms(Z1 / La * half_mean_L)
        z1_CC_inv_mu = 2 * tf.matmul(C_inv_z1, C_inv_mu, transpose_a=True)[:, :, 0]  # [N, M1]

        if feat1 == feat2 and Ka == Kb:
            # in this case Z2==Z1 so we can reuse the Z1 terms
            C_inv_z2, z2_CC_inv_z2  = C_inv_z1, z1_CC_inv_z1
            z2_CC_inv_mu = z1_CC_inv_mu  # [N, M]
            Z2 = Z1
        else:
            # compute terms related to Z2
            Z2, _ = Kb._slice(feat2.Z, p.mu)
            C_inv_z2, z2_CC_inv_z2 = get_cholesky_solve_terms(Z2 / Lb * half_mean_L)
            z2_CC_inv_mu = 2 * tf.matmul(C_inv_z2, C_inv_mu, transpose_a=True)[:, :, 0]  # [N, M2]

        z1_CC_inv_z2 = tf.matmul(C_inv_z1, C_inv_z2, transpose_a=True)  # [N, M1, M2]

        # expand dims for broadcasting
        # along M1
        z2_CC_inv_mu = tf.expand_dims(z2_CC_inv_mu, 1)  # [N, 1, M2]
        z2_CC_inv_z2 = tf.expand_dims(z2_CC_inv_z2, 1)

        # along M2
        z1_CC_inv_mu = tf.expand_dims(z1_CC_inv_mu, 2)  # [N, M1, 1]
        z1_CC_inv_z1 = tf.expand_dims(z1_CC_inv_z1, 2)

        # expanded version of ((Z1 + Z2)-mu) (CCT)-1 ((Z1 + Z2)-mu)
        mahalanobis = mu_CC_inv_mu + z2_CC_inv_z2 + \
                      z1_CC_inv_z1 + 2 * z1_CC_inv_z2 - \
                      z1_CC_inv_mu - z2_CC_inv_mu  # [N, M1, M2]

        exp_mahalanobis = tf.exp(-0.5 * mahalanobis)  # [N, M1, M2]

        if Z1 == Z2:
            # CAVEAT : Compute sqrt(self.K(Z)) explicitly
            # to prevent automatic gradient from
            # being NaN sometimes, see pull request #615
            sqrt_exp_dist = tf.exp(-0.25 * Ka.scaled_square_dist(Z1, None))
        else:
            # Compute exp( -.5 (Z-Z')^top (L_1+L_2)^{-1} (Z-Z') )
            lengthscales_rms = tf.sqrt(La + Lb)
            Z1 = Z1 / lengthscales_rms
            Z1sqr = tf.reduce_sum(tf.square(Z1), axis=1)
            Z2 = Z2 / lengthscales_rms
            Z2sqr = tf.reduce_sum(tf.square(Z2), axis=1)
            dist = -2 * tf.matmul(Z1, Z2, transpose_b=True) \
                   + tf.reshape(Z1sqr, (-1, 1)) + tf.reshape(Z2sqr, (1, -1))
            sqrt_exp_dist = tf.exp(-0.5 * dist)  # M1 x M2

        return Ka.variance * Kb.variance * sqrt_exp_dist * \
               tf.reshape(dets, [N, 1, 1]) * exp_mahalanobis