def diag(A: LinearOperator, k: int = 0, alg: Algorithm = Auto()):
    r""" Extract the (kth) diagonal of a linear operator.

    Can use either the :math:`O(\tfrac{1}{\delta^2})` time stochastic estimation (alg=Hutch())
    or a deterministic :math:`O(n)` time algorithm (alg =Exact()).

    If only unbiased estimates of the diagonal are needed, use the Hutchinson algorithm.

    Args:
        A (LinearOperator): The linear operator to compute the diagonal of.
        k (int, optional): Specify to compute the kth off diagonal diagonal.
        alg (Algorithm, optional): The algorithm to use for the diagonal computation (Hutch or Exact).

    Returns:
        Array: diag
    """
 Extract the (kth) diagonal of a linear operator.

    Can use either the :math:`O(\tfrac{1}{\delta^2})` time stochastic estimation (alg=Hutch())
    or a deterministic :math:`O(n)` time algorithm (alg =Exact()).

    If only unbiased estimates of the diagonal are needed, use the Hutchinson algorithm.

    Args:
        A (LinearOperator): The linear operator to compute the diagonal of.
        k (int, optional): Specify to compute the kth off diagonal diagonal.
        alg (Algorithm, optional): The algorithm to use for the diagonal computation (Hutch or Exact).

    Returns:
        Array: diag
    
cola/linalg/trace/diag_trace.py