Keras - Как применить пользовательскую функцию к партии на последовательной модели?

Question

Я пытаюсь изменить модель, которую я видел на бумаге, фактически меняя входной слой с 1DConv на LSTM. Проблема заключается в том, что после 1DConv они применяют функцию под названием Outer Concatenation. Код выглядит следующим образом.

def OuterConcatenate(x):
    seqLen = int(x.shape[1])
    input2 = x
    mesh = np.mgrid[0:seqLen, 0:seqLen]
    out=input2[mesh]
    output = np.concatenate((out[0],out[1]),axis=-1)
    return output

Фиктивный код, с которым я играл, чтобы проверить это:

import numpy as np
from keras.utils import to_categorical as tc
x = np.random.uniform(0, 1, (2, 10, 46)).astype(np.float32)
vec = np.ones((5))
seq = np.asarray([i for i in np.random.randint(20,size=5)])
seq = tc(seq,20) #tc is keras.utils.to_categorical
seq2 = seq[np.newaxis,...]
#
cell_count = 6
from keras.models import Sequential
from keras.layers import TimeDistributed, Bidirectional, LSTM, Input,Dense,Lambda
model = Sequential()
model.add(Bidirectional(LSTM(cell_count,input_shape=(5,20),return_sequences=True)))
model.add(TimeDistributed(Dense(2,activation='softmax')))
model.add(Lambda(lambda x: OuterConcatenate(x)))
b = model.predict(seq2)

Проблема в том, что я получаю следующую ошибку:

`ValueError: Shape must be rank 1 but is rank 4 for 'lambda_1/stri`ded_slice' (op:StridedSlice') with input shapes: [?,5,2], [1,2,5,5], [1,2,5,5], [1].b.shape

Я хотел бы применить эту функцию OuterConcatenate ко всем x в X (пакет формы (?, 5,2)). Как я могу это сделать?

1. Написание пользовательского слоя?
2. Использование lambda, но изменение функции OuterConcatenate?
3. Реализация функции с помощью методов tf или keras.beckend ?
4. Что-нибудь еще?

Заранее спасибо.