迭代数组#
在 NumPy 1.6 中引入的迭代器对象nditer，提供了许多灵活的方式来系统地访问一个或多个数组的所有元素。本文档首先介绍了一些使用该对象进行 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’) 的元素由于被放入了不同的内存布局，因此以不同的顺序被访问。
控制迭代顺序#
有时，需要按照特定的顺序访问数组的元素，而不管元素在内存中的布局。 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
修改数组值#
默认情况下，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 版本之前，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]])
跟踪索引或多维索引与使用外部循环不兼容，因为它需要每个元素都有不同的索引值。如果您尝试组合这些标志，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
替代循环和元素访问#
为了在迭代过程中更方便地访问其属性，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]
以特定数据类型进行迭代#
有时需要将数组视为与存储类型不同的数据类型。例如，可能希望所有计算都使用 64 位浮点数，即使正在操作的数组是 32 位浮点数。除了编写低级 C 代码外，通常最好让迭代器处理复制或缓冲，而不是在内层循环中自己转换数据类型。
有两种机制可以实现这一点：临时副本和缓冲模式。使用临时副本，会创建一个具有新数据类型的新数组副本，然后对副本进行迭代。通过一种模式可以进行写访问，该模式在所有迭代完成后更新原始数组。临时副本的主要缺点是，临时副本可能会占用大量内存，特别是如果迭代数据类型具有比原始数据类型更大的项大小。
缓冲模式可以缓解内存使用问题，并且比创建临时副本更具缓存友好性。除了需要立即在迭代器外部处理整个数组的特殊情况外，建议使用缓冲而不是临时复制。在 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 有一套处理形状不同的数组的规则，这些规则在函数接受多个组合元素的数时会应用。这被称为广播。当您需要编写这样的函数时，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”参数，用于在提供结果时将其放置。 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])
默认情况下，nditer 对传递为 None 的操作数使用 'allocate' 和 'writeonly' 标志。这意味着我们可以只向迭代器提供两个操作数，它会处理其余的。
添加“out”参数时，我们必须显式地提供这些标志，因为如果有人将数组作为“out”传递，迭代器将默认为“readonly”，我们的内层循环将失败。之所以将输入数组的默认值设为“readonly”，是为了避免因意外触发归约操作而产生的混淆。如果默认值为“readwrite”，则任何广播操作也会触发归约，这是一个将在本文档后面讨论的主题。
同时，我们还可以引入 'no_broadcast' 标志，这将阻止输出被广播。这一点很重要，因为我们只希望每个输出有一个输入值。聚合多个输入值是归约操作，需要特殊处理。它已经会引发错误，因为必须在迭代器标志中显式启用归约，但禁用广播导致的错误消息对最终用户来说更易于理解。要了解如何将 square 函数泛化为归约，请参阅关于 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]]])
请注意，一旦迭代器关闭，我们就无法访问 operands，并且必须使用在上下文管理器内创建的引用。
归约迭代#
每当一个可写操作数的元素少于完整的迭代空间时，该操作数就正在进行归约。 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
当组合归约和分配的操作数时，情况会稍微复杂一些。在开始迭代之前，任何归约操作数都必须初始化为其起始值。以下是我们可以如何做到这一点，计算 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 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 化，我们用专门针对 float64 dtype 的 Cython 代码替换内层循环（y[…] += x*x）。当启用 '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