使用Lambdas在Keras中定制Dot产品

考虑我的输入数据输入a和输入b

np.array([1,2,3,4,5,6,7,8,9,10,11,12]).reshape(2,3,2)
array([[[ 1,  2],
        [ 3,  4],
        [ 5,  6]],

       [[ 7,  8],
        [ 9, 10],
        [11, 12]]])

及

我想要达到的目标

np.einsum('kmn, mk -> mn', input_a, input_b)
array([[ 15,  18],
       [ 45,  52],
       [ 91, 102]])

如何将其转换为keras中的lambda层

到目前为止我都试过了

def tensor_product(x):
    input_a = x[0]
    input_b = x[1]
    y = np.einsum('kmn, mk -> mn', input_a, input_b)
    return y

dim_a = Input(shape=(2,))
dim_b = Input(shape=(2,2,))
layer_3 = Lambda(tensor_product, output_shape=(2,))([dim_a, dim_b])
model = Model(inputs=[input_a, input_b], outputs=layer_3)

多谢各位

from keras.layers import *

def tensor_product(x):
    a = x[0]
    b = x[1]
    b = K.permute_dimensions(b, (1, 0, 2))
    y = K.batch_dot(a, b, axes=1)
    return y

a = Input(shape=(2,))
b = Input(shape=(2,2,))
c = Lambda(tensor_product, output_shape=(2,))([a, b])
model = Model([a, b], c)

from keras.layers import *

def tensor_product(x):
    a = x[0]
    b = x[1]
    b = K.permute_dimensions(b, (1, 0, 2))
    y = K.batch_dot(a, b, axes=1)
    return y

a = Input(shape=(2,))
b = Input(shape=(2,2,))
c = Lambda(tensor_product, output_shape=(2,))([a, b])
model = Model([a, b], c)