Первый подход: использование tf.one_hot
и tf.boolean_mask
:
# shape = (?,1,3)
mask_idx = 1- tf.one_hot(idx,a.shape[1])
# shape = (?,3)
result = tf.boolean_mask(a,mask_idx[:,0,:])
# shape = (?,2,3)
result = tf.reshape(result,shape=(-1,a.shape[1]-1,a.shape[2]))
Второй подход: использование tf.map_fn
:
result = tf.map_fn(lambda x: tf.boolean_mask(x[0],1 - tf.one_hot(tf.squeeze(x[1]),a.shape[1]))
, [a,idx]
, dtype=tf.int32)
Пример:
import tensorflow as tf
a = tf.constant([[[1,2,3],[4,5,6],[7,8,9]],
[[9,8,7],[6,5,4],[3,2,1]],
[[0,8,0],[1,5,4],[3,1,1]]],dtype=tf.int32)
idx = tf.constant([[1],[0],[2]],dtype=tf.int32)
# First approach:
# shape = (?,1,3)
mask_idx = 1- tf.one_hot(idx,a.shape[1])
# shape = (?,3)
result = tf.boolean_mask(a,mask_idx[:,0,:])
# shape = (?,2,3)
result = tf.reshape(result,shape=(-1,a.shape[1]-1,a.shape[2]))
# Second approach:
result = tf.map_fn(lambda x: tf.boolean_mask(x[0],1 - tf.one_hot(tf.squeeze(x[1]),a.shape[1]))
, [a,idx]
, dtype=tf.int32)
with tf.Session() as sess:
print(sess.run(result))
# print
[[[1 2 3]
[7 8 9]]
[[6 5 4]
[3 2 1]]
[[0 8 0]
[1 5 4]]]