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