def linear_solve(a, b):
    f = (lambda y: (high_precision_dot(a, y) - b))
    x0 = jnp.zeros_like(b)
    solution = jnp.linalg.solve(a, b)
    oracle = (lambda func, x0: solution)
    return lax.custom_root(f, x0, oracle, vector_solve)
def linear_solve(a, b):
    f = (lambda y: (high_precision_dot(a, y) - b))
    x0 = jnp.zeros_like(b)
    solution = jnp.linalg.solve(a, b)
    oracle = (lambda func, x0: solution)
    return lax.custom_root(f, x0, oracle, vector_solve)