Как подсчитать элементы, принадлежащие одному классу меток, в пользовательской функции потери keras?

Я ищу способ подсчитать количество вхождений каждого класса в массиве y_true в пользовательской функции потерь и заменить каждый элемент в массиве соответствующим количеством вхождений.

Я уже реализовал решение numpy, но я не могу перевести его в keras (с бэкэндом tf).

Пример ввода:

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

Импорт:

import numpy as np
from keras import backend as k

Непонятная реализация:

def custom_loss(y_true, y_pred):
    bins = np.bincount(y_true)
    y_true_counts = bins[y_true]

    >>> y_true_counts: [2 3 3 3 2 1]

Реализация Кераса:

def custom_loss(y_true, y_pred)
    bins = k.tf.bincount(y_true)
    y_true_counts = bins[y_true]

Хотя решение numpy работает нормально, когда я хочу оценить реализацию keras, я получаю следующую ошибку:

a = custom_loss(y_true, y_pred)
>>> InvalidArgumentError: Shape must be rank 1 but is rank 2 for 'strided_slice_4' (op: 'StridedSlice') with input shapes: [?], [1,6], [1,6], [1].
[...]
----> 3     y_true_counts = bins[y_true]
[...]