docs: def _mix_latent_gp(W, g_mu, g_var, full_cov, full_output_cov):
      r"""
      Takes the mean and variance of an uncorrelated L-dimensional latent GP
      and returns the mean and the variance of the mixed GP, `f = W g`,
      where both f and g are GPs, with W having a shape [P, L]
  
      :param W: [P, L]
      :param g_mu: [..., N, L]
      :param g_var: [..., N, L] (full_cov = False) or [L, ..., N, N] (full_cov = True)
      :return: f_mu and f_var, shape depends on `full_cov` and `full_output_cov`
      """
      f_mu = tf.tensordot(g_mu, W, [[-1], [-1]])  # [..., N, P]
  
      rk = tf.rank(g_var)
      leading_dims = (rk - 3) if full_cov else (rk - 2)
  
      if full_cov and full_output_cov:  # g_var is [L, ..., N, N]
          # this branch is practically never taken
          g_var = _rollaxis_left(g_var, 1)  # [..., N, N, L]
          g_var = tf.expand_dims(g_var, axis=-2)  # [..., N, N, 1, L]
          g_var_W = g_var * W  # [..., N, P, L]
          f_var = tf.tensordot(g_var_W, W, [[-1], [-1]])  # [..., N, N, P, P]
          f_var = misc.leading_transpose(f_var, [..., -4, -2, -3, -1])  # [..., N, P, N, P]
  
      elif full_cov and not full_output_cov:  # g_var is [L, ..., N, N]
          # this branch is practically never taken
          f_var = tf.tensordot(g_var, W**2, [[0], [-1]])  # [..., N, N, P]
          f_var = misc.leading_transpose(f_var, [..., -1, -3, -2])  # [..., P, N, N]
  
      elif not full_cov and full_output_cov:  # g_var is [..., N, L]
          g_var = tf.expand_dims(g_var, axis=-2)  # [..., N, 1, L]
          g_var_W = g_var * W  # [..., N, P, L]
          f_var = tf.tensordot(g_var_W, W, [[-1], [-1]])  # [..., N, P, P]
  
      elif not full_cov and not full_output_cov:  # g_var is [..., N, L]
          W_squared = W**2  # [P, L]
          f_var = tf.tensordot(g_var, W_squared, [[-1], [-1]])  # [..., N, P]
  
      return f_mu, f_var