迭代数组#
迭代器对象 nditer ,在 NumPy 1.6 中引入,提供了许多灵活的方法以系统的方式访问一个或多个数组的所有元素.本页介绍了一些在 Python 中使用该对象进行数组计算的基本方法,然后总结了如何在 Cython 中加速内部循环.由于 nditer 的 Python 接口是 C 数组迭代器 API 的一个相对直接的映射,这些想法也将有助于使用 C 或 C++ 进行数组迭代.
单个数组迭代#
使用 nditer 可以完成的最基本任务是访问数组的每个元素.每个元素使用标准的 Python 迭代器接口逐一提供.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a):
... print(x, end=' ')
...
0 1 2 3 4 5
对于这次迭代,需要注意的一个重要事项是,顺序的选择是为了匹配数组的内存布局,而不是使用标准的C或Fortran排序.这样做是为了访问效率,反映了一个默认情况下只想访问每个元素而不关心特定顺序的想法.我们可以通过迭代我们之前数组的转置来看到这一点,与以C顺序复制该转置进行比较.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a.T):
... print(x, end=' ')
...
0 1 2 3 4 5
>>> for x in np.nditer(a.T.copy(order='C')):
... print(x, end=' ')
...
0 3 1 4 2 5
a 和 a.T 的元素以相同的顺序遍历,即它们在内存中存储的顺序,而 a.T.copy(order=’C’) 的元素以不同的顺序访问,因为它们已经被放入了不同的内存布局中.
控制迭代顺序#
有时按特定顺序访问数组的元素很重要,而不考虑内存中元素的布局.:class:nditer 对象提供了一个 order 参数来控制迭代这一方面.默认情况下,行为如上所述,是 order=’K’ 以保持现有顺序.这可以通过 order=’C’ 覆盖为 C 顺序,通过 order=’F’ 覆盖为 Fortran 顺序.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, order='F'):
... print(x, end=' ')
...
0 3 1 4 2 5
>>> for x in np.nditer(a.T, order='C'):
... print(x, end=' ')
...
0 3 1 4 2 5
修改数组值#
默认情况下,:class:nditer 将输入操作数视为只读对象.要能够修改数组元素,必须使用每个操作数的 ‘readwrite’ 或 ‘writeonly’ 标志指定读写或只写模式.
nditer 随后将产生可写的缓冲区数组,您可以对其进行修改.然而,由于 nditer 必须在迭代完成后将这些缓冲区数据复制回原始数组,您必须通过两种方法之一在迭代结束时发出信号.您可以:
一旦调用了 close 或其上下文退出,nditer 就无法再迭代.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> a
array([[0, 1, 2],
[3, 4, 5]])
>>> with np.nditer(a, op_flags=['readwrite']) as it:
... for x in it:
... x[...] = 2 * x
...
>>> a
array([[ 0, 2, 4],
[ 6, 8, 10]])
如果你正在编写需要支持旧版本 numpy 的代码,请注意在 1.15 之前,:class:nditer 不是一个上下文管理器,也没有 close 方法.相反,它依赖析构函数来启动缓冲区的写回.
使用外部循环#
在目前为止的所有示例中,`a` 的元素是由迭代器一次一个地提供的,因为所有的循环逻辑都在迭代器内部.虽然这很简单方便,但效率不高.更好的方法是将一维的最内层循环移到代码中,外部于迭代器.这样,NumPy 的向量化操作可以用于访问更大的元素块.
nditer 将尝试向内部循环提供尽可能大的块.通过强制 ‘C’ 和 ‘F’ 顺序,我们得到不同的外部循环大小.这种模式通过指定一个迭代器标志来启用.
请注意,使用默认的原生内存顺序时,迭代器能够提供一个单一的一维块,而当强制使用 Fortran 顺序时,它必须提供三个每块包含两个元素的块.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop']):
... print(x, end=' ')
...
[0 1 2 3 4 5]
>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
... print(x, end=' ')
...
[0 3] [1 4] [2 5]
跟踪一个索引或多重索引#
在迭代过程中,您可能希望在计算中使用当前元素的索引.例如,您可能希望按内存顺序访问数组的元素,但使用 C 顺序、Fortran 顺序或多维索引来在不同的数组中查找值.
索引由迭代器对象本身跟踪,并且可以通过 index 或 multi_index 属性访问,具体取决于所请求的内容.下面的示例显示了打印输出,演示了索引的进展:
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> it = np.nditer(a, flags=['f_index'])
>>> for x in it:
... print("%d <%d>" % (x, it.index), end=' ')
...
0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>
>>> it = np.nditer(a, flags=['multi_index'])
>>> for x in it:
... print("%d <%s>" % (x, it.multi_index), end=' ')
...
0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it:
... for x in it:
... x[...] = it.multi_index[1] - it.multi_index[0]
...
>>> a
array([[ 0, 1, 2],
[-1, 0, 1]])
跟踪一个索引或多索引与使用外部循环不兼容,因为它需要每个元素不同的索引值.如果你尝试组合这些标志,:class:nditer 对象将引发异常.
示例
>>> import numpy as np
>>> 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
替代循环和元素访问#
为了在迭代过程中更方便地访问其属性,:class:nditer 有一种替代的迭代语法,它明确地使用迭代器对象本身.使用这种循环结构,可以通过索引访问当前值.其他属性,如跟踪的索引保持不变.下面的示例与上一节中的示例产生相同的结果.
示例
>>> import numpy as np
>>> 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), end=' ')
... is_not_finished = 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), end=' ')
... is_not_finished = it.iternext()
...
0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>
>>> with np.nditer(a, flags=['multi_index'], op_flags=['writeonly']) as it:
... while not it.finished:
... it[0] = it.multi_index[1] - it.multi_index[0]
... is_not_finished = it.iternext()
...
>>> a
array([[ 0, 1, 2],
[-1, 0, 1]])
缓冲数组元素#
当强制一个迭代顺序时,我们观察到外部循环选项可能会以较小的块提供元素,因为元素不能以恒定的步幅按适当的顺序访问.在编写C代码时,这通常是可以的,然而在纯Python代码中,这可能会导致性能显著下降.
通过启用缓冲模式,迭代器提供给内层循环的块可以变得更大,显著减少Python解释器的开销.在强制Fortran迭代顺序的示例中,当启用缓冲时,内层循环可以一次性看到所有元素.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3)
>>> for x in np.nditer(a, flags=['external_loop'], order='F'):
... print(x, end=' ')
...
[0 3] [1 4] [2 5]
>>> for x in np.nditer(a, flags=['external_loop','buffered'], order='F'):
... print(x, end=' ')
...
[0 3 1 4 2 5]
以特定数据类型进行迭代#
有时需要将数组视为与其存储类型不同的数据类型.例如,即使正在操作的数组是32位浮点数,也可能希望在所有计算中使用64位浮点数.除非编写低级C代码,否则通常最好让迭代器处理复制或缓冲,而不是在内层循环中自行转换数据类型.
有两种机制可以实现这一点,即临时副本和缓冲模式.使用临时副本时,会创建整个数组的副本并使用新的数据类型,然后在副本中进行迭代.通过一种模式允许写访问,该模式在所有迭代完成后更新原始数组.临时副本的主要缺点是临时副本可能会消耗大量内存,特别是当迭代数据类型的itemsize大于原始数据类型时.
缓冲模式缓解了内存使用问题,并且比制作临时副本更利于缓存.除了在迭代器外部需要整个数组的特殊情况外,建议使用缓冲而不是临时复制.在NumPy中,缓冲由ufuncs和其他函数使用,以支持具有最小内存开销的灵活输入.
在我们的示例中,我们将使用复杂数据类型的输入数组,以便我们可以对负数取平方根.如果不启用复制或缓冲模式,如果数据类型不精确匹配,迭代器将引发异常.
示例
>>> import numpy as np
>>> a = np.arange(6).reshape(2,3) - 3
>>> for x in np.nditer(a, op_dtypes=['complex128']):
... print(np.sqrt(x), end=' ')
...
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’ 被指定为每个操作数的标志.这样做是为了提供每个操作数的控制.缓冲模式被指定为迭代器标志.
示例
>>> import numpy as np
>>> 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), end=' ')
...
1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['complex128']):
... print(np.sqrt(x), end=' ')
...
1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)
迭代器使用 NumPy 的转换规则来确定是否允许特定转换.默认情况下,它强制执行 ‘安全’ 转换.这意味着,例如,如果您尝试将 64 位浮点数组视为 32 位浮点数组,它将引发异常.在许多情况下,规则 ‘same_kind’ 是最合理的规则,因为它允许从 64 位浮点到 32 位浮点的转换,但不允许从浮点到整数或从复数到浮点的转换.
示例
>>> import numpy as np
>>> a = np.arange(6.)
>>> for x in np.nditer(a, flags=['buffered'], op_dtypes=['float32']):
... print(x, end=' ')
...
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, end=' ')
...
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, end=' ')
...
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’ 转换来允许其他浮点类型也被处理.虽然在只读模式下可以提供一个整数数组,但在读写模式下会引发异常,因为转换回数组会违反转换规则.
示例
>>> import numpy as np
>>> 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 有一套处理不同形状数组的规则,这些规则在函数接受多个按元素组合的操作数时应用.这被称为 广播 .当你需要编写这样的函数时,:class:nditer 对象可以为你应用这些规则.
作为一个例子,我们打印出一个一维和二维数组一起广播的结果.
示例
>>> import numpy as np
>>> a = np.arange(3)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
... print("%d:%d" % (x,y), end=' ')
...
0:0 1:1 2:2 0:3 1:4 2:5
当发生广播错误时,迭代器会引发一个异常,该异常包括输入形状以帮助诊断问题.
示例
>>> import numpy as np
>>> a = np.arange(2)
>>> b = np.arange(6).reshape(2,3)
>>> for x, y in np.nditer([a,b]):
... print("%d:%d" % (x,y), end=' ')
...
Traceback (most recent call last):
...
ValueError: operands could not be broadcast together with shapes (2,) (2,3)
迭代器分配的输出数组#
在NumPy函数中一个常见的情况是根据输入的广播分配输出,并且还有一个可选参数称为’out’,当提供该参数时,结果将被放置在其中.:class:`nditer`对象提供了一种方便的习惯用法,使得支持这种机制变得非常容易.
我们将通过创建一个函数 square 来展示其工作原理,该函数将其输入平方.让我们从一个不包含 ‘out’ 参数支持的最小函数定义开始.
示例
>>> import numpy as np
>>> def square(a):
... with np.nditer([a, None]) as it:
... for x, y in it:
... y[...] = x*x
... return it.operands[1]
...
>>> square([1,2,3])
array([1, 4, 9])
默认情况下,:class:nditer 对作为 None 传递的操作数使用标志 ‘allocate’ 和 ‘writeonly’.这意味着我们只需向迭代器提供两个操作数,其余的它自行处理.
当添加 ‘out’ 参数时,我们必须显式提供这些标志,因为如果有人传入一个数组作为 ‘out’,迭代器将默认设置为 ‘readonly’,我们的内部循环将会失败.输入数组默认设置为 ‘readonly’ 的原因是为了防止无意中触发归约操作的混淆.如果默认设置为 ‘readwrite’,任何广播操作也会触发归约,这一主题在本文档后面会讨论.
在此过程中,让我们也介绍一下 ‘no_broadcast’ 标志,这将防止输出被广播.这一点很重要,因为我们每个输出只想要一个输入值.聚合多个输入值是一种需要特殊处理的缩减操作.由于缩减操作必须在迭代器标志中显式启用,因此它已经会引发错误,但禁用广播导致的错误消息对终端用户来说更容易理解.要了解如何将平方函数推广到缩减操作,请参阅关于 Cython 部分中的平方和函数.
为了完整性,我们还会添加 ‘external_loop’ 和 ‘buffered’ 标志,因为出于性能原因,这些通常是你想要的.
示例
>>> import numpy as np
>>> def square(a, out=None):
... it = np.nditer([a, out],
... flags = ['external_loop', 'buffered'],
... op_flags = [['readonly'],
... ['writeonly', 'allocate', 'no_broadcast']])
... with it:
... 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):
...
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 而不是构造另一个列表.
内层循环中的操作是一个简单的乘法.与外积有关的所有内容都由迭代器设置处理.
示例
>>> import numpy as np
>>> 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])
>>> with it:
... for x, y, z in it:
... z[...] = x*y
... result = it.operands[2] # same as z
...
>>> result
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]]])
请注意,一旦迭代器关闭,我们就无法访问 操作数 ,必须使用在上下文管理器内部创建的引用.
简化迭代#
每当一个可写操作数的元素少于完整的迭代空间时,该操作数正在进行归约.:class:nditer 对象要求任何归约操作数被标记为读写,并且仅在提供 ‘reduce_ok’ 作为迭代器标志时允许归约.
作为一个简单的例子,考虑计算数组中所有元素的总和.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> b = np.array(0)
>>> with np.nditer([a, b], flags=['reduce_ok'],
... op_flags=[['readonly'], ['readwrite']]) as it:
... for x,y in it:
... y[...] += x
...
>>> b
array(276)
>>> np.sum(a)
276
当结合reduction和allocated操作数时,事情会稍微复杂一些.在迭代开始之前,任何reduction操作数必须初始化为起始值.以下是我们如何做到这一点,沿着`a`的最后一个轴求和.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> with it:
... it.operands[1][...] = 0
... for x, y in it:
... y[...] += x
... result = it.operands[1]
...
>>> result
array([[ 6, 22, 38],
[54, 70, 86]])
>>> np.sum(a, axis=2)
array([[ 6, 22, 38],
[54, 70, 86]])
要进行缓冲减少需要在设置期间进行另一个调整.通常,迭代器构造涉及将第一个数据缓冲区从可读数组复制到缓冲区中.任何减少操作数都是可读的,因此可以读入缓冲区.不幸的是,在此缓冲操作完成后对操作数进行初始化将不会反映在迭代开始的缓冲区中,并且会产生垃圾结果.
迭代器标志 “delay_bufalloc” 用于允许迭代器分配的归约操作数与缓冲一起存在.当设置此标志时,迭代器将使其缓冲区保持未初始化状态,直到它收到重置信号,之后它将准备好进行常规迭代.以下是如果我们还启用缓冲时前一个示例的样子.
示例
>>> import numpy as np
>>> a = np.arange(24).reshape(2,3,4)
>>> it = np.nditer([a, None], flags=['reduce_ok',
... 'buffered', 'delay_bufalloc'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, [0,1,-1]])
>>> with it:
... it.operands[1][...] = 0
... it.reset()
... for x, y in it:
... y[...] += x
... result = it.operands[1]
...
>>> result
array([[ 6, 22, 38],
[54, 70, 86]])
将内层循环放入 Cython#
那些希望从他们的低级操作中获得真正良好性能的人应该强烈考虑直接使用C中提供的迭代API,但对于那些不熟悉C或C++的人来说,Cython是一个具有合理性能折衷的好中间地带.对于 nditer 对象,这意味着让迭代器处理广播、dtype转换和缓冲,同时将内部循环交给Cython.
在我们的例子中,我们将创建一个平方和函数.首先,让我们在简单的Python中实现这个函数.我们希望支持类似于numpy :func:`sum`函数的’axis’参数,所以我们需要为`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',
... 'buffered', 'delay_bufalloc'],
... op_flags=[['readonly'], ['readwrite', 'allocate']],
... op_axes=[None, axeslist],
... op_dtypes=['float64', 'float64'])
... with it:
... 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.)
>>> sum_squares_py(a, axis=-1)
array([ 5., 50.])
要使用 Cython 优化这个函数,我们将内部循环（y[…] += x*x）替换为专门针对 float64 dtype 的 Cython 代码.启用 ‘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'])
with it:
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代码和使用NumPy内置sum函数的表达式.:
>>> 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