Does Tensorflow have an inverse of tf.unravel_index?

user_635523 0
Viewed 586

tf.unravel_index takes a flat index into an array and a shape, and returns the tuple that represents that index in the array.

with tf.Session() as sess: 
    t = tf.unravel_index(tf.constant(15), [2, 3, 4]) 
    print(sess.run(t)) # [1 0 3]

Is there an inverse? tf.ravel_index doesn't exist, but maybe it's under a different name?

tensorflow
1 Answers

Here is how you can do the equivalent to np.ravel_multi_index in TensorFlow:

import tensorflow as tf

# multi_index is 2D (num dimensions x num indices), dims is 1D
# Does not check for out of bounds indices
def tf_ravel_multi_index(multi_index, dims):
    strides = tf.cumprod(dims, exclusive=True, reverse=True)
    return tf.reduce_sum(multi_index * tf.expand_dims(strides, 1), axis=0)

# Test
with tf.Graph().as_default(), tf.Session() as sess:
    # Shape tensor
    shape = tf.constant([4, 5, 6])
    # Some flat indices
    idx = tf.constant([23, 56, 4, 17])
    # Unravel indices
    unravel = tf.unravel_index(idx, shape)
    # Ravel indices again
    idx2 = tf_ravel_multi_index(unravel, shape)
    # Check result
    print(*sess.run((unravel, idx2)), sep='\n')
    # [[0 1 0 0]
    #  [3 4 0 2]
    #  [5 2 4 5]]
    # [23 56  4 17]
user_17827920
Related