在张量流中,从张量沿0轴随机(子)采样k项

给定张量 rank>=1 T ,我想随机抽样 k 沿0轴均匀输入。

编辑: 采样应该作为惰性操作作为计算图的一部分,并且应该在每次调用时输出不同的随机项。

例如,给定 T 等级 2 :

T = tf.constant( \
     [[1,1,1,1,1],
      [2,2,2,2,2],
      [3,3,3,3,3],
      ....
      [99,99,99,99,99],
      [100,100,100,100,100]] \
     )

用 k=3 ,可能的输出为:

#output = \
#   [[34,34,34,34,34],
#    [6,6,6,6,6],
#    [72,72,72,72,72]]

如何在tensorflow中实现这一点?