docs: def _expectation(p, kern, none1, none2, none3, nghp=None):
      """
      Compute the expectation:
      <diag(K_{X, X})>_p(X)
          - K_{.,.} :: Linear kernel
  
      :return: N
      """
      with params_as_tensors_for(kern):
          # use only active dimensions
          Xmu, _ = kern._slice(p.mu, None)
          Xcov = kern._slice_cov(p.cov)
  
          return tf.reduce_sum(kern.variance * (tf.matrix_diag_part(Xcov) + tf.square(Xmu)), 1)