Exemplo n.º 1
0
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)
Exemplo n.º 2
0
def test_arrayexpr_derivatives1():

    res = array_derive(X, X)
    assert res == CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, I),
                                          [0, 2, 1, 3])

    cg = CodegenArrayTensorProduct(A, X, B)
    res = array_derive(cg, X)
    assert res == CodegenArrayPermuteDims(
        CodegenArrayTensorProduct(I, A, I, B), [0, 4, 2, 3, 1, 5, 6, 7])

    cg = CodegenArrayContraction(X, (0, 1))
    res = array_derive(cg, X)
    assert res == CodegenArrayContraction(CodegenArrayTensorProduct(I, I),
                                          (1, 3))

    cg = CodegenArrayDiagonal(X, (0, 1))
    res = array_derive(cg, X)
    assert res == CodegenArrayDiagonal(CodegenArrayTensorProduct(I, I), (1, 3))

    cg = ElementwiseApplyFunction(sin, X)
    res = array_derive(cg, X)
    assert res.dummy_eq(
        CodegenArrayDiagonal(
            CodegenArrayTensorProduct(ElementwiseApplyFunction(cos, X), I, I),
            (0, 3), (1, 5)))
Exemplo n.º 3
0
def test_codegen_array_diagonal():
    cg = CodegenArrayDiagonal(M, (1, 0))
    assert cg == CodegenArrayDiagonal(M, (0, 1))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P), (4, 1),
                              (2, 0))
    assert cg == CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P),
                                      (1, 4), (0, 2))
Exemplo n.º 4
0
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))

    # 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))
Exemplo n.º 5
0
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, ))
Exemplo n.º 6
0
def test_normalize_diagonal_permutedims():
    tp = CodegenArrayTensorProduct(M, Q, N, P)
    expr = CodegenArrayDiagonal(
        CodegenArrayPermuteDims(tp, [0, 1, 2, 4, 7, 6, 3, 5]), (2, 4, 5), (6, 7),
        (0, 3))
    result = CodegenArrayDiagonal(tp, (2, 6, 7), (3, 5), (0, 4))
    assert expr == result

    tp = CodegenArrayTensorProduct(M, N, P, Q)
    expr = CodegenArrayDiagonal(CodegenArrayPermuteDims(tp, [0, 5, 2, 4, 1, 6, 3, 7]), (1, 2, 6), (3, 4))
    result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, P, N, Q), (3, 4, 5), (1, 2))
    assert expr == result
Exemplo n.º 7
0
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
Exemplo n.º 8
0
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, [])
Exemplo n.º 9
0
def _(expr: ElementwiseApplyFunction, x: Expr):
    assert get_rank(expr) == 2
    assert get_rank(x) == 2
    fdiff = expr._get_function_fdiff()
    dexpr = array_derive(expr.expr, x)
    tp = CodegenArrayTensorProduct(ElementwiseApplyFunction(fdiff, expr.expr),
                                   dexpr)
    td = CodegenArrayDiagonal(tp, (0, 4), (1, 5))
    return td
Exemplo n.º 10
0
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)))
Exemplo n.º 11
0
def test_diag2contraction_diagmatrix():
    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(CodegenArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3))

    raises(ValueError, lambda: CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M), (1, 2)))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M), (1, 2))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(CodegenArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M, N, b.T), (1, 2), (4, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7), (3, 9))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 2), (4, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (1, 6), (3, 9))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 4), (3, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1), DiagMatrix(a), DiagMatrix(b.T)), (3, 6), (2, 9))
Exemplo n.º 12
0
def test_normalize_diagonal_contraction():
    tp = CodegenArrayTensorProduct(M, N, P, Q)
    expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (1, 3, 4)), (0, 3))
    result = CodegenArrayDiagonal(
        CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q), (0, 6)),
        (0, 2, 3))
    assert expr == result

    expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (0, 1, 2, 3, 7)),
                                   (1, 2, 3))
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q),
                                     (0, 1, 2, 3, 5, 6, 7))
    assert expr == result

    expr = CodegenArrayContraction(CodegenArrayDiagonal(tp, (0, 2, 6, 7)),
                                   (1, 2, 3))
    result = CodegenArrayDiagonal(CodegenArrayContraction(tp, (3, 4, 5)),
                                  (0, 2, 3, 4))
    assert expr == result

    td = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N, P, Q), (0, 3))
    expr = CodegenArrayContraction(td, (2, 1), (0, 4, 6, 5, 3))
    result = CodegenArrayContraction(CodegenArrayTensorProduct(M, N, P, Q),
                                     (0, 1, 3, 5, 6, 7), (2, 4))
    assert expr == result
Exemplo n.º 13
0
def test_diag2contraction_diagmatrix():
    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, a), (1, 2))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(M, OneArray(1), DiagMatrix(a)), (1, 3))

    raises(
        ValueError,
        lambda: CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M), (1, 2)))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M), (1, 2))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1), M, DiagMatrix(a.T)), (1, 4))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a.T, M, N, b.T),
                              (1, 2), (4, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1), M, N, OneArray(1),
                                  DiagMatrix(a.T), DiagMatrix(b.T)), (1, 7),
        (3, 9))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 2),
                              (4, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1),
                                  M, N, OneArray(1), DiagMatrix(a),
                                  DiagMatrix(b.T)), (1, 6), (3, 9))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, M, N, b.T), (0, 4),
                              (3, 7))
    res = _array_diag2contr_diagmatrix(cg)
    assert res.shape == cg.shape
    assert res == CodegenArrayContraction(
        CodegenArrayTensorProduct(OneArray(1),
                                  M, N, OneArray(1), DiagMatrix(a),
                                  DiagMatrix(b.T)), (3, 6), (2, 9))

    I1 = Identity(1)
    x = MatrixSymbol("x", k, 1)
    A = MatrixSymbol("A", k, k)
    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, A.T, I1), (0, 2))
    assert _array_diag2contr_diagmatrix(cg).shape == cg.shape
    assert array2matrix(cg).shape == cg.shape
Exemplo n.º 14
0
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()
Exemplo n.º 15
0
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
Exemplo n.º 16
0
def tensordiagonal(array, *diagonal_axes):
    """
    Diagonalization of an array-like object on the specified axes.

    This is equivalent to multiplying the expression by Kronecker deltas
    uniting the axes.

    The diagonal indices are put at the end of the axes.

    Examples
    ========

    ``tensordiagonal`` acting on a 2-dimensional array by axes 0 and 1 is
    equivalent to the diagonal of the matrix:

    >>> from sympy import Array, tensordiagonal
    >>> from sympy import Matrix, eye
    >>> tensordiagonal(eye(3), (0, 1))
    [1, 1, 1]

    >>> from sympy.abc import a,b,c,d
    >>> m1 = Matrix([[a, b], [c, d]])
    >>> tensordiagonal(m1, [0, 1])
    [a, d]

    In case of higher dimensional arrays, the diagonalized out dimensions
    are appended removed and appended as a single dimension at the end:

    >>> A = Array(range(18), (3, 2, 3))
    >>> A
    [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]], [[12, 13, 14], [15, 16, 17]]]
    >>> tensordiagonal(A, (0, 2))
    [[0, 7, 14], [3, 10, 17]]
    >>> from sympy import permutedims
    >>> tensordiagonal(A, (0, 2)) == permutedims(Array([A[0, :, 0], A[1, :, 1], A[2, :, 2]]), [1, 0])
    True

    """
    if any([len(i) <= 1 for i in diagonal_axes]):
        raise ValueError("need at least two axes to diagonalize")

    from sympy.tensor.array.expressions.array_expressions import _ArrayExpr
    from sympy.codegen.array_utils import _CodegenArrayAbstract, CodegenArrayDiagonal
    from sympy import MatrixSymbol
    if isinstance(array, (_ArrayExpr, _CodegenArrayAbstract, MatrixSymbol)):
        return CodegenArrayDiagonal(array, *diagonal_axes)

    array, remaining_indices, remaining_shape, diagonal_deltas = _util_contraction_diagonal(
        array, *diagonal_axes)

    # Compute the diagonalized array:
    #
    # 1. external for loops on all undiagonalized indices.
    #    Undiagonalized indices are determined by the combinatorial product of
    #    the absolute positions of the remaining indices.
    # 2. internal loop on all diagonal indices.
    #    It appends the values of the absolute diagonalized index and the absolute
    #    undiagonalized index for the external loop.
    diagonalized_array = []
    diagonal_shape = [len(i) for i in diagonal_deltas]
    for icontrib in itertools.product(*remaining_indices):
        index_base_position = sum(icontrib)
        isum = []
        for sum_to_index in itertools.product(*diagonal_deltas):
            idx = array._get_tuple_index(index_base_position +
                                         sum(sum_to_index))
            isum.append(array[idx])

        isum = type(array)(isum).reshape(*diagonal_shape)
        diagonalized_array.append(isum)

    return type(array)(diagonalized_array, remaining_shape + diagonal_shape)
Exemplo n.º 17
0
def test_parsing_of_matrix_expressions():

    expr = M*N
    assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, N), (1, 2))

    expr = Transpose(M)
    assert parse_matrix_expression(expr) == CodegenArrayPermuteDims(M, [1, 0])

    expr = M*Transpose(N)
    assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, CodegenArrayPermuteDims(N, [1, 0])), (1, 2))

    expr = 3*M*N
    res = parse_matrix_expression(expr)
    rexpr = recognize_matrix_expression(res)
    assert expr == rexpr

    expr = 3*M + N*M.T*M + 4*k*N
    res = parse_matrix_expression(expr)
    rexpr = recognize_matrix_expression(res)
    assert expr == rexpr

    expr = Inverse(M)*N
    rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
    assert expr == rexpr

    expr = M**2
    rexpr = recognize_matrix_expression(parse_matrix_expression(expr))
    assert expr == rexpr

    expr = M*(2*N + 3*M)
    res = parse_matrix_expression(expr)
    rexpr = recognize_matrix_expression(res)
    assert expr == rexpr

    expr = Trace(M)
    result = CodegenArrayContraction(M, (0, 1))
    assert parse_matrix_expression(expr) == result

    expr = 3*Trace(M)
    result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M), (0, 1))
    assert parse_matrix_expression(expr) == result

    expr = 3*Trace(Trace(M) * M)
    result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
    assert parse_matrix_expression(expr) == result

    expr = 3*Trace(M)**2
    result = CodegenArrayContraction(CodegenArrayTensorProduct(3, M, M), (0, 1), (2, 3))
    assert parse_matrix_expression(expr) == result

    expr = HadamardProduct(M, N)
    result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, N), (0, 2), (1, 3))
    assert parse_matrix_expression(expr) == result

    expr = HadamardPower(M, 2)
    result = CodegenArrayDiagonal(CodegenArrayTensorProduct(M, M), (0, 2), (1, 3))
    assert parse_matrix_expression(expr) == result

    expr = M**2
    assert isinstance(expr, MatPow)
    assert parse_matrix_expression(expr) == CodegenArrayContraction(CodegenArrayTensorProduct(M, M), (1, 2))
Exemplo n.º 18
0
def _(expr: CodegenArrayDiagonal, x: Expr):
    dsubexpr = array_derive(expr.expr, x)
    rank_x = len(get_shape(x))
    diag_indices = [[j + rank_x for j in i] for i in expr.diagonal_indices]
    return CodegenArrayDiagonal(dsubexpr, *diag_indices)
Exemplo n.º 19
0
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()
Exemplo n.º 20
0
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)
Exemplo n.º 21
0
def test_recognize_diagonalized_vectors():

    a = MatrixSymbol("a", k, 1)
    b = MatrixSymbol("b", k, 1)
    A = MatrixSymbol("A", k, k)
    B = MatrixSymbol("B", k, k)
    C = MatrixSymbol("C", k, k)
    X = MatrixSymbol("X", k, k)
    x = MatrixSymbol("x", k, 1)
    I1 = Identity(1)
    I = Identity(k)

    # Check matrix recognition over trivial dimensions:

    cg = CodegenArrayTensorProduct(a, b)
    assert recognize_matrix_expression(cg) == a * b.T

    cg = CodegenArrayTensorProduct(I1, a, b)
    assert recognize_matrix_expression(cg) == a * I1 * b.T

    # Recognize trace inside a tensor product:

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, B, C), (0, 3),
                                 (1, 2))
    assert recognize_matrix_expression(cg) == Trace(A * B) * C

    # Transform diagonal operator to contraction:

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (1, 2))
    assert cg.transform_to_product() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a)), (1, 2))
    assert recognize_matrix_expression(cg) == A * DiagMatrix(a)

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(a, b), (0, 2))
    assert cg.transform_to_product() == CodegenArrayContraction(
        CodegenArrayTensorProduct(DiagMatrix(a), b), (0, 2))
    assert recognize_matrix_expression(cg).doit() == DiagMatrix(a) * b

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(A, a), (0, 2))
    assert cg.transform_to_product() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a)), (0, 2))
    assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a)

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, I1), (0, 2),
                              (3, 5))
    assert cg.transform_to_product() == CodegenArrayContraction(
        CodegenArrayTensorProduct(I, DiagMatrix(x), I1), (0, 2))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(I, x, A, B), (1, 2),
                              (5, 6))
    assert cg.transform_to_product() == CodegenArrayDiagonal(
        CodegenArrayContraction(
            CodegenArrayTensorProduct(I, DiagMatrix(x), A, B), (1, 2)), (3, 4))

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I1), (1, 2))
    assert isinstance(cg, CodegenArrayDiagonal)
    assert cg.diagonal_indices == ((1, 2), )
    assert recognize_matrix_expression(cg) == x

    cg = CodegenArrayDiagonal(CodegenArrayTensorProduct(x, I), (0, 2))
    assert cg.transform_to_product() == CodegenArrayContraction(
        CodegenArrayTensorProduct(DiagMatrix(x), I), (0, 2))
    assert recognize_matrix_expression(cg).doit() == DiagMatrix(x)

    cg = CodegenArrayDiagonal(x, (1, ))
    assert cg == x

    # Ignore identity matrices with contractions:

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, A, I, I), (0, 2),
                                 (1, 3), (5, 7))
    assert cg.split_multiple_contractions() == cg
    assert recognize_matrix_expression(cg) == Trace(A) * I

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(Trace(A) * I, I, I),
                                 (1, 5), (3, 4))
    assert cg.split_multiple_contractions() == cg
    assert recognize_matrix_expression(cg).doit() == Trace(A) * I

    # Add DiagMatrix when required:

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a), (1, 2))
    assert cg.split_multiple_contractions() == cg
    assert recognize_matrix_expression(cg) == A * a

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (1, 2, 4))
    assert cg.split_multiple_contractions() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a), B), (1, 2), (3, 4))
    assert recognize_matrix_expression(cg) == A * DiagMatrix(a) * B

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, B), (0, 2, 4))
    assert cg.split_multiple_contractions() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a), B), (0, 2), (3, 4))
    assert recognize_matrix_expression(cg) == A.T * DiagMatrix(a) * B

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, b, a.T, B),
                                 (0, 2, 4, 7, 9))
    assert cg.split_multiple_contractions() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a), DiagMatrix(b),
                                  DiagMatrix(a), B), (0, 2), (3, 4), (5, 7),
        (6, 9))
    assert recognize_matrix_expression(
        cg).doit() == A.T * DiagMatrix(a) * DiagMatrix(b) * DiagMatrix(a) * B.T

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(I1, I1, I1),
                                 (1, 2, 4))
    assert cg.split_multiple_contractions() == CodegenArrayContraction(
        CodegenArrayTensorProduct(I1, I1, I1), (1, 2), (3, 4))
    assert recognize_matrix_expression(cg).doit() == Identity(1)

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(I, I, I, I, A),
                                 (1, 2, 8), (5, 6, 9))
    assert recognize_matrix_expression(
        cg.split_multiple_contractions()).doit() == A

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
                                 (1, 2, 4), (5, 6, 8))
    assert cg.split_multiple_contractions() == CodegenArrayContraction(
        CodegenArrayTensorProduct(A, DiagMatrix(a), C, DiagMatrix(a), B),
        (1, 2), (3, 4), (5, 6), (7, 8))
    assert recognize_matrix_expression(
        cg) == A * DiagMatrix(a) * C * DiagMatrix(a) * B

    cg = CodegenArrayContraction(
        CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
        (1, 2, 8), (5, 6, 9))
    assert cg.split_multiple_contractions().dummy_eq(
        CodegenArrayContraction(
            CodegenArrayTensorProduct(a, I1, b, I1, (a.T * b).applyfunc(cos)),
            (1, 2), (3, 8), (5, 6), (7, 9)))
    assert recognize_matrix_expression(cg).dummy_eq(
        MatMul(a, I1, (a.T * b).applyfunc(cos), Transpose(I1), b.T))

    cg = CodegenArrayContraction(
        CodegenArrayTensorProduct(A.T, a, b, b.T, (A * X * b).applyfunc(cos)),
        (1, 2, 8), (5, 6, 9))
    assert cg.split_multiple_contractions().dummy_eq(
        CodegenArrayContraction(
            CodegenArrayTensorProduct(A.T, DiagMatrix(a), b, b.T,
                                      (A * X * b).applyfunc(cos)), (1, 2),
            (3, 8), (5, 6, 9)))
    # assert recognize_matrix_expression(cg)

    # Check no overlap of lines:

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(A, a, C, a, B),
                                 (1, 2, 4), (5, 6, 8), (3, 7))
    assert cg.split_multiple_contractions() == cg

    cg = CodegenArrayContraction(CodegenArrayTensorProduct(a, b, A), (0, 2, 4),
                                 (1, 3))
    assert cg.split_multiple_contractions() == cg
Exemplo n.º 22
0
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()