Как представить данные для LSTM?

У меня есть данные о последовательности, которые говорят мне, какой цвет наблюдался у нескольких субъектов в разные моменты времени. Например:

Я хочу получить прогнозы для наиболее вероятного цвета для следующих 3 временных шагов, а также вероятность появления этого цвета. Например, для ID A я хотел бы знать следующие 3 элемента (время, цвет) в последовательности, а также вероятность появления цвета.

Я понимаю, что LSTM часто используются для прогнозирования этого типа последовательных данных, и что я буду вводить трехмерный массив, например

input =[ 
[[1,1], [2,2], [5,2]], #blue at t=1, red at t=2, red at t=5 for ID A
[[0,0], [3,1], [6,3]], #nothing for first entry, blue at t=3, green at t=6 for ID B
[[0,0], [1,2], [3,4]]
]
  

после сопоставления цветов с числами (Синий->1, Красный->2, Зеленый->3, Оранжевый->4и т.д.). Насколько я понимаю, по умолчанию LSTM просто предсказывает следующий элемент в каждой последовательности, например

output = [
[[7, 2]], #next item is most likely red at t=7
[[9, 3]], # next item is most likely red at t=9
[[6, 2]] 
]

Можно ли изменить вывод моего LSTM, чтобы вместо того, чтобы просто предсказывать время и цвет следующего появления, я мог получить следующие 3 раза, цвета И вероятности появления цвета? Например, такой вывод, как

output = [
[[7, 2, 0.93], [8,2, 0.79], [10,4, 0.67]], 
[[9, 2, 0.88], [11,3, 0.70], [14,3, 0.43]], 
...
]

Я пробовал искать в документации Sequential для Keras, но не уверен, что нашел что-нибудь.

Кроме того, я вижу, что для model.fit() обычно используются TrainX и TrainY, но я также не уверен, какой здесь будет мой TrainY?