How do I do this tensor transformation and preserve the gradients?

Question

In Tensorflow, I have a float tensor T with shape [batch_size, 3]. For example, T[0] = [4, 4, 3].

I want to turn that into a size 5 one hot in order to yield entries from an embedding dictionary. In the above case, that would look like

T[0] = [[0, 0, 0, 0, 1], [0, 0, 0, 0, 1], [0, 0, 0, 1, 0]].

If I can get it into this format, then I can multiply it by the embedding dictionary. However, this is in the middle of the graph and I need the gradients to flow through it. Is there a clever way to use stop_gradient a la How Can I Define Only the Gradient for a Tensorflow Subgraph? to make this work? I'm coming up short.

Answer 1

I was able to solve this in the following way:

expanded = tf.expand_dims(inputs, 2)
embedding_input = tf.cast(tf.one_hot(tf.to_int32(inputs), 5), inputs.dtype)
embedding_input = tf.stop_gradient(embedding_input - expanded) + expanded