Итерация по массивам
Объект итератора nditer, представленный в NumPy 1.6, предоставляет множество гибких способов посещения всех элементов одного или нескольких массивов систематическим образом. Эта страница представляет некоторые базовые способы использования объекта для вычислений над массивами в Python, а затем завершается тем, как можно ускорить внутренний цикл в Cython. Поскольку Python-интерфейс nditer является относительно прямым отображением API итератора C-массивов, эти идеи также помогут при работе с итерацией массивов из C или C++.
Итерация по одному массиву
Наиболее базовая задача, которую можно выполнить с nditer, — посещение каждого элемента массива. Каждый элемент предоставляется по одному с использованием стандартного интерфейса Python-итератора.
Пример
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a): ... print x, ... 0 1 2 3 4 5
Важно помнить, что порядок выбирается для соответствия расположению массива в памяти, а не с использованием стандартного порядка C или Fortran. Это делается для повышения эффективности доступа, отражая идею, что по умолчанию пользователь просто хочет посетить каждый элемент без учета определенного порядка. Мы можем видеть это, проходя по транспонированному массиву по сравнению с копированием этого транспонированного массива в порядке C.
Пример
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a.T): ... print x, ... 0 1 2 3 4 5
>>> for x in np.nditer(a.T.copy(order='C')): ... print x, ... 0 3 1 4 2 5
Элементы как a , так и a.T проходятся в одном и том же порядке, а именно в порядке, в котором они хранятся в памяти, тогда как элементы a.T.copy(order=’C’) проходятся в другом порядке, потому что они были размещены в другом расположении в памяти.
Управление порядком итерации
Иногда важно посещать элементы массива в определенном порядке, независимо от расположения элементов в памяти. Объект nditer предоставляет параметр order, чтобы управлять этим аспектом итерации. По умолчанию, поведение, описанное выше, задается order=’K’ для сохранения существующего порядка. Его можно переопределить с order=’C’ для порядка C и order=’F’ для порядка Fortran.
Пример
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, order='F'): ... print x, ... 0 3 1 4 2 5 >>> for x in np.nditer(a.T, order='C'): ... print x, ... 0 3 1 4 2 5
Изменение значений массива
По умолчанию nditer обрабатывает входной массив как объект только для чтения. Для изменения элементов массива необходимо указать режим чтения-записи или только записи. Это контролируется флагами для каждого операнда.
Обычная присваивание в Python просто изменяет ссылку в локальном или глобальном словаре переменных вместо изменения существующей переменной на месте. Это означает, что простое присваивание x не поместит значение в элемент массива, а вместо этого изменит x с ссылки на элемент массива на ссылку на присвоенное значение. Для фактического изменения элемента массива необходимо индексировать x с помощью многоточия.
Пример
>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
[3, 4, 5]])
>>> for x in np.nditer(a, op_flags=['readwrite']):
... x[...] = 2 * x
...
>>> a
array([[ 0, 2, 4],
[ 6, 8, 10]])
Использование внешнего цикла
Во всех предыдущих примерах элементы a предоставляются итератором по одному, потому что вся логика цикла находится внутри итератора. Хотя это просто и удобно, это не очень эффективно. Лучший подход — переместить одномерный внутренний цикл в свой код, внешний по отношению к итератору. Таким образом, векторизованные операции NumPy могут использоваться на больших кусках посещаемых элементов.
nditer попытается предоставить куски наибольшего размера внутреннему циклу. Принудительно задавая порядок ‘C’ и ‘F’, мы получим разные размеры внешнего цикла. Этот режим включается путем задания флага итератора.
Обратите внимание, что при сохранении родного порядка памяти итератор может предоставить один одномерный кусок, тогда как при принудительном задании порядка Fortran, ему приходится предоставлять три куска по два элемента каждый.
Пример
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, flags=['external_loop']): ... print x, ... [0 1 2 3 4 5]
>>> for x in np.nditer(a, flags=['external_loop'], order='F'): ... print x, ... [0 3] [1 4] [2 5]
Отслеживание индекса или многомерного индекса
Во время итерации вы можете использовать индекс текущего элемента в вычислении. Например, вы можете посещать элементы массива в порядке памяти, но использовать индекс порядка C, Fortran или многомерный индекс для поиска значений в другом массиве.
Протокол Python-итератора не имеет естественного способа запроса этих дополнительных значений от итератора, поэтому мы вводим альтернативную синтаксическую конструкцию для итерации с nditer. Этот синтаксис явно работает с самим объектом итератора, поэтому его свойства доступны в ходе итерации. С этой конструкцией цикла текущее значение доступно по индексу в итераторе, а отслеживаемый индекс — это свойство index или multi_index, в зависимости от запрошенного.
К сожалению, интерактивный интерпретатор Python выводит значения выражений внутри цикла while во время каждой итерации. В примерах, использующих эту конструкцию цикла, мы изменили вывод, чтобы он был более удобочитаемым.
Пример
>>> a = np.arange(6).reshape(2,3) >>> it = np.nditer(a, flags=['f_index']) >>> while not it.finished: ... print "%d <%d>" % (it[0], it.index), ... it.iternext() ... 0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>
>>> it = np.nditer(a, flags=['multi_index']) >>> while not it.finished: ... print "%d <%s>" % (it[0], it.multi_index), ... it.iternext() ... 0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> it = np.nditer(a, flags=['multi_index'], op_flags=['writeonly'])
>>> while not it.finished:
... it[0] = it.multi_index[1] - it.multi_index[0]
... it.iternext()
...
>>> a
array([[ 0, 1, 2],
[-1, 0, 1]])
Отслеживание индекса или многомерного индекса несовместимо с использованием внешнего цикла, так как требуется разное значение индекса на каждый элемент. Если вы попытаетесь объединить эти флаги, объект nditer вызовет исключение.
Пример
>>> a = np.zeros((2,3)) >>> it = np.nditer(a, flags=['c_index', 'external_loop']) Traceback (most recent call last): File "<stdin>", line 1, in <module> ValueError: Iterator flag EXTERNAL_LOOP cannot be used if an index or multi-index is being tracked
Буферизация элементов массива
Принудительное задание порядка итерации показало, что опция внешнего цикла может предоставлять элементы меньшими кусками, поскольку элементы не могут быть посещены в соответствующем порядке с постоянным шагом. При написании кода C это обычно нормально, однако в чистом Python-коде это может привести к значительному снижению производительности.
Включением буферизации можно увеличить размер кусков, предоставляемых итератором внутреннему циклу, значительно уменьшив накладные расходы интерпретатора Python. В примере принудительного задания порядка итерации Fortran внутренний цикл получает все элементы сразу при включенной буферизации.
Пример
>>> a = np.arange(6).reshape(2,3) >>> for x in np.nditer(a, flags=['external_loop'], order='F'): ... print x, ... [0 3] [1 4] [2 5]
>>> for x in np.nditer(a, flags=['external_loop','buffered'], order='F'): ... print x, ... [0 3 1 4 2 5]
Итерация как определённого типа данных
Иногда необходимо обращаться к массиву как к типу данных, отличному от того, в котором он хранится. Например, можно выполнить все вычисления с плавающей точкой двойной точности (64-бит), даже если обрабатываемые массивы имеют плавающую точность одинарной точности (32-бит). За исключением написания низкоуровневого кода C, лучше позволить итератору выполнить копирование или буферизацию, вместо того, чтобы самим преобразовывать тип данных во внутреннем цикле.
Есть два механизма, которые позволяют это сделать: временные копии и режим буферизации. С временными копиями создается копия всего массива с новым типом данных, затем итерация выполняется в этой копии. Доступ для записи разрешен с помощью режима, обновляющего исходный массив после завершения всей итерации. Основной недостаток временных копий заключается в том, что временная копия может потреблять большое количество памяти, особенно если тип данных итерации имеет больший размер элемента, чем исходный.
Режим буферизации устраняет проблему использования памяти и более дружественный к кэшу, чем создание временных копий. За исключением особых случаев, когда весь массив нужен сразу за пределами итератора, буферизация предпочтительнее временного копирования. В NumPy буферизация используется функциями ufunc и другими функциями для поддержки гибких входных данных с минимальными затратами памяти.
В наших примерах мы будем обрабатывать входной массив с комплексным типом данных, чтобы мы могли извлекать квадратные корни из отрицательных чисел. Без включения режима копирования или буферизации итератор вызовет исключение, если тип данных не соответствует точно.
Пример
>>> a = np.arange(6).reshape(2,3) - 3 >>> for x in np.nditer(a, op_dtypes=['complex128']): ... print np.sqrt(x), ... Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: Iterator operand required copying or buffering, but neither copying nor buffering was enabled
В режиме копирования флаг ‘copy’ задается для каждого операнда. Это делается для обеспечения контроля на уровне каждого операнда. Режим буферизации задается как флаг итератора.
Пример
>>> a = np.arange(6).reshape(2,3) - 3 >>> for x in np.nditer(a, op_flags=['readonly','copy'], ... op_dtypes=['complex128']): ... print np.sqrt(x), ... 1.73205080757j 1.41421356237j 1j 0j (1+0j) (1.41421356237+0j)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['complex128']): ... print np.sqrt(x), ... 1.73205080757j 1.41421356237j 1j 0j (1+0j) (1.41421356237+0j)
Итератор использует правила преобразования NumPy, чтобы определить, разрешено ли конкретное преобразование. По умолчанию он применяет преобразование ‘safe’. Это означает, например, что он вызовет исключение, если вы попытаетесь обработать массив с плавающей точкой двойной точности (64-бит) как массив с плавающей точкой одинарной точности (32-бит). Во многих случаях правило ‘same_kind’ является наиболее разумным, так как оно позволит преобразование из 64 в 32-битную плавающую точку, но не из плавающей точки в целое число или из комплексного в плавающую точку.
Пример
>>> a = np.arange(6.)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32']):
... print x,
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('float32') according to the rule 'safe'
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32'], ... casting='same_kind'): ... print x, ... 0.0 1.0 2.0 3.0 4.0 5.0
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['int32'], casting='same_kind'):
... print x,
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('int32') according to the rule 'same_kind'
Следует учитывать преобразования обратно к исходному типу данных при использовании операнда с режимом чтения-записи или только записи. Распространенный случай — реализация внутреннего цикла в терминах плавающей точки двойной точности (64-бит) и использование преобразования ‘same_kind’, чтобы позволить и другим типам с плавающей точкой быть обработаны. Хотя в режиме только чтения может быть предоставлен массив целых чисел, режим чтения-записи вызовет исключение, потому что преобразование обратно в массив нарушит правило преобразования.
Пример
>>> a = np.arange(6)
>>> for x in np.nditer(a, flags=['buffered'], op_flags=['readwrite'],
... op_dtypes=['float64'], casting='same_kind'):
... x[...] = x / 2.0
...
Traceback (most recent call last):
File "<stdin>", line 2, in <module>
TypeError: Iterator requested dtype could not be cast from dtype('float64') to dtype('int64'), the operand 0 dtype, according to the rule 'same_kind'
Итерация по массивам с векторизованными операциями
NumPy имеет набор правил для обработки массивов с различными форматами, которые применяются всякий раз, когда функции принимают несколько операндов, которые объединяются поэлементно. Это называется векторизацией. Объект nditer может применять эти правила за вас, когда вам нужно написать такую функцию.
В качестве примера мы выводим результат векторизованного объединения одномерного и двумерного массивов.
Пример
>>> a = np.arange(3) >>> b = np.arange(6).reshape(2,3) >>> for x, y in np.nditer([a,b]): ... print "%d:%d" % (x,y), ... 0:0 1:1 2:2 0:3 1:4 2:5
При возникновении ошибки векторизации итератор вызовет исключение, которое включает формы входных данных, чтобы помочь диагностировать проблему.
Пример
>>> a = np.arange(2) >>> b = np.arange(6).reshape(2,3) >>> for x, y in np.nditer([a,b]): ... print "%d:%d" % (x,y), ... Traceback (most recent call last): File "<stdin>", line 1, in <module> ValueError: operands could not be broadcast together with shapes (2) (2,3)
Выделение выходных массивов итератором
В функциях NumPy распространен случай, когда выходные массивы выделены на основе векторизации входных данных, а также есть необязательный параметр «out», где результат будет размещен, если он предоставлен. Объект nditer предоставляет удобный idiom, который делает этот механизм очень простым.
Мы покажем, как это работает, создав функцию square, которая возводит свой вход в квадрат. Начнем с минимальной функции без поддержки параметра «out».
Пример
>>> def square(a): ... it = np.nditer([a, None]) ... for x, y in it: ... y[...] = x*x ... return it.operands[1] ... >>> square([1,2,3]) array([1, 4, 9])
По умолчанию nditer использует флаги ‘allocate’ и ‘writeonly’ для операндов, которые переданы как None. Это означает, что мы смогли передать только два операнда итератору, а он справился со всем остальным.
При добавлении параметра ‘out’, мы должны явно указать эти флаги, потому что если кто-то передаёт массив как ‘out’, итератор по умолчанию будет ‘только чтение’, и наш внутренний цикл завершится ошибкой. Причина, по которой ‘только чтение’ является значением по умолчанию для входных массивов, заключается в предотвращении путаницы с непреднамеренным вызовом операции сокращения. Если бы значением по умолчанию было ‘чтение/запись’, любая операция трансляции также вызывала бы сокращение, тема, которая рассматривается позже в этом документе.
Пока мы здесь, давайте также введём флаг ‘no_broadcast’, который предотвратит трансляцию вывода. Это важно, потому что мы хотим получить только одно входное значение для каждого выходного. Агрегирование более одного входного значения — это операция сокращения, которая требует специальной обработки. Она уже вызывала бы ошибку, потому что сокращения должны быть явно разрешены во флагах итератора, но сообщение об ошибке, возникающее при отключении трансляции, намного понятнее для конечных пользователей. Чтобы увидеть, как обобщить функцию квадрата на сокращение, посмотрите на функцию суммы квадратов в разделе о Cython.
Для полноты мы также добавим флаги ‘external_loop’ и ‘buffered’, так как они обычно требуются для производительности.
Пример
>>> def square(a, out=None): ... it = np.nditer([a, out], ... flags = ['external_loop', 'buffered'], ... op_flags = [['readonly'], ... ['writeonly', 'allocate', 'no_broadcast']]) ... for x, y in it: ... y[...] = x*x ... return it.operands[1] ...
>>> square([1,2,3]) array([1, 4, 9])
>>> b = np.zeros((3,)) >>> square([1,2,3], out=b) array([ 1., 4., 9.]) >>> b array([ 1., 4., 9.])
>>> square(np.arange(6).reshape(2,3), out=b) Traceback (most recent call last): File "<stdin>", line 1, in <module> File "<stdin>", line 4, in square ValueError: non-broadcastable output operand with shape (3) doesn't match the broadcast shape (2,3)
Итерация внешнего произведения
Любая бинарная операция может быть расширена до операции над массивом в виде внешнего произведения, как в outer, и объект nditer предоставляет способ достижения этого путём явного сопоставления осей операндов. Это также возможно с индексацией newaxis, но мы покажем вам, как напрямую использовать параметр nditer op_axes для достижения этого без промежуточных представлений.
Мы выполним простое внешнее произведение, разместив размеры первого операнда перед размерами второго операнда. Параметр op_axes требует одного списка осей для каждого операнда и предоставляет отображение осей итератора на оси операнда.
Предположим, что первый операнд одномерный, а второй — двумерный. Итератор будет иметь три измерения, поэтому op_axes будет содержать два списка из 3 элементов. Первый список выбирает одну ось первого операнда и имеет значение -1 для остальных осей итератора, что в итоге даёт [0, -1, -1]. Второй список выбирает две оси второго операнда, но не должен перекрываться с осями, выбранными в первом операнде. Его список — [-1, 0, 1]. Операнд вывода отображается на оси итератора стандартным образом, поэтому мы можем указать None вместо построения другого списка.
Операция во внутреннем цикле — простое умножение. Всё, что связано с внешним произведением, обрабатывается настройкой итератора.
Пример
>>> a = np.arange(3)
>>> b = np.arange(8).reshape(2,4)
>>> it = np.nditer([a, b, None], flags=['external_loop'],
... op_axes=[[0, -1, -1], [-1, 0, 1], None])
>>> for x, y, z in it:
... z[...] = x*y
...
>>> it.operands[2]
array([[[ 0, 0, 0, 0],
[ 0, 0, 0, 0]],
[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],
[[ 0, 2, 4, 6],
[ 8, 10, 12, 14]]])
Итерация сокращения
Всякий раз, когда операнд с возможностью записи имеет меньше элементов, чем полное пространство итерации, этот операнд подвергается сокращению. Объект nditer требует, чтобы любой операнд сокращения был помечен как «чтение/запись», и допускает сокращения только при указании флага итератора ‘reduce_ok’.
Рассмотрим простой пример: вычисление суммы всех элементов в массиве.
Пример
>>> a = np.arange(24).reshape(2,3,4) >>> b = np.array(0) >>> for x, y in np.nditer([a, b], flags=['reduce_ok', 'external_loop'], ... op_flags=[['readonly'], ['readwrite']]): ... y[...] += x ... >>> b array(276) >>> np.sum(a) 276
Дело становится немного сложнее при объединении сокращения и выделенных операндов. Перед началом итерации любой операнд сокращения должен быть инициализирован своими начальными значениями. Вот как мы можем это сделать, суммируя по последней оси a.
Пример
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok', 'external_loop'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> it.operands[1][...] = 0
>>> for x, y in it:
... y[...] += x
...
>>> it.operands[1]
array([[ 6, 22, 38],
[54, 70, 86]])
>>> np.sum(a, axis=2)
array([[ 6, 22, 38],
[54, 70, 86]])
Для буферизованного сокращения требуется ещё одна корректировка во время настройки. Обычно построение итератора включает копирование первого буфера данных из читаемых массивов в буфер. Любой операнд сокращения является читаемым, поэтому он может быть помещён в буфер. К сожалению, инициализация операнда после завершения этой операции буферизации не будет отражена в буфере, с которым начинается итерация, и будут получены ошибочные результаты.
Флаг итератора «delay_bufalloc» предназначен для того, чтобы разрешить операндам сокращения, выделенным итератором, существовать вместе с буферизацией. Когда этот флаг установлен, итератор оставит свои буферы неинициализированными до получения сброса, после чего он будет готов к обычной итерации. Вот как предыдущий пример выглядит, если мы также включим буферизацию.
Пример
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok', 'external_loop',
... 'buffered', 'delay_bufalloc'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> it.operands[1][...] = 0
>>> it.reset()
>>> for x, y in it:
... y[...] += x
...
>>> it.operands[1]
array([[ 6, 22, 38],
[54, 70, 86]])
Размещение внутреннего цикла в Cython
Те, кто хочет действительно хорошей производительности своих операций низкого уровня, должны серьёзно рассмотреть использование API итерации, предоставляемого на C, но для тех, кто не знаком с C или C++, Cython является хорошей промежуточной ступенью с разумными компромиссами по производительности. Для объекта nditer это означает, что итератор отвечает за трансляцию, преобразование типов и буферизацию, а внутренний цикл предоставляется Cython.
В качестве примера мы создадим функцию суммы квадратов. Для начала давайте реализуем эту функцию в обычном Python. Мы хотим поддерживать параметр ‘axis’, аналогичный функции numpy sum, поэтому нам понадобится создать список для параметра op_axes. Вот как это выглядит.
Пример
>>> def axis_to_axeslist(axis, ndim): ... if axis is None: ... return [-1] * ndim ... else: ... if type(axis) is not tuple: ... axis = (axis,) ... axeslist = [1] * ndim ... for i in axis: ... axeslist[i] = -1 ... ax = 0 ... for i in range(ndim): ... if axeslist[i] != -1: ... axeslist[i] = ax ... ax += 1 ... return axeslist ... >>> def sum_squares_py(arr, axis=None, out=None): ... axeslist = axis_to_axeslist(axis, arr.ndim) ... it = np.nditer([arr, out], flags=['reduce_ok', 'external_loop', ... 'buffered', 'delay_bufalloc'], ... op_flags=[['readonly'], ['readwrite', 'allocate']], ... op_axes=[None, axeslist], ... op_dtypes=['float64', 'float64']) ... it.operands[1][...] = 0 ... it.reset() ... for x, y in it: ... y[...] += x*x ... return it.operands[1] ... >>> a = np.arange(6).reshape(2,3) >>> sum_squares_py(a) array(55.0) >>> sum_squares_py(a, axis=-1) array([ 5., 50.])
Для преобразования этой функции в Cython мы заменяем внутренний цикл (y[...] += x*x) кодом Cython, специализированным для типа float64. При включённом флаге ‘external_loop’ массивы, предоставляемые внутреннему циклу, всегда будут одномерными, поэтому требуется очень мало проверок.
Вот список sum_squares.pyx:
import numpy as np
cimport numpy as np
cimport cython
def axis_to_axeslist(axis, ndim):
if axis is None:
return [-1] * ndim
else:
if type(axis) is not tuple:
axis = (axis,)
axeslist = [1] * ndim
for i in axis:
axeslist[i] = -1
ax = 0
for i in range(ndim):
if axeslist[i] != -1:
axeslist[i] = ax
ax += 1
return axeslist
@cython.boundscheck(False)
def sum_squares_cy(arr, axis=None, out=None):
cdef np.ndarray[double] x
cdef np.ndarray[double] y
cdef int size
cdef double value
axeslist = axis_to_axeslist(axis, arr.ndim)
it = np.nditer([arr, out], flags=['reduce_ok', 'external_loop',
'buffered', 'delay_bufalloc'],
op_flags=[['readonly'], ['readwrite', 'allocate']],
op_axes=[None, axeslist],
op_dtypes=['float64', 'float64'])
it.operands[1][...] = 0
it.reset()
for xarr, yarr in it:
x = xarr
y = yarr
size = x.shape[0]
for i in range(size):
value = x[i]
y[i] = y[i] + value * value
return it.operands[1]
На этом компьютере построение файла .pyx в модуль выглядело следующим образом, но вам может потребоваться найти некоторые учебные материалы по Cython, чтобы узнать подробности для вашей конфигурации системы.
$ cython sum_squares.pyx $ gcc -shared -pthread -fPIC -fwrapv -O2 -Wall -I/usr/include/python2.7 -fno-strict-aliasing -o sum_squares.so sum_squares.c
Запуск этого из интерпретатора Python даёт те же ответы, что и наш исходный Python/NumPy код.
Пример
>>> from sum_squares import sum_squares_cy >>> a = np.arange(6).reshape(2,3) >>> sum_squares_cy(a) array(55.0) >>> sum_squares_cy(a, axis=-1) array([ 5., 50.])
Небольшие измерения времени в IPython показывают, что снижение накладных расходов и выделение памяти внутреннего цикла Cython обеспечивают значительное ускорение по сравнению как с простым кодом Python, так и с выражением, использующим встроенную функцию sum в NumPy.
>>> a = np.random.rand(1000,1000) >>> timeit sum_squares_py(a, axis=-1) 10 loops, best of 3: 37.1 ms per loop >>> timeit np.sum(a*a, axis=-1) 10 loops, best of 3: 20.9 ms per loop >>> timeit sum_squares_cy(a, axis=-1) 100 loops, best of 3: 11.8 ms per loop >>> np.all(sum_squares_cy(a, axis=-1) == np.sum(a*a, axis=-1)) True >>> np.all(sum_squares_py(a, axis=-1) == np.sum(a*a, axis=-1)) True
© 2008–2016 NumPy Developers
Licensed under the NumPy License.
https://docs.scipy.org/doc/numpy-1.11.0/reference/arrays.nditer.html
