Tensor索引支持

下载Notebook 下载样例代码 查看源文件

Tensor 支持单层与多层索引取值，赋值以及增强赋值，支持动态图(PyNative)以及静态图(Graph)模式。

索引取值

索引值支持int、bool、None、ellipsis、slice、Tensor、List、Tuple。

  • int索引取值

    支持单层和多层int索引取值，单层int索引取值：tensor_x[int_index]，多层int索引取值：tensor_x[int_index0][int_index1]...。

    int索引取值操作的是第零维，索引值小于第零维长度，在取出第零维对应位置数据后，会消除第零维。

    例如，如果对一个shape为(3, 4, 5)的tensor进行单层int索引取值，取得结果的shape是(4, 5)。

    多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行int索引取值。

    示例如下：

[1]:

import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(2 * 3 * 2).reshape((2, 3, 2)))
data_single = tensor_x[0]
data_multi = tensor_x[0][1]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[0 1]
 [2 3]
 [4 5]]
data_multi:
[2 3]

  • bool索引取值

    支持单层和多层bool索引取值，单层bool索引取值：tensor_x[True]，多层bool索引取值：tensor_x[True][True]...。

    bool索引取值操作的是第零维，在取出所有数据后，会在axis=0轴上扩展一维，对应True/False，该维长度分别为1/0。False将会在shape中引入0，因此暂只支持True。

    例如，对一个shape为(3, 4, 5)的tensor进行单层True索引取值，取得结果的shape是(1, 3, 4, 5)。

    多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行bool索引取值。

    示例如下：

[2]:

import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(2 * 3).reshape((2, 3)))
data_single = tensor_x[True]
data_multi = tensor_x[True][True]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[[0 1 2]
  [3 4 5]]]
data_multi:
[[[[0 1 2]
   [3 4 5]]]]

  • None索引取值

    None索引取值和True索引取值一致，可参考True索引取值，这里不再赘述。

  • ellipsis索引取值

    支持单层和多层ellipsis索引取值，单层ellipsis索引取值：tensor_x[...]，多层ellipsis索引取值：tensor_x[...][...]...。

    ellipsis索引取值操作在所有维度上取出所有数据。一般多作为Tuple索引的组成元素，Tuple索引将于下面介绍。

    例如，对一个shape为(3, 4, 5)的tensor进行ellipsis索引取值，取得结果的shape依然是(3, 4, 5)。

    示例如下：

[3]:

import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(2 * 3).reshape((2, 3)))
data_single = tensor_x[...]
data_multi = tensor_x[...][...]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[0 1 2]
 [3 4 5]]
data_multi:
[[0 1 2]
 [3 4 5]]

  • slice索引取值

    支持单层和多层slice索引取值，单层slice索引取值：tensor_x[slice_index]，多层slice索引取值：tensor_x[slice_index0][slice_index1]...。

    slice索引取值操作的是第零维，取出第零维所切到位置的元素，slice不会降维，即使切到长度为1，区别于int索引取值。

    例如，tensor_x[0:1:1] != tensor_x[0]，因为shape_former = (1,) + shape_latter。

    多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行slice索引取值。

    slice有start、stop和step组成。start默认值为0，stop默认值为该维长度，step默认值为1。

    例如，tensor_x[:] == tensor_x[0:length:1]。

    示例如下：

[4]:

import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(4 * 2 * 2).reshape((4, 2, 2)))
data_single = tensor_x[1:4:2]
data_multi = tensor_x[1:4:2][1:]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[[ 4  5]
  [ 6  7]]

 [[12 13]
  [14 15]]]
data_multi:
[[[12 13]
  [14 15]]]

  • Tensor索引取值

    支持单层和多层Tensor索引取值，单层Tensor索引取值：tensor_x[tensor_index]，多层Tensor索引取值：tensor_x[tensor_index0][tensor_index1]...。

    Tensor索引取值操作的是第零维，取出第零维对应位置的元素。

    索引Tensor数据类型支持int型和bool型。

    当数据类型是int型时，可以是(int8, int16, int32, int64)，值必须为非负数，且小于第零维长度。

    Tensor索引取值得到结果的data_shape = tensor_index.shape + tensor_x.shape[1:]。

    例如，对一个shape为(6, 4, 5)的tensor通过shape为(2, 3)的tensor进行索引取值，取得结果的shape为(2, 3, 4, 5)。

    当数据类型是bool型时，Tensor索引取值得到结果的维度是 tensor_x.ndim - tensor_index.ndim + 1。

    设 tensor_index 中True的数量是 num_true ，tensor_x 的shape是 (N0, N1, ... Ni-1, Ni, Ni+1, ..., Nk), tensor_index 的shape是 (N0, N1, ... Ni-1), 则返回值的shape是 (num_true, Ni+1, Ni+2, ... , Nk) 。

    示例如下：

[5]:

from mindspore import dtype as mstype
import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor([1, 2, 3])
tensor_index = ms.Tensor([True, False, True], dtype=mstype.bool_)
output = tensor_x[tensor_index]
print(output)

[1 3]

多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行Tensor索引取值。

示例如下：

[6]:

from mindspore import dtype as mstype
import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(4 * 2 * 3).reshape((4, 2, 3)))
tensor_index0 = ms.Tensor(np.array([[1, 2], [0, 3]]), mstype.int32)
tensor_index1 = ms.Tensor(np.array([[0, 0]]), mstype.int32)
data_single = tensor_x[tensor_index0]
data_multi = tensor_x[tensor_index0][tensor_index1]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[[[ 6  7  8]
   [ 9 10 11]]

  [[12 13 14]
   [15 16 17]]]


 [[[ 0  1  2]
   [ 3  4  5]]

  [[18 19 20]
   [21 22 23]]]]
data_multi:
[[[[[ 6  7  8]
    [ 9 10 11]]

   [[12 13 14]
    [15 16 17]]]


  [[[ 6  7  8]
    [ 9 10 11]]

   [[12 13 14]
    [15 16 17]]]]]

  • List索引取值

    支持单层和多层List索引取值，单层List索引取值：tensor_x[list_index]，多层List索引取值：tensor_x[list_index0][list_index1]...。

    List索引取值操作的是第零维，取出第零维对应位置的元素。

    索引List数据类型必须是int、bool或两者混合。若数据类型为int，则取值在[-dimension_shape, dimension_shape-1]之间；若数据类型为bool， 则限制bool个数为对应维度长度，筛选对应维度上值为True的元素；若值为前两者混合，则bool类型的True/False将转为int类型的1/0。

    List索引取值得到结果的data_shape = list_index.shape + tensor_x.shape[1:]。

    例如，对一个shape为(6, 4, 5)的tensor通过shape为(3,)的tensor进行索引取值，取得结果的shape为(3, 4, 5)。

    多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行List索引取值。

    示例如下：

[7]:

import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(4 * 2 * 3).reshape((4, 2, 3)))
list_index0 = [1, 2, 0]
list_index1 = [True, False, True]
data_single = tensor_x[list_index0]
data_multi = tensor_x[list_index0][list_index1]
print('data_single:')
print(data_single)
print('data_multi:')
print(data_multi)

data_single:
[[[ 6  7  8]
  [ 9 10 11]]

 [[12 13 14]
  [15 16 17]]

 [[ 0  1  2]
  [ 3  4  5]]]
data_multi:
[[[ 6  7  8]
  [ 9 10 11]]

 [[ 0  1  2]
  [ 3  4  5]]]

  • Tuple索引取值

    索引Tuple的数据类型可以为int、bool、None、slice、ellipsis、Tensor、List、Tuple。支持单层和多层Tuple索引取值，单层Tuple索引取值：tensor_x[tuple_index]，多层Tuple索引取值：tensor_x[tuple_index0][tuple_index1]...。Tuple中包含的List与Tuple包含元素规则与单独的List规则相同，其他元素规则与单独元素也相同。

    索引Tuple中元素按照最终索引Broadcast规则，分为Basic Index、Advanced Index两类。Basic Index包含slice、ellipsis、int与None四种类型，Advanced Index包含bool、Tensor、List、Tuple等类型。索引过程中，所有的Advanced Index将会做Broadcast，若Advaned Index连续，最终broadcast shape将插入在第一个Advanced Index位置；若不连续，则broadcast shape插入在0位置。

    索引里除None扩展对应维度，bool扩展对应维度后与Advanced Index做Broadcast。除ellipsis、bool、None外每个元素操作对应位置维度，即Tuple中第0个元素操作第零维，第1个元素操作第一维，以此类推。每个元素的索引规则与该元素类型索引取值规则一致。

    Tuple索引里最多只有一个ellipsis，ellipsis前半部分索引元素从前往后对应Tensor第零维往后，后半部分索引元素从后往前对应Tensor最后一维往前，其他未指定的维度，代表全取。

    元素里包含的Tensor数据类型必须是int型或bool型，int型可以是(int8, int16, int32, int64)，值必须为非负数，且小于第零维长度。

    例如，tensor_x[0:3, 1, tensor_index] == tensor_x[(0:3, 1, tensor_index)]，因为0:3, 1, tensor_index就是一个Tuple。

    多层索引取值可以理解为，后一层索引取值在前一层索引取值结果上再进行Tuple索引取值。

    示例如下：

[8]:

from mindspore import dtype as mstype
import mindspore as ms
import mindspore.numpy as np
tensor_x = ms.Tensor(np.arange(2 * 3 * 4).reshape((2, 3, 4)))
tensor_index = ms.Tensor(np.array([[1, 2, 1], [0, 3, 2]]), mstype.int32)
data = tensor_x[1, 0:1, tensor_index]
print('data:')
print(data)

data:
[[[13]
  [14]
  [13]]

 [[12]
  [15]
  [14]]]

索引赋值

对于形如: tensor_x[index] = value， index的类型支持int、bool、ellipsis、slice、None、Tensor、List、Tuple。

value的类型支持Number、Tuple、List和Tensor。被赋的值会首先被转换为张量，数据类型与原张量(tensor_x)相符。

当value为Number时，可以理解为将tensor_x[index]索引对应元素都更新为Number。

当value为数组，即只包含Number的Tuple、List或Tensor时，value.shape需要可以与tensor_x[index].shape做广播，将value广播到tensor_x[index].shape后，更新tensor_x[index]对应的值。

当value为Tuple或List时，若value中元素包含Number，Tuple，List 和 Tensor等多种类型，该Tuple 和 List 目前只支持一维。

当value为Tuple或List，且存在Tensor时，非Tensor的元素会首先被转换为Tensor，然后这些Tensor在axis=0轴上打包之后成为一个新的赋值Tensor，这时按照value为Tensor的规则进行赋值。所有Tensor的数据类型必须保持一致。

索引赋值可以理解为对索引到的位置元素按照一定规则进行赋值，所有索引赋值都不会改变原Tensor的shape。

当索引中有多个元素指向原张量的同一个位置时，该值的更新受底层算子限制，可能出现随机的情况。因此暂不支持索引中重复对张量中一个位置的值反复更新。详情请见:TensorScatterUpdate 算子介绍

当前只支持单层索引(tensor_x[index] = value)，多层索引(tensor_x[index1][index2]... = value)暂不支持。

  • int索引赋值

    支持单层int索引赋值：tensor_x[int_index] = u。

    示例如下：

[9]:

import mindspore.numpy as np
tensor_x = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_y = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_x[1] = 88.0
tensor_y[1] = np.array([66, 88, 99]).astype(np.float32)
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)

tensor_x:
[[ 0.  1.  2.]
 [88. 88. 88.]]
tensor_y:
[[ 0.  1.  2.]
 [66. 88. 99.]]

  • bool索引赋值

    支持单层bool索引赋值：tensor_x[bool_index] = u。

    示例如下：

[10]:

import mindspore.numpy as np
tensor_x = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_y = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_z = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_x[True] = 88.0
tensor_y[True] = np.array([66, 88, 99]).astype(np.float32)
tensor_z[True] = (66, 88, 99)
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)

tensor_x:
[[88. 88. 88.]
 [88. 88. 88.]]
tensor_y:
[[66. 88. 99.]
 [66. 88. 99.]]
tensor_z:
[[66. 88. 99.]
 [66. 88. 99.]]

  • ellipsis索引赋值

    支持单层ellipsis索引赋值，单层ellipsis索引赋值：tensor_x[...] = u。

    示例如下：

[11]:

import mindspore.numpy as np
tensor_x = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_y = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_z = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_x[...] = 88.0
tensor_y[...] = np.array([[22, 44, 55], [22, 44, 55]])
tensor_z[...] = ([11, 22, 33], [44, 55, 66])
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)

tensor_x:
[[88. 88. 88.]
 [88. 88. 88.]]
tensor_y:
[[22. 44. 55.]
 [22. 44. 55.]]
tensor_z:
[[11. 22. 33.]
 [44. 55. 66.]]

  • slice索引赋值

    支持单层slice索引赋值：tensor_x[slice_index] = u。

    示例如下：

[12]:

import mindspore.numpy as np
tensor_x = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_y = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_z = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_k = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_x[0:1] = 88.0
tensor_y[0:2] = 88.0
tensor_z[0:2] = np.array([[11, 12, 13], [11, 12, 13]]).astype(np.float32)
tensor_k[0:2] = ([11, 12, 13], (14, 15, 16))
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)
print('tensor_k:')
print(tensor_k)

tensor_x:
[[88. 88. 88.]
 [ 3.  4.  5.]
 [ 6.  7.  8.]]
tensor_y:
[[88. 88. 88.]
 [88. 88. 88.]
 [ 6.  7.  8.]]
tensor_z:
[[11. 12. 13.]
 [11. 12. 13.]
 [ 6.  7.  8.]]
tensor_k:
[[11. 12. 13.]
 [14. 15. 16.]
 [ 6.  7.  8.]]

  • None索引赋值

    支持单层None索引赋值：tensor_x[none_index] = u。

    示例如下：

[13]:

import mindspore.numpy as np
tensor_x = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_y = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_z = np.arange(2 * 3).reshape((2, 3)).astype(np.float32)
tensor_x[None] = 88.0
tensor_y[None] = np.array([66, 88, 99]).astype(np.float32)
tensor_z[None] = (66, 88, 99)
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)

tensor_x:
[[88. 88. 88.]
 [88. 88. 88.]]
tensor_y:
[[66. 88. 99.]
 [66. 88. 99.]]
tensor_z:
[[66. 88. 99.]
 [66. 88. 99.]]

  • Tensor索引赋值

    支持单层Tensor索引赋值，即tensor_x[tensor_index] = u。

    当前支持索引Tensor为 int 型和 bool 型。

    int 型示例如下：

[14]:

import mindspore.numpy as np
tensor_x = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_y = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_z = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_index = np.array([[2, 0, 2], [0, 2, 0], [0, 2, 0]], np.int32)
tensor_x[tensor_index] = 88.0
tensor_y[tensor_index] = np.array([11.0, 12.0, 13.0]).astype(np.float32)
tensor_z[tensor_index] = [11, 12, 13]
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)

tensor_x:
[[88. 88. 88.]
 [ 3.  4.  5.]
 [88. 88. 88.]]
tensor_y:
[[11. 12. 13.]
 [ 3.  4.  5.]
 [11. 12. 13.]]
tensor_z:
[[11. 12. 13.]
 [ 3.  4.  5.]
 [11. 12. 13.]]

bool 型示例如下：

[15]:

from mindspore import dtype as mstype
import mindspore as ms
tensor_x = ms.Tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8]], mstype.float32)
tensor_index = ms.Tensor([True, False, True], mstype.bool_)
tensor_x[tensor_index] = -1
print(tensor_x)

[[-1. -1. -1.]
 [ 3.  4.  5.]
 [-1. -1. -1.]]

  • List索引赋值

    支持单层List索引赋值：tensor_x[list_index] = u。

    List索引赋值和List索引取值对索引的支持一致。

    示例如下：

[16]:

import mindspore.numpy as np
tensor_x = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_y = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_index = np.array([[0, 1], [1, 0]]).astype(np.int32)
tensor_x[[0, 1]] = 88.0
tensor_y[[True, False, False]] = np.array([11, 12, 13]).astype(np.float32)
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)

tensor_x:
[[88. 88. 88.]
 [88. 88. 88.]
 [ 6.  7.  8.]]
tensor_y:
[[11. 12. 13.]
 [ 3.  4.  5.]
 [ 6.  7.  8.]]

  • Tuple索引赋值

    支持单层Tuple索引赋值：tensor_x[tuple_index] = u。

    Tuple索引赋值和Tuple索引取值对索引的支持一致，但不支持Tuple中包含None。

    示例如下：

[17]:

import mindspore.numpy as np
tensor_x = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_y = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_z = np.arange(3 * 3).reshape((3, 3)).astype(np.float32)
tensor_index = np.array([0, 1]).astype(np.int32)
tensor_x[1, 1:3] = 88.0
tensor_y[1:3, tensor_index] = 88.0
tensor_z[1:3, tensor_index] = np.array([11, 12]).astype(np.float32)
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)
print('tensor_z:')
print(tensor_z)

tensor_x:
[[ 0.  1.  2.]
 [ 3. 88. 88.]
 [ 6.  7.  8.]]
tensor_y:
[[ 0.  1.  2.]
 [88. 88.  5.]
 [88. 88.  8.]]
tensor_z:
[[ 0.  1.  2.]
 [11. 12.  5.]
 [11. 12.  8.]]

索引增强赋值

增强索引赋值，支持+=、-=、*=、/=、%=、**=、//=七种类型，index与value的规则约束与索引赋值相同。索引值支持int、bool、ellipsis、slice、None、Tensor、List、Tuple八种类型，赋值支持Number、Tensor、Tuple、List四种类型。

索引增强赋值可以理解为对索引到的位置元素按照一定规则进行取值，取值所得再与value进行操作符运算，最终将运算结果进行赋值，所有索引增强赋值都不会改变原Tensor的shape。

当索引中有多个元素指向原张量的同一个位置时，该值的更新受底层算子限制，可能出现随机的情况。因此暂不支持索引中重复对张量中一个位置的值反复更新。详情请见:TensorScatterUpdate 算子介绍。

目前索引中包含 True、False 和 None的情况暂不支持。

  • 规则与约束

    与索引赋值相比，增加了取值与运算的过程。取值过程中index的约束规则与索引取值中index相同，支持int、bool、Tensor、Slice、Ellipsis、None、List与Tuple。上述几种类型的数据中所包含int值，需在[-dim_size, dim_size-1]闭合区间内。 运算过程中value的约束规则与索引赋值中value的约束规则相同，value类型需为(Number、Tensor、List、Tuple)之一，且value类型不是Number时， value的形状需要可以广播到tensor_x[index]的形状。

    示例如下：

[18]:

import mindspore as ms
tensor_x = ms.Tensor(np.arange(3 * 4).reshape(3, 4).astype(np.float32))
tensor_y = ms.Tensor(np.arange(3 * 4).reshape(3, 4).astype(np.float32))
tensor_x[[0, 1], 1:3] += 2
tensor_y[[1], ...] -= [4, 3, 2, 1]
print('tensor_x:')
print(tensor_x)
print('tensor_y:')
print(tensor_y)

tensor_x:
[[ 0.  3.  4.  3.]
 [ 4.  7.  8.  7.]
 [ 8.  9. 10. 11.]]
tensor_y:
[[ 0.  1.  2.  3.]
 [ 0.  2.  4.  6.]
 [ 8.  9. 10. 11.]]