tensorflow2之切片

• tf.gather：按axis和indices获取索引对应的tensor

a = tf.Variable([[1,2,3,4,5], [6,7,8,9,10], [11,12,13,14,15]])
idx_a = tf.Variable([0, 2])
tf.gather(a, idx_a)
# 返回： [[1,2,3,4,5],  [11,12,13,14,15]]

tf.gather(a, idx_a, axis=1)
# 返回： [[1,3],  [6,8], [11,13]]

• tf.gather_nd：按indices获取params对应的元素

a = tf.Variable([[1,2,3,4,5], [6,7,8,9,10], [11,12,13,14,15]])
idx_b = tf.Variable([[0, 1], [0, 3], [2, 4]])
tf.gather_nd(a, idx_b)
# 返回： [2, 4, 15]

idx_c = [[0], [1]]
tf.gather_nd(a, idx_c)
# 返回：[[1,2,3,4,5], [6,7,8,9,10]]

• tf.boolean_mask：有点像np.where

c = [[1, 2], [3, 4], [5, 6]]
mask = np.array([True, True, False])
tf.boolean_mask(c, mask)
# 返回：[[1, 2], [3, 4]]

c = tf.constant([[-1,1,-1],[2,2,-2],[3,-3,3]],dtype=tf.float32)
tf.boolean_mask(c,c<0)
c[c<0]
#  返回结果相同：均为获取张量`c`小于零的元素，返回一维数组

• tf.where：
  • 返回True的idx
  • True则返回x对应的元素，否则返回y对应的元素

tf.where([True, False, True, True])
# 返回： [[0], [2], [3]]
tf.where(([[True, False], [False, True], [True, True]]))
# 返回：[[0, 0], [1, 1], [2, 0], [2, 1]]

tf.where([True, False, False, True], [1,2,3,4], [100,200,300,400])
# 返回：[1, 200, 300, 4]]
tf.where([True, False, False, True], [1,2,3,4], [100])
# 返回：[1, 100, 100, 4]
tf.where([True, False, False, True], [1,2,3,4], 100)
# 返回：[1, 100, 100, 4]
tf.where([True, False, False, True], 1, 100)
# 返回：[1, 100, 100, 1]

c = tf.constant([[-1,1,-1],[2,2,-2],[3,-3,3]],dtype=tf.float32)
# 返回符合条件的位置坐标
tf.where(c < 0)
# 将小于零的值替换为nan
tf.where(c<0,tf.fill(c.shape,np.nan),c) 

• tf.scatter_nd：替换，更新原来的tensor

indices = tf.constant([[4], [3], [1], [7]])
updates = tf.constant([9, 10, 11, 12])
# 用update来填充indices至目标shape
tf.scatter_nd(indices, updates, shape=tf.constant([8]))

# 除[0,0],[2,1]外，其它位置元素全部用0填充
tf.scatter_nd([[0,0],[2,1]],[c[0,0],c[2,1]],c.shape)