kronecker_product_test.py 文件源码

def test_kronecker_product():
    matrix_list = []
    matrix1 = torch.Tensor([
        [1, 2, 3],
        [4, 5, 6],
    ])
    matrix2 = torch.Tensor([
        [1, 2],
        [4, 3],
    ])
    matrix_list.append(matrix1)
    matrix_list.append(matrix2)
    res = kronecker_product(matrix_list)

    actual = torch.Tensor([
        [1, 2, 2, 4, 3, 6],
        [4, 3, 8, 6, 12, 9],
        [4, 8, 5, 10, 6, 12],
        [16, 12, 20, 15, 24, 18]
    ])

    assert(torch.equal(res, actual))