def hutchinson_diag_estimate(A: LinearOperator, k=0, bs=100, tol=3e-2, max_iters=10000, pbar=False, rand='normal',
                             key=None):
    """ Extract the (kth) diagonal of a linear operator using stochastic estimation

    Args:
        A (LinearOperator): Linear operator.
        k (int, optional): Index of the diagonal to extract (default 0).
        bs (int, optional): Chunk size (default 100).
        tol (float, optional): Tolerance (default 3e-2).
        max_iters (int, optional): Maximum number of iterations (default 10000).
        pbar (bool, optional): Flag for showing progress bar.
        key (xnp.PRNGKey, optional): Random key (default None).

    Returns:
        Array: Extracted diagonal elements.
        Info: Dictionary with information about the method used.
    """
    bs = min(100, A.shape[0])
    # lazily create chunks of the identity matrix of size bs
    xnp = A.xnp
    assert tol > 1e-3, "tolerance chosen too high for stochastic diagonal estimation"
    assert rand in ['normal', 'rademacher'], "rand must be 'normal' or 'rademacher'"
    key = xnp.PRNGKey(42) if key is None else key

    @xnp.jit
    def body(state):
        # TODO: fix randomness when using with JAX
        i, diag_sum, diag_sumsq, key = state
        key = xnp.next_key(key)
        z = xnp.randn(A.shape[0], bs, dtype=A.dtype, key=key, device=A.device)
        if rand == 'rademacher':
            z = xnp.sign(z)
        z2 = xnp.roll(z, -k, 0)
        z2 = xnp.update_array(z2, 0, slice(0, abs(k)) if k <= 0 else slice(-abs(k), None))
        slc = slice(abs(k), None) if -k > 0 else slice(None, -abs(k) or None)
        estimator = ((A @ z) * z2)[slc]
        return i + 1, diag_sum + estimator.sum(-1), diag_sumsq + (estimator**2).sum(-1), key

    def err(state):
        i, diag_sum, diag_sumsq, _ = state
        mean = diag_sum / (i * bs)
        stderr = xnp.sqrt((diag_sumsq / (i * bs) - mean**2) / (i * bs))
        return xnp.mean(stderr / xnp.maximum(xnp.abs(mean), .1 * xnp.ones_like(mean)))

    def cond(state):
        return (state[0] == 0) | ((state[0] < max_iters) & (err(state) > tol))

    while_loop, infos = xnp.while_loop_winfo(err, tol, max_iters, pbar=pbar)
    zeros = xnp.zeros((A.shape[0] - abs(k), ), dtype=A.dtype, device=A.device)
    n, diag_sum, *_ = while_loop(cond, body, (0, zeros, zeros, key))
    mean = diag_sum / (n * bs)
    return mean, infos
 Extract the (kth) diagonal of a linear operator using stochastic estimation

    Args:
        A (LinearOperator): Linear operator.
        k (int, optional): Index of the diagonal to extract (default 0).
        bs (int, optional): Chunk size (default 100).
        tol (float, optional): Tolerance (default 3e-2).
        max_iters (int, optional): Maximum number of iterations (default 10000).
        pbar (bool, optional): Flag for showing progress bar.
        key (xnp.PRNGKey, optional): Random key (default None).

    Returns:
        Array: Extracted diagonal elements.
        Info: Dictionary with information about the method used.
    
cola/linalg/trace/diagonal_estimation.py