docs: def tmatrix_cov(C, row=None):
      r"""Covariance tensor for the non-reversible transition matrix ensemble
  
      Normally the covariance tensor cov(p_ij, p_kl) would carry four indices
      (i,j,k,l). In the non-reversible case rows are independent so that
      cov(p_ij, p_kl)=0 for i not equal to k. Therefore the function will only
      return cov(p_ij, p_ik).
  
      Parameters
      ----------
      C : (M, M) ndarray
          Count matrix
      row : int (optional)
          If row is given return covariance matrix for specified row only
  
      Returns
      -------
      cov : (M, M, M) ndarray
          Covariance tensor
  
      """
  
      if row is None:
          alpha = C + 1.0  # Dirichlet parameters
          alpha0 = alpha.sum(axis=1)  # Sum of paramters (per row)
  
          norm = alpha0 ** 2 * (alpha0 + 1.0)
  
          """Non-normalized covariance tensor"""
          Z = -alpha[:, :, np.newaxis] * alpha[:, np.newaxis, :]
  
          """Correct-diagonal"""
          ind = np.diag_indices(C.shape[0])
          Z[:, ind[0], ind[1]] += alpha0[:, np.newaxis] * alpha
  
          """Covariance matrix"""
          cov = Z / norm[:, np.newaxis, np.newaxis]
  
          return cov
  
      else:
          alpha = C[row, :] + 1.0
          return dirichlet_covariance(alpha)