TensorFlow: bincount with axis option

Question

In TensorFlow, I can get the count of each element in an array with tf.bincount:

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

this returns

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

Is there a way to do this on a 2D tensor? i.e.

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

that returns

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

Answer 1

A simple way I found to do this is to take advantage of broadcasting to compare all values in the tensor against the pattern [0, 1, ..., length - 1], and then count the number of "hits" along the desired axis.

Namely:

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)

returns:

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