Keras Lambda层中具有3-D张量的索引3-D张量

Question

我试图在Keras Lambda层的Tensorflow中使用2D张量索引3D张量。我将以下函数用于Lambda层：

def testFunc(x):
   indexMatrix = K.cast(x[0], dtype = "int32")
   valueMatrix = x[1]
   idndex = tf.meshgrid(*[tf.range(s) for s in indexMatrix.get_shape()],indexing='ij')
   idx = tf.stack(index + [indexMatrix], axis=-1)
   output = tf.gather_nd(valueMatrix, idx)    
   return  output

并且我使用Lambda层，例如：

layer = Lambda(testFunc, output_shape)([model.layers[0].output,model.layers[-1].output])

运行代码时，在图形构造期间出现错误，因为indexMatrix和valueMatrix的尺寸未知，并且tf.range（）需要已知的尺寸。错误是

"Cannot convert an unknown Dimension to a Tensor: ?" 

有人知道我可以解决这个问题吗？谢谢