How do I select multiple rows with Tensorflow?

0
votes

I have an tensor of probabilities, and I want to select the highest probability per column.

import tensorflow as tf

x = tf.random.normal(mean=0, stddev=1, shape=(5, 4))

<tf.Tensor: shape=(5, 4), dtype=float32, numpy=
array([[ 0.29182285, -0.30140358, -1.6745052 ,  0.04754949],
       [-0.4013166 ,  0.98574334, -0.33083338, -1.2661675 ],
       [-0.49397126, -0.19175254,  1.4922557 , -0.9857015 ],
       [ 0.57078785,  0.30159777,  0.34956333,  0.17599773],
       [ 1.0847769 , -0.58611375, -0.10727347, -0.13440256]],
      dtype=float32)>

I'm trying to select the highest probabilities per column like I would with numpy and it doesn't work:

x[tf.argmax(x, axis=0), :]

TypeError: Only integers, slices (:), ellipsis (...), tf.newaxis (None) and scalar tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor: shape=(4,), dtype=int32, numpy=array([4, 1, 2, 3])>

What would be the correct way to do this?

2
pythonnumpytensorflow

2 Answers

1
votes

Convert Tensor to numpy first before using it as matrix

x.numpy()[tf.argmax(x, axis=0), :]
1
votes

Today I learned that tf.gather() does that.

tf.gather(x, tf.argmax(x, axis=0))

<tf.Tensor: shape=(4, 4), dtype=float32, numpy=
array([[ 1.8891758 ,  0.7073202 , -0.78521085, -2.7632885 ],
       [-1.3851309 ,  1.4023514 , -0.9735394 , -0.81982684],
       [-0.22595228, -0.7155944 ,  0.37807527,  2.2081604 ],
       [-0.22595228, -0.7155944 ,  0.37807527,  2.2081604 ]],
      dtype=float32)>