MXNet numpy 模块的索引机制基础切片、高级索引与布尔掩码的实现解析【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet本文以 MXNet 官方文档中的 NumPy 兼容索引参考 arrays.indexing.rst 为主体系统讲解mxnet.np.ndarray的三类索引方式——基础切片basic slicing、高级索引advanced indexing和布尔掩码索引boolean mask indexing——并深入 mxnet.numpy.multiarray 的源码说明每种索引何时返回共享内存的视图、何时返回副本以及布尔索引为何只支持单个布尔数组等限制背后的实现原因帮助你在 MXNet 数值计算mx.np环境中正确、高效地编写数组选取与赋值代码。索引的三种类型与总体分发机制:class:ndarray 可以使用标准的 Pythonx[obj]语法进行索引其中x是数组obj是选取对象。MXNet 的 numpy 兼容模块遵循 NumPy 的语义提供三种索引类型基础切片basic slicingobj是slice对象start:stop:step语法、整数或由它们组成的元组高级索引advanced indexingobj是非元组的序列对象、整型或布尔型ndarray或至少包含一个序列/整型/布尔型ndarray的元组布尔掩码索引boolean mask indexingobj是布尔类型的数组通常来自比较运算符的结果。具体发生哪一种取决于obj的内容。在 Python 中x[(exp1, exp2, ..., expN)]与x[exp1, exp2, ..., expN]完全等价后者只是前者的语法糖。从源码结构看__getitem__在 multiarray.py#L644 中实现它先处理若干快速路径单个整数、单个切片、纯整型 ndarray 元组等然后调用 indexing_key_expand_implicit_axes 把索引中的None、Ellipsis展开为显式轴向再由 get_indexing_dispatch_code 判定分派码最终走入三个分支之一_NDARRAY_BASIC_INDEXING→_get_np_basic_indexing_NDARRAY_ADVANCED_INDEXING→_get_np_advanced_indexing单个布尔索引 →_get_np_boolean_indexing在进入分派之前被优先拦截以获得更高速度这与文档三种索引类型的划分一一对应。基础切片与索引基本规则基础切片把 Python 的一维切片概念扩展到 N 维。当obj是slice对象、整数或二者的元组时发生基础切片Ellipsis和newaxis对象也可以穿插其中。用N个整数索引的最简单情况返回一个数组标量。与 Python 一样所有索引从 0 开始对于第i个索引n_i有效范围是0 n_i d_i其中d_i是数组 shape 的第i个元素。负索引从数组末尾计数n_i 0表示n_i d_i。所有由基础切片产生的数组只要取出的元素在内存中连续就始终是对原数组的视图view。这一点在源码中有直接对应_get_np_basic_indexing先调用_basic_indexing_slice_is_contiguous判断连续性若连续则通过 C APIMXNDArraySlice64/MXNDArraySlice在展平后的数组上做低层切片并构造共享 handle 的视图multiarray.py#L513-L543若不连续则退化为_npi.slice产生一个新数组。因此切片后改视图会改原数组改非连续切片则不会这一行为可以直接从该代码路径得到解释。i:j:k语义与默认值标准序列切片规则逐维适用包括 step。需要记住的要点切片语法为i:j:ki是起始索引j是停止索引k是步长k ! 0。它选出m个元素索引值为i、i k、...、i (m - 1)k其中m由j - i除以k的商q和余数r决定j - i q k r且i (m - 1) k j。 x np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) x[1:7:2] array([1, 3, 5])负的i和j解释为n i和n jn为该维元素数负的k让索引朝更小的方向步进。 x[-2:10] array([8, 9]) x[-3:3:-1] array([7, 6, 5, 4])设n是被切片维的元素数i缺省时k 0默认为 0k 0默认为n - 1j缺省时k 0默认为nk 0默认为-n - 1k缺省为 1。::与:相同表示沿该轴选取全部索引。 x[5:] array([5, 6, 7, 8, 9])若选取元组中的对象数少于N维则后续各维自动补:。 x np.array([[[1],[2],[3]], [[4],[5],[6]]]) x.shape (2, 3, 1) x[1:2] array([[[4], [5], [6]]])Ellipsis会扩展为补足选取元组覆盖所有维度所需的:个数多数情况下扩展后元组长度为x.ndim。选取中最多只能出现一个Ellipsis。 x[...,0] array([[1, 2, 3], [4, 5, 6]])选取元组中的每个newaxis对象会在结果中增加一个长度为 1 的新维度位置即newaxis在元组中的位置。newaxis是None的别名在 MXNet 中同样直接定义为None见 utils.py#L58可直接用None替换。 x[:,np.newaxis,:,:].shape (2, 1, 3, 1)整数i返回与i:i1相同的值但结果对象的维数降低 1。具体地选取元组中第p个元素为整数其余均为:时返回 N - 1 维的子数组当N 1时返回ndim0的标量ndarray。若选取元组除第p项为切片i:j:k外全为:则返回 N 维数组由对元素i、ik、... 的整数索引子数组拼接而成。多个非:项的基础切片等价于逐个应用单个非:项的切片其余非:项临时替换为:。因此x[ind1,...,ind2,:]在基础切片下等价于x[ind1][...,ind2,:]。注意上述结论不适用于高级索引。可以用切片给数组赋值但与 list 不同永远不能借此扩大数组。x[obj] value中 value 的形状必须能广播到x[obj]的形状。在实现上整数轴向会被统一转换为切片形式后再进入_basic_indexing_key_int_to_slice并在 _get_np_basic_indexing#L506-L511 中做越界检查越界时抛出IndexError。用 slice 对象构造通用索引一个切片元组总可以显式构造为obj并用于x[obj]记法。构造中可用slice(...)对象代替[start:stop:step]语法。例如x[1:10:5,::-1]也可以写成obj (slice(1,10,5), slice(None,None,-1)) x[obj]这对编写可作用于任意维数数组的通用代码很有用。高级索引当obj是非元组序列对象、整型或布尔型ndarray或元组中至少含一个序列/整型/布尔型ndarray时触发高级索引。高级索引总是返回数据的副本对照基础切片在部分情况下返回视图。警告x[(1,2,3),]与x[(1,2,3)]有本质区别。后者等价于x[1,2,3]触发基础索引前者是单个长度为 3 的序列触发高级索引。必须理解产生这种差异的原因Python 元组解析。另外x[[1,2,3]]触发高级索引而x[[1,2,slice(None)]]在官方 NumPy 中出于已废弃的 Numeric 兼容考虑触发基础切片这种组合在 MXNet 的 numpy 模块中当前不支持。纯整型数组索引当索引由与数组维数相同的整型数组组成时行为直观但不同于切片。高级索引总是先广播、再作为整体迭代result[i_1, ..., i_M] x[ind_1[i_1, ..., i_M], ind_2[i_1, ..., i_M], ..., ind_N[i_1, ..., i_M]]结果形状与广播后的索引数组ind_1, ..., ind_N的形状完全相同。典型例子——从每行选取一个特定元素行索引为[0, 1, 2]列索引[0, 1, 0]指定每行取哪一列 x np.array([[1, 2], [3, 4], [5, 6]]) x[[0, 1, 2], [0, 1, 0]] array([1, 4, 5])在源码中全元组都是整型 ndarray的情况会走快速路径_npi.advanced_indexing_multiplemultiarray.py#L840-L844一般情况则归一化为索引张量后统一调用_npi.gather_nd完成这解释了结果形状与广播后的索引形状一致以及总是副本两个性质。混合高级索引与基础索引当索引中至少含有一个切片:、Ellipsis或newaxis或数组维数多于高级索引个数时行为更复杂可以理解为为每个高级索引元素拼接其索引结果。最简单的情况是只有一个高级索引单个高级索引可以替换一个切片结果数组相同但它是副本且内存布局可能不同。能用切片时优先用切片。 x[1:2, 1:3] array([[4, 5]]) x[1:2, [1, 2]] array([[4, 5]])理解混合索引最方便的方式是从结果形状入手。索引操作分两部分基础索引不含整数定义的子空间以及高级索引部分的子空间。需要区分两种情形高级索引之间被切片、Ellipsis或newaxis分隔例如x[arr1, :, arr2]此时高级索引产生的维度排在结果数组的前面基础子空间维度排在其后。所有高级索引彼此相邻例如x[..., arr1, arr2, :]但x[arr1, :, 1]不属于此情形因为整数1在此处被视为高级索引。此时高级索引产生的维度会插回结果数组中它原本出现的位置正是这一规则使简单高级索引看起来与切片一致。形状演算示例设x.shape为 (10,20,30)ind是形状 (2,3,4) 的整型索引数组则result x[...,ind,:]的形状为(10,2,3,4,30)(20,) 的基础子空间被替换成 (2,3,4) 的广播索引子空间。令i, j, k遍历 (2,3,4) 子空间则result[...,i,j,k,:] x[...,ind[i,j,k],:]。该例结果与x.take(ind, axis-2)相同。设x.shape为 (10,20,30,40,50)ind_1与ind_2可广播到 (2,3,4)x[:,ind_1,ind_2]形状为(10,2,3,4,40,50)高级索引相邻原位替换 (20,30) 子空间而x[:,ind_1,:,ind_2]形状为(2,3,4,10,30,50)高级索引被基础项分隔索引子空间被前置。此时没有明确的插入点因此统一放在开头需要时可随时用transpose把子空间移到任意位置。注意这种分隔情形无法用take复现。布尔掩码索引当obj是布尔类型的数组对象例如比较运算符的返回值时发生布尔索引。单个布尔索引数组实际上等价于x[obj.nonzero()]——obj.nonzero()返回一个长度为obj.ndim的整型索引数组元组标出obj中所有True元素的位置但当obj.shape x.shape时直接走布尔路径更快。规则若obj.ndim x.ndimx[obj]返回一维数组按行主序C 顺序填入x中对应True位置的元素。若obj的True位置超出x的边界抛出索引错误若obj比x小则等价于用False填充。当前限制布尔索引只支持单个布尔 ndarray作为索引暂不支持布尔数组与其他索引项组合的复合索引。# 选取所有行和 2 的行 x np.array([[0, 1], [1, 1], [2, 2]], dtypenp.int32) rowsum x.sum(-1) x[rowsum 2] array([[0, 1], [1, 1]], dtypeint32) # 若 rowsum 也带两维则报错 rowsum x.sum(-1, keepdimsTrue) rowsum.shape (3, 1) x[rowsum 2] # fail IndexError: boolean index did not match indexed array along dimension 1源码印证了这条限制__getitem__中只要 key 是单个布尔ndarray就在进入通用分派之前直接调用_get_np_boolean_indexingmultiarray.py#L803-L804。该方法先校验key_ndim ndim且前key_ndim个维度形状严格相等不匹配时抛出上面示例中的IndexErrormultiarray.py#L610-L619校验通过后将两端的维度展平为一维调用底层_npi.boolean_mask一次完成筛选再 reshape 回(0,) shape[key_ndim:]的形态——这比先nonzero再高级索引的两步法少了一次索引数组的构造与gather_nd开销即文档所说更快的实现来源。赋值方向同样存在快速路径单个布尔掩码赋值走_set_np_boolean_indexingmultiarray.py#L627-L641根据 value 是数值还是ndarray分别调用_npi.boolean_mask_assign_scalar或_npi.boolean_mask_assign_tensor原地写入。而混合了布尔数组与其他索引项的复合赋值则不在支持范围内与文档的限制说明一致。详细注意事项Detailed notes以下细节对日常索引不重要但边界场景下必须知道文档明确列出的注意点高级赋值的迭代顺序无保证。若同一次高级赋值中某个元素被写入多次最终结果不可预测。空元组索引x[()]是对零维数组的完整标量索引x为零维时x[()]返回一个scalarndarray而x[...]总是返回视图。源码中对应快速路径isinstance(key, tuple) and len(key) 0时直接返回自身multiarray.py#L810-L811零维数组则只接受()作为索引否则抛出IndexError(scalar tensor can only accept()as index)multiarray.py#L807-L808。索引中出现零维数组时MXNet 与 NumPy 不同它不被视为完整的整型索引因此不会触发高级索引。零维布尔数组不满足nonzero等价关系。当高级索引操作的结果不含任何元素、但某个个体索引越界时MXNet当前不会像 NumPy 那样抛出IndexError。使用限制与适用前提结合源码可以看到几个 MXNet numpy 索引相对官方 NumPy 的差异点使用时需注意不支持x[[1,2,slice(None)]]这类列表含 slice的基础切片兼容写法文档明确说明官方 NumPy 因已废弃的 Numeric 兼容才支持它MXNet 未实现。布尔索引只接受单个布尔 ndarray布尔数组与整数/切片混合的复合索引不受支持。HybridBlock符号/延迟计算模式下限制更严格用 numpy 数组作索引会直接抛出TypeError(Indexing with a numpy array is not supported in HybridBlock.)multiarray.py#L793-L795元组索引在该模式下仅支持 slice 与整数项且为保持与命令式模式的行为一致会主动放弃共享内存路径出现不支持的索引类型会抛出TypeError(The type of indexing used is not supported in HybridBlock.)multiarray.py#L883-L884。因此在mx.gluon.HybridBlock中编写动态形状代码时应优先使用整数与切片的基础索引。__setitem__不支持把mx.nd.NDArray直接赋给mx.np.ndarray会抛TypeError赋值的 value 必须是数值、可广播的数组或mxnet.np.ndarray本身multiarray.py#L916-L967。掌握视图 vs 副本基础切片连续时共享内存、高级索引恒为副本、Ellipsis/newaxis的形状规则、混合索引的两种维度排布情形以及布尔索引的单掩码限制就足以在 MXNet 的mx.np数值计算环境中正确编写绝大多数索引与赋值逻辑。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mx/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考