Sample from ragged tensor-CodePudding

I have a raggedTensor of row_lens going from 1 to up to 10k. I would like to select elements randomly from it with an upper limit on the number per row in a scalable way. Like in this example:

vect = [[1,2,3],[4,5][6],[7,8,9,10,11,12,13]]
limit = 3
sample(vect, limit)

-> output: [[1,2,3],[4,5],[6],[7,9,11]]

My idea was to select * in case len_row < limit and randomly in the other case. I wonder if this can be done with less than batch_size complexity with some tensorflow operations?

CodePudding user response：

You can try using tf.map_fn in graph mode:

import tensorflow as tf

vect = tf.ragged.constant([[1,2,3],[4,5],[6],[7,8,9,10,11,12,13]])

@tf.function
def sample(x, samples=3):
  length = tf.shape(x)[0]
  x = tf.cond(tf.less_equal(length, samples), lambda: x, lambda: tf.gather(x, tf.random.shuffle(tf.range(length))[:samples]))
  return x

c = tf.map_fn(sample, vect)

<tf.RaggedTensor [[1, 2, 3], [4, 5], [6], [12, 7, 9]]>

Note that tf.vectorized_map would probably be faster, but there is a current bug regarding this function and ragged tensors. The use of tf.while_loop is also an option.