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

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

    def loss(a, b):
        matvec = partial(high_precision_dot, a)
        x = lax.custom_linear_solve(matvec, b, explicit_jacobian_solve)
        return jnp.sum(x)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    jtu.check_grads(loss, (a, b), order=2, modes=['fwd'], atol={np.float32: 0.002, np.float64: 1e-11})
    jtu.check_grads(api.vmap(loss), (a[None, :, :], b[None, :]), order=2, modes=['fwd'], atol={np.float32: 0.002, np.float64: 1e-11})
    with self.assertRaisesRegex(TypeError, 'transpose_solve required'):
        api.grad(loss)(a, b)
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_linear_solve_without_transpose_solve(self):

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

    def loss(a, b):
        matvec = partial(high_precision_dot, a)
        x = lax.custom_linear_solve(matvec, b, explicit_jacobian_solve)
        return jnp.sum(x)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    jtu.check_grads(loss, (a, b), order=2, modes=['fwd'], atol={np.float32: 0.002, np.float64: 1e-11})
    jtu.check_grads(api.vmap(loss), (a[None, :, :], b[None, :]), order=2, modes=['fwd'], atol={np.float32: 0.002, np.float64: 1e-11})
    with self.assertRaisesRegex(TypeError, 'transpose_solve required'):
        api.grad(loss)(a, b)