4

I'm looking for a way in tensorflow to, given two inputs:

  1. input1, a 3D tensor of shape (batch_size, x, y)
  2. input2, a 1D tensor of shape (batch_size,) whose values are all in the range [0, y - 1] (inclusive).

return a 2D tensor of shape (batch_size, x) such that the ith element in the output is equal to the input2[i]-th column of the ith element in input1.

Example: If input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]] (so shape of input1 is (3, 2, 2)) and

input2 = [0, 1, 1], then the output I want is [[1,3], [6,8], [10,12]].

Explanation: The 0th element in the output is [1,3] because the 0th element in input2 is 0; so, it becomes the 0th column in the 0th element of input1. The last element in the output is [6,8], because the last element in input2 is 1; so, it becomes the 1st column in the last element of input1.

Attempts:

I tried using tf.one_hot to accomplish this, (tf.reduce_sum(input1 * tf.one_hot(input2, y), 2)) but Tensorflow became unhappy when doing the multiplication, saying "ValueError: Dimensions must be equal, but are 2 and 3 for 'mul' (op: 'Mul') with input shapes: [3,2,2], [3,2]."

Any help would be super appreciated, thanks!

Eskapp
  • 3,419
  • 2
  • 22
  • 39
Kevin
  • 133
  • 8
  • Hi Kevin, I'm a software engineer at Google and I'd like to include a screenshot of this question in a research paper we're writing, to showcase some representative questions about TensorFlow. If that's alright with you, could you release your question under a permissive license? A comment saying e.g. "I license this StackOverflow question under the Apache License 2.0" would be sufficient for us. Thanks, David – David Bieber Apr 07 '20 at 19:47
  • I license this StackOverflow question under the Apache License 2.0. Out of curiosity, could you drop a link to the paper when it's completed? Best of luck. – Kevin Apr 08 '20 at 21:11
  • Thanks Kevin. The paper is now available at [https://arxiv.org/abs/2003.09040](https://arxiv.org/abs/2003.09040). You can read more about TF-Coder [here](https://blog.tensorflow.org/2020/08/introducing-tensorflow-coder-tool.html), and you can use it yourself [on Colab here](https://colab.sandbox.google.com/github/google-research/tensorflow-coder/blob/master/TF-Coder_Colab.ipynb#scrollTo=DCvkqaXs4fyU&line=1&uniqifier=1). – David Bieber Aug 28 '20 at 14:43

1 Answers1

1

You can use tf.map_fn() to achieve it.

import tensorflow as tf
import numpy as np

input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]]
input2 = [0, 1, 1]

tf_input1 = tf.placeholder(shape=(None,2,2),dtype=tf.int32)
tf_input2 = tf.placeholder(shape=(None),dtype=tf.int32)

result = tf.map_fn(lambda x: x[0][:,x[1]], [tf_input1,tf_input2], dtype=tf.int32)

with tf.Session()as sess:
    result = sess.run(result,feed_dict={tf_input1:np.array(input1)
        ,tf_input2:np.array(input2)})
    print(result)

# print
[[ 1  3]
 [ 6  8]
 [10 12]]

Edit

tf.map_fn() is slow compared to the vectorized operation. I added a matrix multiplication operation.

# shape= (3,2,1)
result = tf.cast(tf.expand_dims(tf.one_hot(input2, 2),-1),tf.int32)
# shape= (3,2)
result = tf.squeeze(tf.matmul(tf_input1, result))
giser_yugang
  • 6,058
  • 4
  • 21
  • 44