docs: def base_conditional(Kmn, Kmm, Knn, f, *, full_cov=False, q_sqrt=None, white=False):
      r"""
      Given a g1 and g2, and distribution p and q such that
        p(g2) = N(g2;0,Kmm)
        p(g1) = N(g1;0,Knn)
        p(g1|g2) = N(g1;0,Knm)
      And
        q(g2) = N(g2;f,q_sqrt*q_sqrt^T)
      This method computes the mean and (co)variance of
        q(g1) = \int q(g2) p(g1|g2)
      :param Kmn: M x [...] x N
      :param Kmm: M x M
      :param Knn: [...] x N x N  or  N
      :param f: M x R
      :param full_cov: bool
      :param q_sqrt: None or R x M x M (lower triangular)
      :param white: bool
      :return: N x R  or R x N x N
      """
      logger.debug("base conditional")
      # compute kernel stuff
      num_func = tf.shape(f)[-1]  # R
      N = tf.shape(Kmn)[-1]
      M = tf.shape(f)[-2]
  
      # get the leadings dims in Kmn to the front of the tensor
      # if Kmn has rank two, i.e. [M, N], this is the identity op.
      K = tf.rank(Kmn)
      perm = tf.concat([tf.reshape(tf.range(1, K-1), [K-2]), # leading dims (...)
                        tf.reshape(0, [1]),  # [M]
                        tf.reshape(K-1, [1])], 0)  # [N]
      Kmn = tf.transpose(Kmn, perm)  # ... x M x N
  
      leading_dims = tf.shape(Kmn)[:-2]
      Lm = tf.cholesky(Kmm)  # [M,M]
  
      # Compute the projection matrix A
      Lm = tf.broadcast_to(Lm, tf.concat([leading_dims, tf.shape(Lm)], 0))  # [...,M,M]
      A = tf.matrix_triangular_solve(Lm, Kmn, lower=True)  # [...,M,N]
      # compute the covariance due to the conditioning
      if full_cov:
          fvar = Knn - tf.matmul(A, A, transpose_a=True)  # [...,N,N]
          cov_shape = tf.concat([leading_dims, [num_func, N, N]], 0)
          fvar = tf.broadcast_to(tf.expand_dims(fvar, -3), cov_shape)  # [...,R,N,N]
      else:
          fvar = Knn - tf.reduce_sum(tf.square(A), -2)  # [...,N]
          cov_shape = tf.concat([leading_dims, [num_func, N]], 0) # [...,R,N]
          fvar = tf.broadcast_to(tf.expand_dims(fvar, -2), cov_shape)  # [...,R,N]
  
      # another backsubstitution in the unwhitened case
      if not white:
          A = tf.matrix_triangular_solve(tf.matrix_transpose(Lm), A, lower=False)
  
      # construct the conditional mean
      f_shape = tf.concat([leading_dims, [M, num_func]], 0)  # [...,M,R]
      f = tf.broadcast_to(f, f_shape)  # [...,M,R]
      fmean = tf.matmul(A, f, transpose_a=True)  # [...,N,R]
  
      if q_sqrt is not None:
          if q_sqrt.get_shape().ndims == 2:
              LTA = A * tf.expand_dims(tf.transpose(q_sqrt), 2)  # R x M x N
          elif q_sqrt.get_shape().ndims == 3:
              L = q_sqrt
              L = tf.broadcast_to(L, tf.concat([leading_dims, tf.shape(L)], 0))
  
              shape = tf.concat([leading_dims, [num_func, M, N]], 0)
              A_tiled = tf.broadcast_to(tf.expand_dims(A, -3), shape)
              LTA = tf.matmul(L, A_tiled, transpose_a=True)  # R x M x N
          else:  # pragma: no cover
              raise ValueError("Bad dimension for q_sqrt: %s" %
                               str(q_sqrt.get_shape().ndims))
          if full_cov:
              fvar = fvar + tf.matmul(LTA, LTA, transpose_a=True)  # R x N x N
          else:
              fvar = fvar + tf.reduce_sum(tf.square(LTA), -2)  # R x N
  
      if not full_cov:
          fvar = tf.matrix_transpose(fvar)  # N x R
  
      return fmean, fvar