def linear_solve(a, b):
    f = (lambda x: (high_precision_dot(a, x) - b))
    factors = jsp.linalg.cho_factor(a)
    cho_solve = (lambda f, b: jsp.linalg.cho_solve(factors, b))

    def pos_def_solve(g, b):
        return lax.custom_linear_solve(g, b, cho_solve, symmetric=True)
    return lax.custom_root(f, b, cho_solve, pos_def_solve)
def linear_solve(a, b):
    f = (lambda x: (high_precision_dot(a, x) - b))
    factors = jsp.linalg.cho_factor(a)
    cho_solve = (lambda f, b: jsp.linalg.cho_solve(factors, b))

    def pos_def_solve(g, b):
        return lax.custom_linear_solve(g, b, cho_solve, symmetric=True)
    return lax.custom_root(f, b, cho_solve, pos_def_solve)