python之TensorFlow:带轴选项的 bincount

在 TensorFlow 中，我可以使用 tf.bincount 获取数组中每个元素的计数:

x = tf.placeholder(tf.int32, [None]) 
freq = tf.bincount(x) 
tf.Session().run(freq, feed_dict = {x:[2,3,1,3,7]}) 

返回

Out[45]: array([0, 1, 1, 2, 0, 0, 0, 1], dtype=int32) 

有没有办法在二维张量上做到这一点？即

x = tf.placeholder(tf.int32, [None, None]) 
freq = tf.axis_bincount(x, axis = 1) 
tf.Session().run(freq, feed_dict = {x:[[2,3,1,3,7],[1,1,2,2,3]]}) 

返回

[[0, 1, 1, 2, 0, 0, 0, 1],[0, 2, 2, 1, 0, 0, 0, 0]] 

请您参考如下方法:

我发现一个简单的方法是利用广播将张量中的所有值与模式 [0, 1, ..., length - 1] 进行比较，然后计算沿所需轴的“命中”次数。

即:

def bincount(arr, length, axis=-1): 
  """Count the number of ocurrences of each value along an axis.""" 
  mask = tf.equal(arr[..., tf.newaxis], tf.range(length)) 
  return tf.math.count_nonzero(mask, axis=axis - 1 if axis < 0 else axis) 
 
x = tf.convert_to_tensor([[2,3,1,3,7],[1,1,2,2,3]]) 
bincount(x, tf.reduce_max(x) + 1, axis=1) 

返回:

<tf.Tensor: id=406, shape=(2, 8), dtype=int64, numpy= 
array([[0, 1, 1, 2, 0, 0, 0, 1], 
       [0, 2, 2, 1, 0, 0, 0, 0]])>