Определение ранга матрицы без встроенных функций на Python

Пишу алгоритм поиска ранга матрицы без использования встроенной функции в numpy, но пока что не выходит. Помогите, пожалуйста, найти ошибку. Код откомментировал, насколько это возможно.

import numpy as np
a = np.array([[1, 1, 1, 1],
              [0, 0, 1, 1],
              [0, 0, 1, 1],
              [0, 0, 1, 1]])
#Копия исходной матрицы
b = np.copy(a)
#Количество нулевых строк
count = 0
#Вспомогательная переменная - индикатор,
#была ли найдена строка с ненулевым элементом
key = True
#Проходим по строкам матрицы и приводим матрицу к трапецевидной
for k in range(a.shape[0]):
      #Если элемент на главной диагонали в исходной
      #строке - нуль, то ищем строку, где элемент
      #того же столбца не нулевой, и меняем строки
      #местами, или, если такую строку найти не удалось,
      #прерываем цикл
      if (b[k][k] == 0):
           key = False
           #Идём по строкам, расположенным ниже исходной
           for i in range(k+1, a.shape[0]):
                  #Если нашли строку, где в том же столбце имеется не нулевой элемент
                  if (b[i][k] != 0):
                        key = True
                        #Меняем найденную и исходную строки местами (функция самописная)
                        b = swap(b, k, i)
                        break
      #Если на каком-то этапе не нашли столбца с ненулевыми элементами
      #ниже k-ой позиции - прерываем цикл        
      if (key == False):
            break

      #Идём по строкам, которые расположены ниже исходной
      for i in range(k+1, a.shape[0]):
            
            multi = b[i][k] / b[k][k]

            for j in range(a.shape[1]):
                 
                 b[i][j] -= multi*b[k][j]
#Считаем количество нулевых строк
for i in range(a.shape[0]):
      key = False
      for j in range(a.shape[1]):
            if b[i][j] != 0:
              key = True
      if (key == False):
        count+=1
print(a.shape[0] - count)