maxframe.tensor.take#

maxframe.tensor.take(a, indices, axis=None, out=None)[源代码]#

沿轴从张量中提取元素。

当 axis 不为 None 时，此函数的作用与“花式”索引（使用张量索引数组）相同；但如果你需要沿着某个轴提取元素，它可能更容易使用。例如调用 mt.take(arr, indices, axis=3) 等价于 arr[:,:,:,indices,...]。

如果不使用花式索引来解释，它等价于以下使用 ndindex 的方式，其中将 ii、jj 和 kk 设置为索引元组:

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
Nj = indices.shape
for ii in ndindex(Ni):
    for jj in ndindex(Nj):
        for kk in ndindex(Nk):
            out[ii + jj + kk] = a[ii + (indices[jj],) + kk]
参数:

  • a (array_like (Ni..., M, Nk...)) -- 源张量。

  • indices (array_like (Nj...)) -- 要提取的值的索引。也允许标量作为索引。

  • axis (int, optional) -- 选择值的轴。默认情况下，使用展平后的输入张量。

  • out (Tensor, optional (Ni..., Nj..., Nk...)) -- 如果提供此参数，结果将被放置在该张量中。它应该具有适当的形状和数据类型。

  • mode ({'raise', 'wrap', 'clip'}, optional) -- 指定超出边界的索引的行为。 * 'raise' -- 抛出错误（默认） * 'wrap' -- 循环处理 * 'clip' -- 裁剪到范围 'clip' 模式表示所有过大的索引都将被替换为指向该轴最后一个元素的索引。请注意，这会禁用负数索引。

返回:

out -- 返回的张量与 a 的类型相同。

返回类型:

Tensor (Ni..., Nj..., Nk...)

参见

compress

使用布尔掩码提取元素

Tensor.take

等效方法

备注

通过消除上述描述中的内层循环，并使用 s_ 构造简单切片对象，可以用对每个一维切片应用花式索引的方式来表达 take:

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
    for kk in ndindex(Nj):
        out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices]

因此，它等价于（但比）以下使用 apply_along_axis 的方式更快:

out = mt.apply_along_axis(lambda a_1d: a_1d[indices], axis, a)

示例

>>> import maxframe.tensor as mt
>>> a = [4, 3, 5, 7, 6, 8]
>>> indices = [0, 1, 4]
>>> mt.take(a, indices).execute()
array([4, 3, 6])

在此示例中，如果 a 是一个张量，则可以使用“花式”索引。

>>> a = mt.array(a)
>>> a[indices].execute()
array([4, 3, 6])

如果 indices 不是一维的，输出也具有这些维度。

>>> mt.take(a, [[0, 1], [2, 3]]).execute()
array([[4, 3],
       [5, 7]])