def solve(a, b):

    def solve(matvec, x):
        return jsp.linalg.solve(a, x)

    def tr_solve(matvec, x):
        return jsp.linalg.solve(a.T, x)
    matvec = partial(high_precision_dot, a)
    return lax.custom_linear_solve(matvec, b, solve, tr_solve)
def solve(a, b):

    def solve(matvec, x):
        return jsp.linalg.solve(a, x)

    def tr_solve(matvec, x):
        return jsp.linalg.solve(a.T, x)
    matvec = partial(high_precision_dot, a)
    return lax.custom_linear_solve(matvec, b, solve, tr_solve)