def positive_definite_solve(a, b):
    factors = jsp.linalg.cho_factor(a)

    def solve(matvec, x):
        return jsp.linalg.cho_solve(factors, x)
    matvec = partial(high_precision_dot, a)
    return lax.custom_linear_solve(matvec, b, solve, symmetric=True)
def positive_definite_solve(a, b):
    factors = jsp.linalg.cho_factor(a)

    def solve(matvec, x):
        return jsp.linalg.cho_solve(factors, x)
    matvec = partial(high_precision_dot, a)
    return lax.custom_linear_solve(matvec, b, solve, symmetric=True)