Pick random entry from two ragged tensors

How do I pick random entries from two ragged tensors? For example,

c = tf.ragged.constant([[1, 2, 3], [4, 5]])
v = tf.ragged.constant([[10., 20., 30.], [40., 50.]])

r = tf.random.uniform([1, 1], maxval=2, dtype=tf.int32)

with tf.Session() as sess:
    print(sess.run([tf.gather_nd(c, r), tf.gather_nd(v, r)]))

gives me either [1, 2, 3] and [10., 20., 30.], or [4, 5] and [40., 50.]. But now I want to pick a random number between 0 and the length of the returned list, and get the corresponding entries from both lists. And then I want to batch this entire process.

Thank you in advance for your help!

Solution

Here is an example based on the values you've given: (I am using TF 1.13)

import tensorflow as tf
tf.enable_eager_execution() # you can use a normal Session, but this is to show intermediate output

c = tf.ragged.constant([[1, 2, 3], [4, 5]])
v = tf.ragged.constant([[10., 20., 30.], [40., 50.]])

r = tf.random.uniform([1, 1], maxval=2, dtype=tf.int32)

a = tf.gather_nd(c, r)
b = tf.gather_nd(v, r)
print(a)
print(b)

# Output example
#<tf.RaggedTensor [[1, 2, 3]]>
#<tf.RaggedTensor [[10.0, 20.0, 30.0]]>

# Lengths
l_a = tf.squeeze(a.row_lengths())
l_b = tf.squeeze(b.row_lengths())
print(l_a)
print(l_b)

#Output example
#tf.Tensor(3, shape=(), dtype=int64)
#tf.Tensor(3, shape=(), dtype=int64)

#Random index between 0 and length
rand_idx_a = tf.random.uniform([1],minval=0,maxval=l_a,dtype=tf.int64)
rand_idx_b = tf.random.uniform([1],minval=0,maxval=l_b,dtype=tf.int64)
print(rand_idx_a)
print(rand_idx_b)

#Output example
#tf.Tensor([0], shape=(1,), dtype=int64)
#tf.Tensor([2], shape=(1,), dtype=int64)

#Convert ragged tensor to tensor of shape [1,n]
t_a = a.to_tensor()
t_b = b.to_tensor()
print(t_a)
print(t_b)

#Read from extracted tensor using random index
rand_a = tf.gather_nd(tf.squeeze(t_a),rand_idx_a) #removes dimension of 1
rand_b = tf.gather_nd(tf.squeeze(t_b),rand_idx_b)
print(rand_a)
print(rand_b)

#Output example
#tf.Tensor([[1 2 3]], shape=(1, 3), dtype=int32)
#tf.Tensor([[10. 20. 30.]], shape=(1, 3), dtype=float32)
#tf.Tensor(1, shape=(), dtype=int32)
#tf.Tensor(30.0, shape=(), dtype=float32)

All of these operations can be batched easily depending on your input.