Преобразование numpy в тензорные операции с одного горячего на разреженный

Question

new_one_hot = None
for i in range(num_channels):
    new_channel = one_hot[:,:,i] * (i+1)
    if new_one_hot is None:
        new_one_hot = new_channel
    else:
        new_one_hot = np.dstack((new_one_hot, new_channel))

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

Я немного застрял - как я могу преобразовать вышеперечисленное в операции тензорного потока? Могу ли я использовать for-l oop? Я не вижу, чтобы тензорный поток использовал циклы for.

А как насчет индексации и dstack? Немного смущен этим.

Вот моя попытка ниже:

def convert_one_hot_to_sparse_tf(one_hot):
    num_channels = one_hot.get_shape()[2]
    new_one_hot = None
    for i in range(num_channels):
        new_channel = one_hot[:,:,i] * tf.constant(i+1)
        if new_one_hot is None:
            new_one_hot = new_channel
        else:
            new_one_hot = tf.stack([new_one_hot, new_channel], axis=-1)

    sparse_one_hot = tf.math.argmax(new_one_hot, axis=-1)
    return sparse_one_hot

но я получаю "Фигуры должны быть одинакового ранга, но имеют 3 и 2 Из слияния фигуры 0 с другими фигурами . для 'stack_3' (op: 'Pack') с формами ввода: [ширина, высота, 2], [ширина, высота]. "

Ввод функции должен быть [ширина, высота, num_channels ] Мой желаемый результат будет иметь форму [ширина, высота]