def test_custom_linear_solve_complex(self):

    def solve(a, b):

        def solve(matvec, x):
            return jsp.linalg.solve(a, x)

        def tr_solve(matvec, x):
            return jsp.linalg.solve(a.T, x)
        matvec = partial(high_precision_dot, a)
        return lax.custom_linear_solve(matvec, b, solve, tr_solve)
    rng = np.random.RandomState(0)
    a = ((0.5 * rng.randn(2, 2)) + (0.5j * rng.randn(2, 2)))
    b = ((0.5 * rng.randn(2)) + (0.5j * rng.randn(2)))
    jtu.check_grads(solve, (a, b), order=2, rtol=0.01)
def test_custom_linear_solve_complex(self):

    def solve(a, b):

        def solve(matvec, x):
            return jsp.linalg.solve(a, x)

        def tr_solve(matvec, x):
            return jsp.linalg.solve(a.T, x)
        matvec = partial(high_precision_dot, a)
        return lax.custom_linear_solve(matvec, b, solve, tr_solve)
    rng = np.random.RandomState(0)
    a = ((0.5 * rng.randn(2, 2)) + (0.5j * rng.randn(2, 2)))
    b = ((0.5 * rng.randn(2)) + (0.5j * rng.randn(2)))
    jtu.check_grads(solve, (a, b), order=2, rtol=0.01)