def loss(a, b):
    matvec = partial(high_precision_dot, a)
    x = lax.custom_linear_solve(matvec, b, explicit_jacobian_solve)
    return jnp.sum(x)
def loss(a, b):
    matvec = partial(high_precision_dot, a)
    x = lax.custom_linear_solve(matvec, b, explicit_jacobian_solve)
    return jnp.sum(x)