def test_custom_linear_solve_cholesky(self):

    def positive_definite_solve(a, b):
        factors = jsp.linalg.cho_factor(a)

        def solve(matvec, x):
            return jsp.linalg.cho_solve(factors, x)
        matvec = partial(high_precision_dot, a)
        return lax.custom_linear_solve(matvec, b, solve, symmetric=True)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    expected = jnp.linalg.solve(np.asarray(posify(a)), b)
    actual = positive_definite_solve(posify(a), b)
    self.assertAllClose(expected, actual)
    actual = api.jit(positive_definite_solve)(posify(a), b)
    self.assertAllClose(expected, actual)
    jtu.check_grads((lambda x, y: positive_definite_solve(posify(x), y)), (a, b), order=2, rtol=0.01)
def test_custom_linear_solve_cholesky(self):

    def positive_definite_solve(a, b):
        factors = jsp.linalg.cho_factor(a)

        def solve(matvec, x):
            return jsp.linalg.cho_solve(factors, x)
        matvec = partial(high_precision_dot, a)
        return lax.custom_linear_solve(matvec, b, solve, symmetric=True)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    expected = jnp.linalg.solve(np.asarray(posify(a)), b)
    actual = positive_definite_solve(posify(a), b)
    self.assertAllClose(expected, actual)
    actual = api.jit(positive_definite_solve)(posify(a), b)
    self.assertAllClose(expected, actual)
    jtu.check_grads((lambda x, y: positive_definite_solve(posify(x), y)), (a, b), order=2, rtol=0.01)