Matrix operation in python

matrix_ops.py

A = [[1, 2],
     [3, 4]] # 2 x 2

B = [[5, 6],
     [7, 8]] # 2 x 2

D = [[5, 6],
     [7, 8],
     [9, 10]] # 2 x 3
     
def add_matrix(A, B):
    if len(A) != len(B) or len(A[0]) != len(B[0]):
        raise ValueError("Matrices must be the same size")
    return [[A[i][j] + B[i][j] for j in range(len(A[0]))] for i in range(len(A))]

def scalar_product(x, A):
    return [[x * A[i][j] for j in range(len(A[0]))] for i in range(len(A))]

def matrix_product(A, B):
    if len(A[0]) != len(B):
        raise ValueError("Number of columns in A must be equal to the number of rows in B")
    return [[sum(A[i][k] * B[k][j] for k in range(len(A[0]))) for j in range(len(B[0]))] for i in range(len(A))]
    
def transpose(A):
    return [[A[j][i] for j in range(len(A))] for i in range(len(A[0]))]

C = add_matrix(A, B)
print(C)  # Output: [[6, 8], [10, 12]]

x = 2
C = scalar_product(x, A)
print(C)  # Output: [[2, 4], [6, 8]]

C = matrix_product(A, D)
print(C)  # Output: [[47, 52], [109, 122]]


D_transpose = transpose(D)
print(D_transpose)  # Output: [[5, 7, 9], [6, 8, 10]]