How to make a Tensor that has the number of times its indexes are in another Tensor?-CodePudding

Say I have a Tensor I of shape (batch_size x N) where values are integers < max_indexes.

I want, without using any loop, to create a Tensor A of shape (batch_size x max_indexes) where A[batch_num,i] is the number of times i appears in I[batch_num].

For example (batch_size=3, N=4, max_indexes=9):

If I have

>>> I
<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[2, 6, 5, 5],
       [1, 2, 0, 3],
       [4, 2, 1, 0]])>

I want

>>> A
<tf.Tensor: shape=(3, 9), dtype=int32, numpy=
array([[0, 0, 1, 0, 0, 2, 1, 0, 0],
       [1, 1, 1, 1, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 1, 0, 0, 0, 0]])>

Anyone has any idea to solve this? I did this already using a for loop*, but it is a bit time-consuming...

*For loop implementation:

repeat_range = tf.transpose(tf.repeat([tf.range(max_indexes)],batch_size,0))
A = tf.zeros((batch_size,max_indexes))
for i in range(I.shape[1]):
    A  = tf.transpose(tf.cast(repeat_range == I[:,i],tf.float32))

CodePudding user response：

Found the answer! Need to use tf.math.bincount:

A = tf.math.bincount(I, axis=-1)

Note, if you need the second dimension of A being max_indexes, you can pad like this:

A = tf.math.bincount(a, axis=-1, minlength=max_indexes, maxlength=max_indexes)