@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_linear_solve_zeros(self):

    def explicit_jacobian_solve(matvec, b):
        return lax.stop_gradient(jnp.linalg.solve(api.jacobian(matvec)(b), b))

    def matrix_free_solve(matvec, b):
        return lax.custom_linear_solve(matvec, b, explicit_jacobian_solve, explicit_jacobian_solve)

    def linear_solve(a, b):
        return matrix_free_solve(partial(high_precision_dot, a), b)
    rng = np.random.RandomState(0)
    a = rng.randn(3, 3)
    b = rng.randn(3)
    jtu.check_grads((lambda x: linear_solve(x, b)), (a,), order=2, rtol={np.float32: 0.005})
    jtu.check_grads((lambda x: linear_solve(a, x)), (b,), order=2, rtol={np.float32: 0.005})
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_linear_solve_zeros(self):

    def explicit_jacobian_solve(matvec, b):
        return lax.stop_gradient(jnp.linalg.solve(api.jacobian(matvec)(b), b))

    def matrix_free_solve(matvec, b):
        return lax.custom_linear_solve(matvec, b, explicit_jacobian_solve, explicit_jacobian_solve)

    def linear_solve(a, b):
        return matrix_free_solve(partial(high_precision_dot, a), b)
    rng = np.random.RandomState(0)
    a = rng.randn(3, 3)
    b = rng.randn(3)
    jtu.check_grads((lambda x: linear_solve(x, b)), (a,), order=2, rtol={np.float32: 0.005})
    jtu.check_grads((lambda x: linear_solve(a, x)), (b,), order=2, rtol={np.float32: 0.005})