如何制作一个张量，其索引在另一个张量中的次数？

Question

假设我有一个形状为 (batch_size x N) 的张量 I，其中值是整数 < max_indexes.

我想不使用任何循环，创建一个形状为 (batch_size x max_indexes) 的张量 A，其中 A[batch_num,i]是i在I[batch_num]中出现的次数。

例如(batch_size=3, N=4, max_indexes=9):

如果我有

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

我要

>>> 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]])>

有人知道解决这个问题吗？我已经使用 for 循环* 完成了此操作，但这有点耗时...

*For循环实现：

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))

Answer 1

找到答案了！需要使用tf.math.bincount:

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

注意，如果你需要A的第二个维度是max_indexes，你可以这样填充：

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

如何制作一个张量，其索引在另一个张量中的次数？

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

python

indexing

counting

tensorflow