我可以在没有 for 循环的情况下对 3 阶张量应用 tf.math.bincount 吗?

Can I apply tf.math.bincount on a rank 3 tensor without for loop?

我想获取 3 维张量的出现次数(显示在相应索引的张量中,就像 tf.math.bincount 一样)。

对于二维张量,您可以简单地这样做:

T = tf.round(25*tf.random.uniform((5,8)))
bincounts = tf.cast(tf.math.bincount(T, axis=-1),tf.float32)

但在 3d 张量上,我发现的唯一方法是在三维上循环,如下所示:

third_dim  = 10
T = tf.round(25*tf.random.uniform((5,8,third_dim)))
bincounts = []
for i in range(third_dim):
    bincounts.append(tf.math.bincount(T[:,:,i], axis=-1))
bincounts = tf.stack(bincounts,-1)

有谁知道有没有办法直接在所有维度上应用这样的函数?

我找到了一个方法:

对重塑后的张量应用 bincounts,然后重塑回您想要的形状:

third_dim  = 10
T = tf.round(25*tf.random.uniform((5,8,third_dim)))
T2 = tf.reshape(T,(5*8,third_dim))
bincounts2 = tf.math.bincount(T2, axis=-1)
bincounts = tf.reshape(bincounts2, [5,8,bincounts2.shape[-1]])