def test_sparse_ops():
T.set_backend("numpy")
size = 5
x_sparse = T.random([size, size], format='coo', density=0.1)
x_dense = T.to_numpy(x_sparse)
assert isinstance(x_sparse, sparse._coo.core.COO)
assert isinstance(x_dense, np.ndarray)
y_dense = T.random([size, size])
y_sparse = T.tensor(y_dense, format="coo")
assert isinstance(y_sparse, sparse._coo.core.COO)
assert isinstance(y_dense, np.ndarray)
# test einsum
einsum1 = T.einsum("ab,bc->ac", x_dense, y_dense)
einsum2 = T.einsum("ab,bc->ac", x_dense, y_sparse)
einsum3 = T.einsum("ab,bc->ac", x_sparse, y_sparse)
assert float_eq(einsum1, einsum2)
assert float_eq(einsum1, einsum3)
# test solve_tri, first change matrices to full-rank ones
x_sparse += T.tensor(T.identity(size), format="coo")
y_sparse += T.tensor(T.identity(size), format="coo")
x_dense += T.identity(size)
y_dense += T.identity(size)
out1 = T.solve_tri(x_sparse, y_sparse)
out2 = T.solve_tri(x_dense, y_dense)
assert float_eq(out1, out2)
def test_jacobian_einsum(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
x1 = ad.Variable(name="x1", shape=[3, 3, 3])
x2 = ad.Variable(name="x2", shape=[3, 3, 3])
y = ad.einsum("ikl,jkl->ijk", x1, x2)
jacobian_x1, jacobian_x2 = ad.jacobians(y, [x1, x2])
executor = ad.Executor([y, jacobian_x1, jacobian_x2])
x1_val = T.random((3, 3, 3))
x2_val = T.random((3, 3, 3))
y_val, jacobian_x1_val, jacobian_x2_val = executor.run(feed_dict={
x1: x1_val,
x2: x2_val,
})
I = T.identity(3)
expected_jacobian_x1_val = T.einsum("im,kn,jno->ijkmno", I, I, x2_val)
expected_jacobian_x2_val = T.einsum("jm,kn,ino->ijkmno", I, I, x1_val)
assert isinstance(y, ad.Node)
assert T.array_equal(y_val, T.einsum("ikl,jkl->ijk", x1_val, x2_val))
assert T.array_equal(jacobian_x1_val, expected_jacobian_x1_val)
assert T.array_equal(jacobian_x2_val, expected_jacobian_x2_val)
def init_rand_cp(dim, size, rank):
X = T.random([size for _ in range(dim)])
A_list = []
for i in range(dim):
A_list.append(T.random((size, rank)))
return A_list, X
def init_rand_tucker(dim, size, rank):
assert size > rank
X = T.random([size for _ in range(dim)])
core = T.random([rank for _ in range(dim)])
A_list = []
for i in range(dim):
# for Tucker, factor matrices are orthogonal
mat, _, _ = T.svd(T.random((size, rank)))
A_list.append(mat[:, :rank])
return A_list, core, X
def test_large_matmul_chain(backendopt):
n = 60
size = 3
for datatype in backendopt:
T.set_backend(datatype)
# build the graph of x_1 @ ... @ x_n
x_list = [
ad.Variable(name=f"x{i}", shape=[size, size]) for i in range(n)
]
prev_char = chr(192)
left_char = prev_char
for i in range(n):
new_char = chr(ord(prev_char) + 1)
x_list[i].subscripts = f"{prev_char}{new_char}"
prev_char = new_char
right_char = prev_char
input_subs = ','.join([node.subscripts for node in x_list])
einsum_subscripts = input_subs + '->' + left_char + right_char
out = ad.einsum(einsum_subscripts, *x_list)
# decompose the large einsum, and rewrite the einsum expression of the
# generated einsum tree so there's no unicode character
out = optimize(out)
executor = ad.Executor([out])
x_val_list = [T.random([size, size]) for _ in range(n)]
out_val, = executor.run(feed_dict=dict(zip(x_list, x_val_list)))
out_val_matmul = x_val_list[0]
for i in range(1, n):
out_val_matmul = out_val_matmul @ x_val_list[i]
assert float_eq(out_val, out_val_matmul, tol=1e-2)
def gen_dict(input_nodes):
"""Generates a random dict for executor to use.
"""
feed_dict = {}
for i_node in input_nodes:
feed_dict[i_node] = T.random(i_node.shape)
return feed_dict
def test_hessian_quadratic(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
x = ad.Variable(name="x", shape=[3])
H = ad.Variable(name="H", shape=[3, 3])
y = ad.einsum("i,ij,j->", x, H, x)
hessian = ad.hessian(y, [x])
executor = ad.Executor([hessian[0][0]])
x_val = T.random([3])
H_val = T.random((3, 3))
hessian_val, = executor.run(feed_dict={x: x_val, H: H_val})
assert T.array_equal(hessian_val, H_val + T.transpose(H_val))
def test_HinverseG(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
N = 10
T.seed(1224)
A = T.random([N, N])
A = T.transpose(A) @ A
A = A + T.identity(N)
b = T.random([N])
def hess_fn(x):
return [T.einsum("ab,b->a", A, x[0])]
error_tol = 1e-9
x, = conjugate_gradient(hess_fn, [b], error_tol)
assert (T.norm(T.abs(T.einsum("ab,b->a", A, x) - b)) <= 1e-4)
def test_executor_debug_orthonormal(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
A = ad.Matrix(name="A", shape=[3, 3], orthonormal='row')
out = ad.einsum("ab,bc->ac", A, A)
A_val, _, _ = T.svd(T.random((3, 3)))
executor = ad.Executor([out])
executor.run(feed_dict={A: A_val}, debug=True)
def test_tensorinv_matrix():
for datatype in backends:
T.set_backend(datatype)
x = ad.Variable(name="x", shape=[3, 3])
inv_x = ad.tensorinv(x)
executor = ad.Executor([inv_x])
x_val = T.random([3, 3])
inv_x_val, = executor.run(feed_dict={x: x_val})
assert T.array_equal(inv_x_val, T.inv(x_val))
def test_simpledot():
def testfunc(w, b, x):
return np.dot(w, x) + b + np.ones(5), x
T.set_backend('jax')
w = T.random((5, 10))
b = T.random((5, ))
x = T.random((10, ))
inputs = [w, b, x]
out_nodes, variables = make_graph(testfunc, *inputs)
executor = ad.Executor(out_nodes)
feed_dict = dict(zip(variables, inputs))
outvals = executor.run(feed_dict=feed_dict)
expect_outvals = testfunc(*inputs)
for outval, expect_outval in zip(outvals, expect_outvals):
assert T.norm(outval - expect_outval) < 1e-6
def test_sparse_einsum_graph():
T.set_backend("taco")
size = 5
coo = formats.SparseFormat([formats.compressed, formats.compressed])
csc = formats.SparseFormat([formats.compressed, formats.dense])
x1 = ad.Variable(name="x1", shape=[size, size], format=coo)
x2 = ad.Variable(name="x2", shape=[size, size])
y = ad.einsum('ik,kj->ij', x1, x2, out_format=csc)
executor = ad.Executor([y])
x1_val = T.random([size, size], format='coo', density=0.1)
x2_val = T.random([size, size])
y_val, = executor.run(feed_dict={x1: x1_val, x2: x2_val}, debug=True)
expected_yval = T.einsum("ab,bc->ac", x1_val, x2_val)
assert float_eq(y_val, expected_yval)
assert isinstance(y_val, pt.pytensor.taco_tensor.tensor)
def test_einsum():
def testfunc(a, b):
# Note: because our executor output is always a list, here a list is also
# returned to make them consistent.
return np.einsum('ij,jk->ik', a, b),
T.set_backend('jax')
a = T.random((5, 10))
b = T.random((10, 5))
inputs = [a, b]
out_nodes, variables = make_graph(testfunc, *inputs)
executor = ad.Executor(out_nodes)
feed_dict = dict(zip(variables, inputs))
outvals = executor.run(feed_dict=feed_dict)
expect_outvals = testfunc(*inputs)
for outval, expect_outval in zip(outvals, expect_outvals):
assert T.norm(outval - expect_outval) < 1e-6
def test_mul():
def testfunc(w, b, x):
# Note: because our executor output is always a list, here a list is also
# returned to make them consistent.
return w * x + b,
T.set_backend('jax')
w = T.random((5, 10))
b = T.random((5, 10))
x = T.random((5, 10))
inputs = [w, b, x]
out_nodes, variables = make_graph(testfunc, *inputs)
executor = ad.Executor(out_nodes)
feed_dict = dict(zip(variables, inputs))
outvals = executor.run(feed_dict=feed_dict)
expect_outvals = testfunc(*inputs)
for outval, expect_outval in zip(outvals, expect_outvals):
assert T.norm(outval - expect_outval) < 1e-6
def test_tensorinv_odd_dim(backendopt):
for datatype in backendopt:
T.set_backend(datatype)
x = ad.Variable(name="x", shape=[24, 8, 3])
inv_x = ad.tensorinv(x, ind=1)
assert inv_x.shape == [8, 3, 24]
assert inv_x.input_indices_length == 2
executor = ad.Executor([inv_x])
x_val = T.random([24, 8, 3])
inv_x_val, = executor.run(feed_dict={x: x_val})
assert T.array_equal(inv_x_val, T.tensorinv(x_val, ind=1))
def test_executor_debug_symmetry(backendopt):
for datatype in backendopt:
if datatype == "taco":
# Taco addition (line 76) will output sparse matrix even though the input is dense.
# This will make the format check fail.
continue
T.set_backend(datatype)
A = ad.Variable(name="A", shape=[3, 3], symmetry=[[0, 1]])
out = ad.einsum("ab,bc->ac", A, A)
A_val = T.random((3, 3))
A_val += T.transpose(A_val)
executor = ad.Executor([out])
executor.run(feed_dict={A: A_val}, debug=True)
def test_cpd():
def testfunc(A, B, C, X):
T = np.einsum("ia,ja,ka->ijk", A, B, C)
res = T - X
return np.einsum("ijk,ijk->", res, res),
size = 3
rank = 2
T.set_backend('jax')
X = T.random((size, size, size))
A = T.random((size, rank))
B = T.random((size, rank))
C = T.random((size, rank))
inputs = [A, B, C, X]
out_nodes, variables = make_graph(testfunc, *inputs)
executor = ad.Executor(out_nodes)
feed_dict = dict(zip(variables, inputs))
outvals = executor.run(feed_dict=feed_dict)
expect_outvals = testfunc(*inputs)
for outval, expect_outval in zip(outvals, expect_outvals):
assert T.norm(outval - expect_outval) < 1e-6