def lanczos(A: LinearOperator, start_vector: Array = None, max_iters=100, tol=1e-7, pbar=False, key=None):
    """
    Computes the Lanczos decomposition of a the operator A, A = Q T Q^*.

    Args:
        A (LinearOperator): A symmetric linear operator of size (n, n).
        start_vector (Array, optional): An initial vector to start Lanczos of size (n, ) or (b, n)
         Defaults to a random probe vector (n, ).
        max_iters (int, optional): The maximum number of iterations to run.
        tol (float, optional): Stopping criteria.
        key (PNRGKey, optional): PRNGKey for random number generation.

    Returns:
        tuple:
            - Q (Array): Orthogonal matrix of size (n, max_iters) or (b, n, max_iters) if start vector has a batch dim
            - T (Array): Tridiagonal matrix of size (max_iters, max_iters) or (b, max_iters, max_iters) if batch dim
            - info (dict): General information about the iterative procedure.

    Notes:
        On jax backend, T will be padded with 0s to size of max_iters since jax does not
         support dynamic sized outputs if traced.

    https://en.wikipedia.org/wiki/Lanczos_algorithm
    """
    xnp = A.xnp
    max_iters = min(max_iters, A.shape[0])
    if start_vector is None:
        key = xnp.PRNGKey(42) if key is None else key
        start_vector = xnp.randn(A.shape[0], dtype=A.dtype, device=A.device, key=key)

    if len(start_vector.shape) == 1:
        rhs = start_vector[:, None]
    else:
        rhs = start_vector

    init_val = init_lanczos(xnp, rhs, max_iters=max_iters, dtype=A.dtype)
    vec, beta, alpha, i, info = lanczos_fact(A, init_val, max_iters, tol)
    alpha, beta, Q, iters = alpha[..., 1:-1], beta, vec[..., 1:-1], i - 1
    if xnp.__name__.find("jax") < 0:
        alpha, beta, Q = alpha[..., :iters - 1], beta[..., :iters], Q[..., :iters]
    if len(start_vector.shape) == 1:
        alpha, beta = alpha[0], beta[0]
        T = Tridiagonal(alpha, beta, alpha)
        return Unitary(Dense(Q[0, :, :])), T, info
    else:
        T = xnp.vmap(Tridiagonal)(alpha, beta, alpha)
        Q = Unitary(xnp.vmap(Dense)(Q))
        return Q, T, info

    Computes the Lanczos decomposition of a the operator A, A = Q T Q^*.

    Args:
        A (LinearOperator): A symmetric linear operator of size (n, n).
        start_vector (Array, optional): An initial vector to start Lanczos of size (n, ) or (b, n)
         Defaults to a random probe vector (n, ).
        max_iters (int, optional): The maximum number of iterations to run.
        tol (float, optional): Stopping criteria.
        key (PNRGKey, optional): PRNGKey for random number generation.

    Returns:
        tuple:
            - Q (Array): Orthogonal matrix of size (n, max_iters) or (b, n, max_iters) if start vector has a batch dim
            - T (Array): Tridiagonal matrix of size (max_iters, max_iters) or (b, max_iters, max_iters) if batch dim
            - info (dict): General information about the iterative procedure.

    Notes:
        On jax backend, T will be padded with 0s to size of max_iters since jax does not
         support dynamic sized outputs if traced.

    https://en.wikipedia.org/wiki/Lanczos_algorithm
    
cola/linalg/decompositions/lanczos.py