In [1]:
A = [[1, 2,3 ]]
B = [[1, 2], [3, 4], [5, 6]]

def shape(A):
    num_rows = len(A)
    num_cols = len(A[0] if A else 0)
    return num_rows, num_cols

print(shape(A))
print(shape(B))


(1, 3)
(3, 2)

In [8]:
def get_row(A, i):
    return A[i]


def get_col(A, i):
    return [A_i[j] for A_i in A]


def make_matrix(num_rows, num_cols, entry_fn):
    return [[entry_fn(i, j) for j in range(num_cols)] for i in range(num_rows)]


def is_diagonal(i, j):
    return 1 if i == j else 0
            

identity_matrix = make_matrix(5, 5, is_diagonal)

from pprint import pprint
pprint(identity_matrix)
pprint(shape(identity_matrix))


[[1, 0, 0, 0, 0],
 [0, 1, 0, 0, 0],
 [0, 0, 1, 0, 0],
 [0, 0, 0, 1, 0],
 [0, 0, 0, 0, 1]]
(5, 5)

In [ ]: