Whoever wrote this function was clearly in a hurry. They forgot to finish the assertion and completely neglected to test it, or they surely would have found the bug. Be kind and help them out. def trace(A): "Return trace of a square matrix."" #todo: put in assert statemtn to check A is really square matrix! return sum( [A[i] for i in range(A. shape[0])])