docs: 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