def linear_solve(a, b):
    a_factors = jsp.linalg.lu_factor(a)
    at_factors = jsp.linalg.lu_factor(a.T)

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

    def transpose_solve(vecmat, x):
        return jsp.linalg.lu_solve(at_factors, x)
    return lax.custom_linear_solve(partial(high_precision_dot, a), b, solve, transpose_solve)
def linear_solve(a, b):
    a_factors = jsp.linalg.lu_factor(a)
    at_factors = jsp.linalg.lu_factor(a.T)

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

    def transpose_solve(vecmat, x):
        return jsp.linalg.lu_solve(at_factors, x)
    return lax.custom_linear_solve(partial(high_precision_dot, a), b, solve, transpose_solve)