TensorFlow: bincount с опцией оси

Решение для этого задается для массива numpy: Применяет bincount к каждой строке массива 2D numpy. Сделайте каждую строку уникальной, добавив row_id * (max + 1) в каждую строку, а затем найдите bincount для сплющенного 1d-массива и затем соответствующим образом bincount его.

Для TF внесите следующие изменения:

x = tf.placeholder(tf.int32, [None, None])

max_x_plus_1 = tf.reduce_max(x)+1
ids = x + max_x_plus_1*tf.range(tf.shape(x)[0])[:,None]
out = tf.reshape(tf.bincount(tf.layers.flatten(ids), 
                 minlength=max_x_plus_1*tf.shape(x)[0]), [-1, N])

tf.Session().run(out, 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]])>

Я сам нуждался в этом и написал для него небольшую функцию, так как нет официальной реализации.

def bincount(tensor, minlength=None, axis=None):
    if axis is None:
        return tf.bincount(tensor, minlength=minlength)
    else:
        if not hasattr(axis, "__len__"):
            axis = [axis]

        other_axis = [x for x in range(0, len(tensor.shape)) if x not in axis]
        swap = tf.transpose(tensor, [*other_axis, *axis])
        flat = tf.reshape(swap, [-1, *np.take(tensor.shape.as_list(), axis)])
        count = tf.map_fn(lambda x: tf.bincount(x, minlength=minlength), flat)
        res = tf.reshape(count, [*np.take([-1 if a is None else a for a in tensor.shape.as_list()], other_axis), minlength])
        return res

Там много обработки для различных крайних случаев.

Суть этого решения состоит в следующем:

swap = tf.transpose(tensor, [*other_axis, *axis])
flat = tf.reshape(swap, [-1, *np.take(tensor.shape.as_list(), axis)])
count = tf.map_fn(lambda x: tf.bincount(x, minlength=minlength), flat)

• Операция transpose перемещает все оси, которые вы хотите bincount к концу тензора. Например, если бы у вас была матрица, которая выглядит как [100, 50, 20] с осью [0, 1, 2] и вы хотели бы bincount к оси 1, эта операция поменяет местами ось 1 до конца, и вы получите [100, 20, 50] матрица.
• Операция reshape выравнивает все остальные оси, для которых не требуется bincount к одному измерению/оси.
• Операция map_fn отображает bincount на каждую запись сплющенного измерения/оси.

Вы должны указать параметр minlength. Это необходимо, чтобы все результаты bincount имели одинаковую длину (иначе матрица не будет иметь правильную форму). Вероятно, это максимальное значение для вашего tensor. Для меня было лучше передать его как параметр, так как у меня уже было это значение, и мне не нужно было его извлекать, но вы также можете рассчитать его с помощью tf.reduce_max(tensor).

Полное решение дополнительно изменяет его, чтобы восстановить другие оси. Он также поддерживает несколько осей и одну ось None в тензоре (для дозирования).

tf.bincount() принимает массив как аргумент, но он агрегирует счетчик по массиву и не работает вдоль некоторой оси в настоящий момент. Например:

In [27]: arr
Out[27]: 
array([[2, 3, 1, 3, 7],
       [1, 1, 2, 2, 3]], dtype=int32)

In [28]: x = tf.placeholder(tf.int32, [None, None])
    ...: freq = tf.bincount(x)
    ...: tf.Session().run(freq, feed_dict = {x:arr})

# aggregates the count across the whole array
Out[28]: array([0, 3, 3, 3, 0, 0, 0, 1], dtype=int32)
# 0 occurs 0 times
# 1 occurs 3 times
# 2 occurs 3 times
# 3 occurs 3 times and so on..

Итак, по крайней мере, на данный момент нет способа передать информацию о оси в tf.bincount().

Однако немного неэффективным способом было бы передать одну строку за раз до tf.bincount() и получить результаты. И затем, наконец, объедините полученные результирующие 1D массивы как массив желаемой размерности.

Я не уверен, что это самый эффективный способ, но в любом случае это один из способов циклы над тензором (вдоль оси 0)

In [3]: arr = np.array([[2, 3, 1, 3, 7], [1, 1, 2, 2, 3]], dtype=np.int32)
In [4]: sess = tf.InteractiveSession()

In [5]: for idx, row in enumerate(tf.unstack(arr)):
   ...:     freq = tf.bincount(row)
   ...:     print(freq.eval())
   ...:     
[0 1 1 2 0 0 0 1]
[0 2 2 1]