calculating mode and mode-count along an axis of a given tensor in tensorflow

Question

I have a 2-D array as:

c=np.array([[1, 3, 0, 0, 3],
            [1, 3, 1, 0, 2],
            [1, 3, 1, 2, 2]])

I want to calculate mode and mode-count along axis=0.
Therefore result should look like this:

mode = [1,3,1,0,2], mode-count=[3,3,2,2,2]

I have searched in the TensorFlow website, but cannot find any useful API other than tf.unique_with_counts , which expects a 1-D tensor.

I do not want to run a loop over each column of array c in order to use tf.unique_with_counts for computing mode and mode-count.
Any suggestions with examples are most welcome.

Answer 1

TensorFlow Probability has a tfp.stats.count_integers function that can make this fairly straightforward:

import tensorflow as tf
import tensorflow_probability as tfp

def mode_and_counts(x, axis=-1):
    x = tf.convert_to_tensor(x)
    dt = x.dtype
    # Shift input in case it has negative values
    m = tf.math.reduce_min(x)
    x2 = x - m
    # minlength should not be necessary but may fail without it
    # (reported here https://github.com/tensorflow/probability/issues/962)
    c = tfp.stats.count_integers(x2, axis=axis, dtype=dt,
                                 minlength=tf.math.reduce_max(x2) + 1)
    # Find the values with largest counts
    idx = tf.math.argmax(c, axis=0, output_type=dt)
    # Get the modes by shifting by the subtracted minimum
    modes = idx + m
    # Get the number of counts
    counts = tf.math.reduce_max(c, axis=0)
    # Alternatively, you could reuse the indices obtained before
    # with something like this:
    #counts = tf.transpose(tf.gather_nd(tf.transpose(c), tf.expand_dims(idx, axis=-1),
    #                                   batch_dims=tf.rank(c) - 1))
    return modes, counts

# Test
x = tf.constant([[1, 3, 0, 0, 3],
                 [1, 3, 1, 0, 2],
                 [1, 3, 1, 2, 2]])
tf.print(*mode_and_counts(x, axis=0), sep='\n')
# [1 3 1 0 2]
# [3 3 2 2 2]
tf.print(*mode_and_counts(x, axis=1), sep='\n')
# [0 1 1]
# [2 2 2]

Answer 2

c=np.array([[1, 3, 0, 0, 3],
            [1, 3, 1, 0, 2],
            [1, 3, 1, 2, 2]])
c = tf.constant(c)

mode

tf.map_fn(lambda x: tf.unique_with_counts(x).y[tf.argmax(tf.unique_with_counts(x).count, output_type=tf.int32)], tf.transpose(c))

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

mode-count

tf.map_fn(lambda x: tf.reduce_max(tf.unique_with_counts(x).count), tf.transpose(c))

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

I do this simply using unique_with_counts in a map_fn