python矩阵乘法

# 矩阵乘法
def mul_matrix(A, B):
    A_rows = len(A)
    A_cols = len(A[0])
    B_rows = len(B)
    B_cols = len(B[0])

    if A_cols != B_rows:
        raise Exception('Matrix A columns must equal to Matrix B rows!')
    res = []
    for i in range(A_rows):
        temp = []
        for k in range(B_cols):
            s = 0
            for j in range(A_cols):
                s += A[i][j] * B[j][k]
            temp.append(s)
        res.append(temp)
    return res


if __name__ == '__main__':
    a = [[1], [2]]
    b = [[3, 3]]
    print(mul_matrix(a, b))