Как правильно транспонировать матрицу

Как правильно сделать функцию из моего кода и чтобы она работала полностью? Код написал, но он не работает с некоторыми входными данными. А именно при нулях (строка = 0, колонка = 0) и при очень больших матрицах(в тестах валится на 950x764 матрице). Функцию, которую пытался написать и код - ниже. Помогите пожалуйста довести до ума.

Условие: Есть матрица размера m × n. Нужно написать функцию, которая её транспонирует. Транспонированная матрица получается из исходной заменой строк на столбцы.

Формат ввода: В первой строке задано число n — количество строк матрицы. Во второй строке задано m — число столбцов, m и n не превосходят 1000. В следующих n строках задана матрица. Числа в ней не превосходят по модулю 1000.

Формат вывода: Напечатайте транспонированную матрицу в том же формате, который задан во входных данных. Каждая строка матрицы выводится на отдельной строке, элементы разделяются пробелами.

Пример вводы/вывода:

Ввод

4

3

1 2 3

0 2 6

7 4 1

2 7 0

Вывод

1 0 7 2

2 2 4 7

3 6 1 0

Просто код:

rows = int(input().strip())
colums = int(input().strip())
a = [[0]*colums for _ in range(rows)]
for i in range(rows):
    a[i] = [int(j) for j in input().strip().split(" ")]
rows_count = len(a)
colums_count = len(a[0])
new_matrix = [[0] * rows_count for _ in range(colums_count)]
for i in range(rows_count):
    for j in range(colums_count):
        new_matrix[j][i] = a[i][j]
for row in new_matrix:
    print(*row)

Пытался создать функцию из кода выше, но безуспешно. Много ошибок:

def transpose(rows, colums):

    rows_count = len(a)
    colums_count = len(a[0])
    new_matrix = [[0] * rows_count for _ in range(colums_count)]
    for i in range(rows_count):
        for j in range(colums_count):
            new_matrix[j][i] = a[i][j]
    for row in new_matrix:
        print(*row)


if __name__ == '__main__':
    rows = int(input().strip())
    colums = int(input().strip())
    a = [[0] * colums for _ in range(rows)]
    for i in range(rows):
        a[i] = [int(j) for j in input().strip().split(" ")]
    print(transpose(rows, colums))

Правка после комментариев Sergey и A_Vaclav

def transpose(a):
    for i in range(rows):
        a[i] = [int(j) for j in input().strip().split(' ')]
    rows_count = len(a)
    colums_count = len(a[0])
    new_matrix = [[0] * rows_count for _ in range(colums_count)]
    for i in range(rows_count):
        for j in range(colums_count):
            new_matrix[j][i] = a[i][j]
    for row in new_matrix:
        print(*row)


if __name__ == '__main__':
    rows = int(input().strip())
    colums = int(input().strip())
    a = [colums for _ in range(rows)]
    print(transpose(a))

Результат работы функции:

Ввод:

4

3

1 2 3

2 4 5

6 78 8

6 8 9

Вывод:

1 2 6 6

2 4 78 8

3 5 8 9

None

Проверка случая, когда на вход подаются нули:

def transpose(a):
    rows_count = len(a)
    colums_count = len(a)
    if rows_count == colums_count == 0:
        print('None')
    new_matrix = []
    for j in range(colums_count):
        tmp = []
        for i in range(rows_count):
            tmp.append(a[i][j])

        new_matrix.append(tmp)
    return new_matrix


if __name__ == '__main__':
    rows = int(input().strip())
    colums = int(input().strip())
    a = [[0] * colums for _ in range(rows)]
    for i in range(rows):
        a[i] = [int(j) for j in input().strip().split(' ')]

    # print(transpose(a))
    for row in transpose(a):
        if row != 0:
            print(*row)

Входные

0

0

Выходные

None

Входные:

4

5

4 5 6 7 8

5 6 7 7 8

6 7 8 8 9

8 9 6 4 3

Выходные: (выводит без последней колонки)

4 5 6 8

5 6 7 9

6 7 8 6

7 7 8 4

Если я делаю так colums_count = len(a[0]), то матрица транспонируется нормально, но при нулях ошибка! NoneType is not iterable. Пытался в конце под main() задать условие в цикле, что вывод только если row != 0, но это не помогло