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

    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)
    rng = np.random.RandomState(0)
    a = rng.randn(3, 3)
    b = rng.randn(3)
    expected = jnp.linalg.solve(a, b)
    actual = linear_solve(a, b)
    self.assertAllClose(expected, actual)
    jtu.check_grads(linear_solve, (a, b), order=2, rtol=0.002)
    jtu.check_grads(api.jit(linear_solve), (a, b), order=2, rtol={np.float32: 0.002})
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_linear_solve_lu(self):

    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)
    rng = np.random.RandomState(0)
    a = rng.randn(3, 3)
    b = rng.randn(3)
    expected = jnp.linalg.solve(a, b)
    actual = linear_solve(a, b)
    self.assertAllClose(expected, actual)
    jtu.check_grads(linear_solve, (a, b), order=2, rtol=0.002)
    jtu.check_grads(api.jit(linear_solve), (a, b), order=2, rtol={np.float32: 0.002})