如何在稀疏张量的每一行中获得值的数量?

Question

我有一个稀疏张量如下:

st = tf.sparse.from_dense([[1, 0, 2, 5], [3, 0, 0, 4], [0, 0, 0, 0], [1, 1, 3, 0], [1, 2, 2, 2]])
print(st)

SparseTensor(indices=tf.Tensor(
[[0 0]
[0 2]
[0 3]
[1 0]
[1 3]
[3 0]
[3 1]
[3 2]
[4 0]
[4 1]
[4 2]
[4 3]], shape=(12, 2), dtype=int64), values=tf.Tensor([1 2 5 3 4 1 1 3 1 2 2 2], shape=(12,), dtype=int32), dense_shape=tf.Tensor([5 4], shape=(2,), dtype=int64))

我想将这个稀疏张量转换为形状为(5, 1)的另一个1D张量，其中唯一的列表示每一行中值的数量(或大小)。

例如，对于上面的稀疏张量，期望的1D张量为[3, 2, 0, 3, 4]。

我该怎么做?

我试着通过TensorFlow API文档，但找不到任何尝试。

Answer 1

您可以在索引上使用bin count。

tf.math.bincount(tf.cast(st.indices[:,0], tf.int32))