def test_contraction_tp_additions():
a = CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M)
)
tp = CodegenArrayTensorProduct(P, a, Q)
expr = CodegenArrayContraction(tp, (3, 4))
expected = CodegenArrayTensorProduct(
P,
CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M), (1, 2)),
),
Q
)
assert expr == expected
assert recognize_matrix_expression(expr) == CodegenArrayTensorProduct(P, M*N + N*M, Q)
expr = CodegenArrayContraction(tp, (1, 2), (3, 4), (5, 6))
result = CodegenArrayContraction(
CodegenArrayTensorProduct(
P,
CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M), (1, 2)),
),
Q
), (1, 2), (3, 4))
assert expr == result
assert recognize_matrix_expression(expr) == P*(M*N + N*M)*Q
def test_codegen_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
assert recognize_matrix_expression(expr) == M + Transpose(M)
expr = M[i,j] + N[i,j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N
expr = M[i,j] + N[j,i]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M + N.T
expr = M[i,j]*N[k,l] + N[i,j]*M[k,l]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M))
expr = (M*N*P)[i, j]
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*N*P
expr = Sum(M[i,j]*(N*P)[j,m], (j, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*N*P
expr = Sum((P[j, m] + P[m, j])*(M[i,j]*N[m,n] + N[i,j]*M[m,n]), (j, 0, k-1), (m, 0, k-1))
p1, p2 = _codegen_array_parse(expr)
assert recognize_matrix_expression(p1) == M*P*N + M*P.T*N + N*P*M + N*P.T*M
def test_codegen_array_parse():
expr = M[i, j]
assert _codegen_array_parse(expr) == (M, (i, j))
expr = M[i, j] * N[k, l]
assert _codegen_array_parse(expr) == (CodegenArrayTensorProduct(M, N),
(i, j, k, l))
expr = M[i, j] * N[j, k]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k, j))
expr = Sum(M[i, j] * N[j, k], (j, 0, k - 1))
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N), (1, 2)), (i, k))
expr = M[i, j] + N[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(M, N),
(i, j))
expr = M[i, j] + N[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(N, Permutation([1, 0]))), (i, j))
expr = M[i, j] + M[j, i]
assert _codegen_array_parse(expr) == (CodegenArrayElementwiseAdd(
M, CodegenArrayPermuteDims(M, Permutation([1, 0]))), (i, j))
expr = (M * N * P)[i, j]
assert _codegen_array_parse(expr) == (CodegenArrayContraction(
CodegenArrayTensorProduct(M, N, P), (1, 2), (3, 4)), (i, j))
expr = expr.function # Disregard summation in previous expression
ret1, ret2 = _codegen_array_parse(expr)
assert ret1 == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P),
(1, 2), (3, 4))
assert str(ret2) == "(i, j, _i_1, _i_2)"
expr = KroneckerDelta(i, j) * M[i, k]
assert _codegen_array_parse(expr) == (M, ({i, j}, k))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * M[i, l]
assert _codegen_array_parse(expr) == (M, ({i, j, k}, l))
expr = KroneckerDelta(j, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, k})))
expr = KroneckerDelta(j, m) * KroneckerDelta(
m, k) * (M[i, j] * N[k, l] + N[i, j] * M[k, l])
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(M, N),
CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
Permutation(0, 2)(1, 3))),
(1, 2)), (i, l, frozenset({j, m, k})))
expr = KroneckerDelta(i, j) * KroneckerDelta(j, k) * KroneckerDelta(
k, m) * M[i, 0] * KroneckerDelta(m, n)
assert _codegen_array_parse(expr) == (M, ({i, j, k, m, n}, 0))
expr = M[i, i]
assert _codegen_array_parse(expr) == (CodegenArrayDiagonal(M,
(0, 1)), (i, ))
def test_array_expr_zero_array():
za1 = ZeroArray(k, l, m, n)
zm1 = ZeroMatrix(m, n)
za2 = ZeroArray(k, m, m, n)
zm2 = ZeroMatrix(m, m)
zm3 = ZeroMatrix(k, k)
assert CodegenArrayTensorProduct(M, N, za1) == ZeroArray(k, k, k, k, k, l, m, n)
assert CodegenArrayTensorProduct(M, N, zm1) == ZeroArray(k, k, k, k, m, n)
assert CodegenArrayContraction(za1, (3,)) == ZeroArray(k, l, m)
assert CodegenArrayContraction(zm1, (1,)) == ZeroArray(m)
assert CodegenArrayContraction(za2, (1, 2)) == ZeroArray(k, n)
assert CodegenArrayContraction(zm2, (0, 1)) == 0
assert CodegenArrayDiagonal(za2, (1, 2)) == ZeroArray(k, n, m)
assert CodegenArrayDiagonal(zm2, (0, 1)) == ZeroArray(m)
assert CodegenArrayPermuteDims(za1, [2, 1, 3, 0]) == ZeroArray(m, l, n, k)
assert CodegenArrayPermuteDims(zm1, [1, 0]) == ZeroArray(n, m)
assert CodegenArrayElementwiseAdd(za1) == za1
assert CodegenArrayElementwiseAdd(zm1) == ZeroArray(m, n)
tp1 = CodegenArrayTensorProduct(MatrixSymbol("A", k, l), MatrixSymbol("B", m, n))
assert CodegenArrayElementwiseAdd(tp1, za1) == tp1
tp2 = CodegenArrayTensorProduct(MatrixSymbol("C", k, l), MatrixSymbol("D", m, n))
assert CodegenArrayElementwiseAdd(tp1, za1, tp2) == CodegenArrayElementwiseAdd(tp1, tp2)
assert CodegenArrayElementwiseAdd(M, zm3) == M
assert CodegenArrayElementwiseAdd(M, N, zm3) == CodegenArrayElementwiseAdd(M, N)
def _(expr: CodegenArrayTensorProduct, x: Expr):
args = expr.args
addend_list = []
for i, arg in enumerate(expr.args):
darg = array_derive(arg, x)
if darg == 0:
continue
args_prev = args[:i]
args_succ = args[i + 1:]
shape_prev = reduce(operator.add, map(get_shape, args_prev), ())
shape_succ = reduce(operator.add, map(get_shape, args_succ), ())
addend = CodegenArrayTensorProduct(*args_prev, darg, *args_succ)
tot1 = len(get_shape(x))
tot2 = tot1 + len(shape_prev)
tot3 = tot2 + len(get_shape(arg))
tot4 = tot3 + len(shape_succ)
perm = [i for i in range(tot1, tot2)] + \
[i for i in range(tot1)] + [i for i in range(tot2, tot3)] + \
[i for i in range(tot3, tot4)]
addend = CodegenArrayPermuteDims(addend, _af_invert(perm))
addend_list.append(addend)
if len(addend_list) == 1:
return addend_list[0]
elif len(addend_list) == 0:
return S.Zero
else:
return CodegenArrayElementwiseAdd(*addend_list)
def test_nested_array_elementwise_add():
cg = CodegenArrayContraction(
CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M)), (1, 2))
result = CodegenArrayElementwiseAdd(
CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayContraction(CodegenArrayTensorProduct(N, M), (1, 2)))
assert cg == result
cg = CodegenArrayDiagonal(
CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(M, N),
CodegenArrayTensorProduct(N, M)), (1, 2))
result = CodegenArrayElementwiseAdd(
CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2)),
CodegenArrayDiagonal(CodegenArrayTensorProduct(N, M), (1, 2)))
assert cg == result
def test_remove_trivial_dims():
# Tensor Product:
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b)) == (a * b.T,
[1, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T, b)) == (a * b.T,
[0, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, b.T)) == (a * b.T,
[1, 2])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a.T,
b.T)) == (a * b.T,
[0, 2])
assert _remove_trivial_dims(CodegenArrayTensorProduct(
I, a.T, b.T)) == (a * b.T, [0, 1, 2, 4])
assert _remove_trivial_dims(CodegenArrayTensorProduct(
a.T, I, b.T)) == (a * b.T, [0, 2, 3, 4])
assert _remove_trivial_dims(CodegenArrayTensorProduct(a, I)) == (a, [2, 3])
assert _remove_trivial_dims(CodegenArrayTensorProduct(I, a)) == (a, [0, 1])
assert _remove_trivial_dims(CodegenArrayTensorProduct(
a.T, b.T, c, d)) == (CodegenArrayTensorProduct(a * b.T,
c * d.T), [0, 2, 5, 7])
assert _remove_trivial_dims(CodegenArrayTensorProduct(
a.T, I, b.T, c, d,
I)) == (CodegenArrayTensorProduct(a * b.T, c * d.T,
I), [0, 2, 3, 4, 7, 9])
# Addition:
cg = CodegenArrayElementwiseAdd(CodegenArrayTensorProduct(a, b),
CodegenArrayTensorProduct(c, d))
assert _remove_trivial_dims(cg) == (a * b.T + c * d.T, [1, 3])
# Permute Dims:
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, b),
Permutation(3)(1, 2))
assert _remove_trivial_dims(cg) == (a * b.T, [2, 3])
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(a, I, b),
Permutation(5)(1, 2, 3, 4))
assert _remove_trivial_dims(cg) == (a * b.T, [1, 2, 4, 5])
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, b, a),
Permutation(5)(1, 2, 4, 5, 3))
assert _remove_trivial_dims(cg) == (b * a.T, [0, 3, 4, 5])
# Diagonal:
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
# Contraction:
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, a), (1, 2))
assert _remove_trivial_dims(cg) == (cg, [])
def test_codegen_extra():
if not np:
skip("NumPy not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = np.matrix([[1, 2], [3, 4]])
mb = np.matrix([[1, -2], [-1, 3]])
mc = np.matrix([[2, 0], [1, 2]])
md = np.matrix([[1, -1], [4, 7]])
cg = CodegenArrayTensorProduct(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.einsum(ma, [0, 1], mb, [2, 3])).all()
cg = CodegenArrayElementwiseAdd(M, N)
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == ma + mb).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
f = lambdify((M, N, P), cg, 'numpy')
assert (f(ma, mb, mc) == ma + mb + mc).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
f = lambdify((M, N, P, Q), cg, 'numpy')
assert (f(ma, mb, mc, md) == ma + mb + mc + md).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
f = lambdify((M, ), cg, 'numpy')
assert (f(ma) == ma.T).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 2, 3, 0])
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.transpose(np.einsum(ma, [0, 1], mb, [2, 3]),
(1, 2, 3, 0))).all()
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2))
f = lambdify((M, N), cg, 'numpy')
assert (f(ma, mb) == np.diagonal(np.einsum(ma, [0, 1], mb, [2, 3]),
axis1=1,
axis2=2)).all()
def test_codegen_recognize_matrix_expression():
expr = CodegenArrayElementwiseAdd(M, CodegenArrayPermuteDims(M, [1, 0]))
rec = _recognize_matrix_expression(expr)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [M])])
assert _unfold_recognized_expr(rec) == M + Transpose(M)
expr = M[i, j] + N[i, j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, N])
assert _unfold_recognized_expr(rec) == M + N
expr = M[i, j] + N[j, i]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [M, _RecognizeMatOp(Transpose, [N])])
assert _unfold_recognized_expr(rec) == M + N.T
expr = M[i, j] * N[k, l] + N[i, j] * M[k, l]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(
MatAdd, [_RecognizeMatMulLines([M, N]),
_RecognizeMatMulLines([N, M])])
#assert _unfold_recognized_expr(rec) == TensorProduct(M, N) + TensorProduct(N, M) maybe?
expr = (M * N * P)[i, j]
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatMulLines([_RecognizeMatOp(MatMul, [M, N, P])])
assert _unfold_recognized_expr(rec) == M * N * P
expr = Sum(M[i, j] * (N * P)[j, m], (j, 0, k - 1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatMul, [M, N, P])
assert _unfold_recognized_expr(rec) == M * N * P
expr = Sum((P[j, m] + P[m, j]) * (M[i, j] * N[m, n] + N[i, j] * M[m, n]),
(j, 0, k - 1), (m, 0, k - 1))
p1, p2 = _codegen_array_parse(expr)
rec = _recognize_matrix_expression(p1)
assert rec == _RecognizeMatOp(MatAdd, [
_RecognizeMatOp(MatMul, [
M,
_RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), N
]),
_RecognizeMatOp(MatMul, [
N,
_RecognizeMatOp(MatAdd, [P, _RecognizeMatOp(Transpose, [P])]), M
])
])
assert _unfold_recognized_expr(
rec) == M * (P + P.T) * N + N * (P + P.T) * M
def test_codegen_array_recognize_matrix_mul_lines():
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M), (0, 1))
assert recognize_matrix_expression(cg) == Trace(M)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 1),
(2, 3))
assert recognize_matrix_expression(cg) == Trace(M) * Trace(N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3),
(1, 2))
assert recognize_matrix_expression(cg) == Trace(M * N)
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 2),
(1, 3))
assert recognize_matrix_expression(cg) == Trace(M * N.T)
cg = parse_indexed_expression((M * N * P)[i, j])
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_matrix_expression(M * N * P)
assert recognize_matrix_expression(cg) == M * N * P
cg = parse_indexed_expression((M * N.T * P)[i, j])
assert recognize_matrix_expression(cg) == M * N.T * P
cg = parse_matrix_expression(M * N.T * P)
assert recognize_matrix_expression(cg) == M * N.T * P
cg = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2),
(5, 6))
assert recognize_matrix_expression(cg) == CodegenArrayTensorProduct(
M * N, P * Q)
expr = -2 * M * N
elem = expr[i, j]
cg = parse_indexed_expression(elem)
assert recognize_matrix_expression(cg) == -2 * M * N
a = MatrixSymbol("a", k, 1)
b = MatrixSymbol("b", k, 1)
c = MatrixSymbol("c", k, 1)
cg = CodegenArrayPermuteDims(
CodegenArrayContraction(
CodegenArrayTensorProduct(
a,
CodegenArrayElementwiseAdd(
CodegenArrayTensorProduct(b, c),
CodegenArrayTensorProduct(c, b),
)), (2, 4)), [0, 1, 3, 2])
assert recognize_matrix_expression(cg) == a * (b.T * c + c.T * b)
za = ZeroArray(m, n)
assert recognize_matrix_expression(za) == ZeroMatrix(m, n)
cg = CodegenArrayTensorProduct(3, M)
assert recognize_matrix_expression(cg) == 3 * M
def test_codegen_array_shape():
expr = CodegenArrayTensorProduct(M, N, P, Q)
assert expr.shape == (k, k, k, k, k, k, k, k)
Z = MatrixSymbol("Z", m, n)
expr = CodegenArrayTensorProduct(M, Z)
assert expr.shape == (k, k, m, n)
expr2 = CodegenArrayContraction(expr, (0, 1))
assert expr2.shape == (m, n)
expr2 = CodegenArrayDiagonal(expr, (0, 1))
assert expr2.shape == (m, n, k)
exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0])
assert exprp.shape == (m, k, n, k)
expr3 = CodegenArrayTensorProduct(N, Z)
expr2 = CodegenArrayElementwiseAdd(expr, expr3)
assert expr2.shape == (k, k, m, n)
# Contraction along axes with discordant dimensions:
raises(ValueError, lambda: CodegenArrayContraction(expr, (1, 2)))
# Also diagonal needs the same dimensions:
raises(ValueError, lambda: CodegenArrayDiagonal(expr, (1, 2)))
def test_codegen_array_doit():
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
M = M.as_explicit()
N = N.as_explicit()
P = P.as_explicit()
Q = Q.as_explicit()
expr = CodegenArrayTensorProduct(M, N, P, Q)
assert expr.doit() == tensorproduct(M, N, P, Q)
expr2 = CodegenArrayContraction(expr, (0, 1))
assert expr2.doit() == tensorcontraction(tensorproduct(M, N, P, Q), (0, 1))
expr2 = CodegenArrayDiagonal(expr, (0, 1))
#assert expr2 = ... # TODO: not implemented
expr = CodegenArrayTensorProduct(M, N)
exprp = CodegenArrayPermuteDims(expr, [2, 1, 3, 0])
assert exprp.doit() == permutedims(tensorproduct(M, N), [2, 1, 3, 0])
expr = CodegenArrayElementwiseAdd(M, N)
assert expr.doit() == M + N
def test_codegen_array_flatten():
# Flatten nested CodegenArrayTensorProduct objects:
expr1 = CodegenArrayTensorProduct(M, N)
expr2 = CodegenArrayTensorProduct(P, Q)
expr = CodegenArrayTensorProduct(expr1, expr2)
assert expr == CodegenArrayTensorProduct(M, N, P, Q)
assert expr.args == (M, N, P, Q)
# Flatten mixed CodegenArrayTensorProduct and CodegenArrayContraction objects:
cg1 = CodegenArrayContraction(expr1, (1, 2))
cg2 = CodegenArrayContraction(expr2, (0, 3))
expr = CodegenArrayTensorProduct(cg1, cg2)
assert expr == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 2), (4, 7))
expr = CodegenArrayTensorProduct(M, cg1)
assert expr == CodegenArrayContraction(CodegenArrayTensorProduct(M, M, N), (3, 4))
# Flatten nested CodegenArrayContraction objects:
cgnested = CodegenArrayContraction(cg1, (0, 1))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (0, 3), (1, 2))
cgnested = CodegenArrayContraction(CodegenArrayTensorProduct(cg1, cg2), (0, 3))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 6), (1, 2), (4, 7))
cg3 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4))
cgnested = CodegenArrayContraction(cg3, (0, 1))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 5), (1, 3), (2, 4))
cgnested = CodegenArrayContraction(cg3, (0, 3), (1, 2))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 7), (1, 3), (2, 4), (5, 6))
cg4 = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7))
cgnested = CodegenArrayContraction(cg4, (0, 1))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 2), (1, 5), (3, 7))
cgnested = CodegenArrayContraction(cg4, (0, 1), (2, 3))
assert cgnested == CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 2), (1, 5), (3, 7), (4, 6))
cg = CodegenArrayDiagonal(cg4)
assert cg == cg4
assert isinstance(cg, type(cg4))
# Flatten nested CodegenArrayDiagonal objects:
cg1 = CodegenArrayDiagonal(expr1, (1, 2))
cg2 = CodegenArrayDiagonal(expr2, (0, 3))
cg3 = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4))
cg4 = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7))
cgnested = CodegenArrayDiagonal(cg1, (0, 1))
assert cgnested == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2), (0, 3))
cgnested = CodegenArrayDiagonal(cg3, (1, 2))
assert cgnested == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 3), (2, 4), (5, 6))
cgnested = CodegenArrayDiagonal(cg4, (1, 2))
assert cgnested == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (1, 5), (3, 7), (2, 4))
cg = CodegenArrayElementwiseAdd(M, N)
cg2 = CodegenArrayElementwiseAdd(cg, P)
assert isinstance(cg2, CodegenArrayElementwiseAdd)
assert cg2.args == (M, N, P)
assert cg2.shape == (k, k)
def test_codegen_extra():
if not tf:
skip("TensorFlow not installed")
graph = tf.Graph()
with graph.as_default():
session = tf.compat.v1.Session()
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = tf.constant([[1, 2], [3, 4]])
mb = tf.constant([[1, -2], [-1, 3]])
mc = tf.constant([[2, 0], [1, 2]])
md = tf.constant([[1, -1], [4, 7]])
cg = CodegenArrayTensorProduct(M, N)
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,cd", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ij,kl", ma, mb))
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N)
assert tensorflow_code(cg) == 'tensorflow.math.add(M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(ma + mb)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
assert tensorflow_code(cg) == \
'tensorflow.math.add(tensorflow.math.add(M, N), P)'
f = lambdify((M, N, P), cg, 'tensorflow')
y = session.run(f(ma, mb, mc))
c = session.run(ma + mb + mc)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
assert tensorflow_code(cg) == \
'tensorflow.math.add(' \
'tensorflow.math.add(tensorflow.math.add(M, N), P), Q)'
f = lambdify((M, N, P, Q), cg, 'tensorflow')
y = session.run(f(ma, mb, mc, md))
c = session.run(ma + mb + mc + md)
assert (y == c).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
assert tensorflow_code(cg) == 'tensorflow.transpose(M, [1, 0])'
f = lambdify((M, ), cg, 'tensorflow')
y = session.run(f(ma))
c = session.run(tf.transpose(ma))
assert (y == c).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N),
[1, 2, 3, 0])
assert tensorflow_code(cg) == \
'tensorflow.transpose(' \
'tensorflow.linalg.einsum("ab,cd", M, N), [1, 2, 3, 0])'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.transpose(tf.einsum("ab,cd", ma, mb), [1, 2, 3, 0]))
assert (y == c).all()
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2))
assert tensorflow_code(cg) == \
'tensorflow.linalg.einsum("ab,bc->acb", M, N)'
f = lambdify((M, N), cg, 'tensorflow')
y = session.run(f(ma, mb))
c = session.run(tf.einsum("ab,bc->acb", ma, mb))
assert (y == c).all()
def _(expr: CodegenArrayElementwiseAdd, x: Expr):
return CodegenArrayElementwiseAdd(
*[array_derive(arg, x) for arg in expr.args])
def test_codegen_extra():
if not torch:
skip("Torch not installed")
M = MatrixSymbol("M", 2, 2)
N = MatrixSymbol("N", 2, 2)
P = MatrixSymbol("P", 2, 2)
Q = MatrixSymbol("Q", 2, 2)
ma = torch.tensor([[1, 2], [3, 4]])
mb = torch.tensor([[1, -2], [-1, 3]])
mc = torch.tensor([[2, 0], [1, 2]])
md = torch.tensor([[1, -1], [4, 7]])
cg = CodegenArrayTensorProduct(M, N)
assert torch_code(cg) == 'torch.einsum("ab,cd", [M, N])'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ij,kl", ma, mb)
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N)
assert torch_code(cg) == 'torch.add(M, N)'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = ma + mb
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P)
assert torch_code(cg) == 'torch.add(torch.add(M, N), P)'
f = lambdify((M, N, P), cg, 'torch')
y = f(ma, mb, mc)
c = ma + mb + mc
assert (y == c).all()
cg = CodegenArrayElementwiseAdd(M, N, P, Q)
assert torch_code(cg) == 'torch.add(torch.add(torch.add(M, N), P), Q)'
f = lambdify((M, N, P, Q), cg, 'torch')
y = f(ma, mb, mc, md)
c = ma + mb + mc + md
assert (y == c).all()
cg = CodegenArrayPermuteDims(M, [1, 0])
assert torch_code(cg) == 'M.permute(1, 0)'
f = lambdify((M, ), cg, 'torch')
y = f(ma)
c = ma.T
assert (y == c).all()
cg = CodegenArrayPermuteDims(CodegenArrayTensorProduct(M, N), [1, 2, 3, 0])
assert torch_code(
cg) == 'torch.einsum("ab,cd", [M, N]).permute(1, 2, 3, 0)'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ab,cd", ma, mb).permute(1, 2, 3, 0)
assert (y == c).all()
cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (1, 2))
assert torch_code(cg) == 'torch.einsum("ab,bc->acb", [M, N])'
f = lambdify((M, N), cg, 'torch')
y = f(ma, mb)
c = torch.einsum("ab,bc->acb", ma, mb)
assert (y == c).all()
