Керас, слой, преобразующий выходные данные слоя Softmax в категориальный массив?

#tensorflow #keras

Вопрос:

Мне нужен слой, который преобразует softmax в категориальный вектор (а также выбирает категорию «случайным образом», используя вероятность softmax), вот мой код на данный момент:

 from tensorflow import keras
from tensorflow.keras import layers

controller_input = keras.Input(shape=(1,1,))

x, x2, xc = layers.LSTM(100, return_sequences=True, return_state=True)(controller_input)

x = layers.Dense(3, activation="softmax")(x) # output shape (None, 1, 3)
 

Теперь предположим , что значение x равно [ 0.1 0.2 0.7 ] , мне нужен слой, который преобразует его в категориальный вектор

Вот эквивалент с Numpy :

 x = [0.1, 0.2, 0.7] 
choice = np.random.choice(len(x), p=x)   # let say that choice = 2
vector = utils.to_categorical(choice, len(x), dtype="int") # then vector = [ 0 0 1 ]
 

Кроме того, форма вывода этого нового слоя должна совпадать с формой вывода плотного слоя ((None, 1, 3)) .

Есть идеи?