Tensorflow从每行中随机采样

Question

假设我有一个形状为(m, n)的张量A，我想从每一行中随机采样k个元素(没有替换)，得到一个形状为(m, k)的张量B。如何在tensorflow中做到这一点？

下面是一个例子：

A：[1,2,3，4,5,6，7,8,9，10,11,12]

k：2

B：[1,3,5,6,9,8,12,10]

Answer 1

这是实现这一点的一种方法：

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    tf.random.set_random_seed(0)
    a = tf.constant([[1,2,3], [4,5,6], [7,8,9], [10,11,12]], tf.int32)
    k = tf.constant(2, tf.int32)
    # Tranpose, shuffle, slice, undo transpose
    aT = tf.transpose(a)
    aT_shuff = tf.random.shuffle(aT)
    at_shuff_k = aT_shuff[:k]
    result = tf.transpose(at_shuff_k)
    print(sess.run(result))
    # [[ 3  1]
    #  [ 6  4]
    #  [ 9  7]
    #  [12 10]]