In tensorflow, randomly (sub)sample k entries from a tensor along 0-axis

Question

Given a tensor of rank>=1 T, I would like to randomly sample k entries from it, uniformly, along the 0-axis.

EDIT: The sampling should be part of the computation graph as a lazy operation and should output different random entries every time it is called.

For instance, given T of rank 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]] \
     )

With k=3, a possible output would be:

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

How can this be achieved in tensorflow?

Answer 1

You can use a random shuffle on an array of indices:

Take the first sample_num indices, and use them to pick slices of your input.

idxs = tf.range(tf.shape(inputs)[0])
ridxs = tf.random.shuffle(idxs)[:sample_num]
rinput = tf.gather(inputs, ridxs)