title: How to get the indices of the maximum value in a tensor using Tensorflow?
body: <p>I am looking for the best and optimized way (without loops) to get the indices of the maximum value of tensor (Rank 2 tensor) using TensorFlow 1.14, in Google Colab i tried this code using TF2.0:</p>
  <pre><code>import tensorflow as tf
  A = tf.constant([[0.2,0.8],[0.3,0.9],[0.4,0.7],[0.5,0.4]])
  b = tf.math.argmax(A,0)
  bb = b.numpy()
  </code></pre>
  <p>The indices of the max here is [1,1], but the problem is i have to give the axis as input, and it doesn't give me the right one even i change the axis.</p>