tensorflow scatter_nd函数

函数原型

tf.scatter_nd(
    indices, updates, shape, name=None
)

函数说明

将张量updates的形状扩展到新的形状shape，张量在新形状shape中的坐标为indices，其他的元素值为0。

参数indices表示张量的索引，参数updates表示分散到输出张量的值，参数shape表示输出张量的形状。

函数使用

1、一维张量

>>> indices = tf.constant([[1], [3], [5], [7]])
>>> updates = tf.constant([2, 3, 4, 5])
>>> shape = tf.constant([8])
>>> tf.scatter_nd(indices, updates, shape)
<tf.Tensor: shape=(8,), dtype=int32, numpy=array([0, 2, 0, 3, 0, 4, 0, 5])>

2、二维张量

>>> indices = tf.constant([[1], [2]])
>>> updates = tf.constant([[1, 2], [3, 4]])
>>> updates
<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4]])>
>>> shape = tf.constant([4, 2])
>>> tf.scatter_nd(indices, updates, shape)
<tf.Tensor: shape=(4, 2), dtype=int32, numpy=
array([[0, 0],
       [1, 2],
       [3, 4],
       [0, 0]])>

3、三维张量

>>> indices = tf.constant([[0], [2]])
>>> updates = tf.constant([[[5, 5, 5, 5], [6, 6, 6, 6],
                            [7, 7, 7, 7], [8, 8, 8, 8]],
                           [[5, 5, 5, 5], [6, 6, 6, 6],
                            [7, 7, 7, 7], [8, 8, 8, 8]]])
>>> updates
<tf.Tensor: shape=(2, 4, 4), dtype=int32, numpy=
array([[[5, 5, 5, 5],
        [6, 6, 6, 6],
        [7, 7, 7, 7],
        [8, 8, 8, 8]],

       [[5, 5, 5, 5],
        [6, 6, 6, 6],
        [7, 7, 7, 7],
        [8, 8, 8, 8]]])>
>>> shape = tf.constant([4, 4, 4])
>>> tf.scatter_nd(indices, updates, shape)
<tf.Tensor: shape=(4, 4, 4), dtype=int32, numpy=
array([[[5, 5, 5, 5],
        [6, 6, 6, 6],
        [7, 7, 7, 7],
        [8, 8, 8, 8]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[5, 5, 5, 5],
        [6, 6, 6, 6],
        [7, 7, 7, 7],
        [8, 8, 8, 8]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]]])>