def richardson_iteration(matvec, b, omega=0.1, tolerance=1e-06):

    def cond(x):
        return (jnp.linalg.norm((matvec(x) - b)) > tolerance)

    def body(x):
        return (x + (omega * (b - matvec(x))))
    return lax.while_loop(cond, body, b)
def richardson_iteration(matvec, b, omega=0.1, tolerance=1e-06):

    def cond(x):
        return (jnp.linalg.norm((matvec(x) - b)) > tolerance)

    def body(x):
        return (x + (omega * (b - matvec(x))))
    return lax.while_loop(cond, body, b)