Как индексировать список тензоров в определении модели Keras

Question

Предположим, у нас есть следующее определение модели keras

def build_model(ls):
    x = Input(shape=(1,), name='input_index')
    gs = [K.constant(x) for x in ls]

    idx = Lambda(lambda x: K.cast(x, 'int32'))(x)
    y = Lambda(lambda x: gs[x[0, 0]])(idx)

    model = Model(inputs=x, outputs=y)
    return model

Здесь мы передаем список ndarray:ls для инициализации модели. Внутри определения модели gs преобразует ls в список постоянного тензора. Входные данные x дают индекс i, и модель должна вернуть соответствующий тензор gs[i]. Тем не менее, приведенная выше модель верна.

Можно ли как-нибудь сделать это в керасе?