def test_matmul(vector):
a, U, V, P, K, Y = get_matrices(vector=vector, include_dense=True)
# First compute the expected value
expect = np.dot(K, Y)
# Then solve using celerite
Z = np.empty_like(Y)
value = driver.matmul(a, U, V, P, Y, Z)
# Make sure that no copy is made if possible
assert np.allclose(value, Z)
# Check that the solution is correct
assert np.allclose(value, expect)
def test_matmul_fwd(vector):
a, U, V, P, Y = get_matrices(vector=vector)
X0 = driver.matmul(a, U, V, P, Y, np.empty_like(Y))
X = np.empty_like(Y)
Z = np.empty_like(Y)
if vector:
F = np.empty_like(U)
else:
F = np.empty((U.shape[0], U.shape[1] * Y.shape[1]))
G = np.empty_like(F)
X, Z, F, G = backprop.matmul_fwd(a, U, V, P, Y, X, Z, F, G)
assert np.allclose(X0, X)
def test_matmul(vector):
a, U, V, P, Y = get_matrices(vector=vector)
X = driver.matmul(a, U, V, P, Y, np.copy(Y))
check_op(ops.matmul, [a, U, V, P, Y], [X])
check_op(jit(ops.matmul), [a, U, V, P, Y], [X])