def test_solve_upper_fwd(vector): x, c, a, U, V, Y = get_matrices(vector=vector) if vector: Y = Y[:, None] d, W = driver.factor(x, c, a, U, V, a, V) Z0 = driver.solve_upper(x, c, U, W, Y, np.copy(Y)) Z = np.empty_like(Y) F = np.empty((U.shape[0], U.shape[1], Y.shape[1])) Z, F = backprop.solve_upper_fwd(x, c, U, W, Y, Z, F) assert np.allclose(Z0, Z)
def test_solve_upper(vector): x, c, a, U, V, K, Y = get_matrices(vector=vector, include_dense=True) if vector: Y = Y[:, None] # First compute the expected value expect = np.linalg.solve(np.linalg.cholesky(K).T, Y) # Then solve using celerite d, W = driver.factor(x, c, a, U, V, a, V) Y /= np.sqrt(d)[:, None] value = driver.solve_upper(x, c, U, W, Y, Y) # Make sure that no copy is made if possible assert np.allclose(value, Y) # Check that the solution is correct assert np.allclose(value, expect)
def test_solve_upper(): x, c, a, U, V, Y = get_matrices() d, W = driver.factor(x, c, a, U, V, np.copy(a), np.copy(V)) Z = driver.solve_upper(x, c, U, W, Y, np.zeros_like(Y)) check_op(ops.solve_upper, [x, c, U, W, Y], [Z]) check_op(jit(ops.solve_upper), [x, c, U, W, Y], [Z])