Argmax on a tensor and ceiling in Tensorflow

Question

Suppose I have a tensor in Tensorflow that its values are like:

A = [[0.7, 0.2, 0.1],[0.1, 0.4, 0.5]]

How can I change this tensor into the following:

B = [[1, 0, 0],[0, 0, 1]] 

In other words I want to just keep the maximum and replace it with 1.
Any help would be appreciated.

Camilacamile

Answer 1

I think that you can solve it with a one-liner:

import tensorflow as tf
import numpy as np

x_data = [[0.7, 0.2, 0.1],[0.1, 0.4, 0.5]]
# I am using hard-coded dimensions for simplicity
x = tf.placeholder(dtype=tf.float32, name="x", shape=(2,3))

session = tf.InteractiveSession()

session.run(tf.one_hot(tf.argmax(x, 1), 3), {x: x_data})

The result is the one that you expect:

Out[6]: 
array([[ 1.,  0.,  0.],
       [ 0.,  0.,  1.]], dtype=float32)

Dominations

It was smart! Thanks – Camilacamile