如何在自定义Keras图层中查找输入批处理的特定行

Question

我想编写一个自定义Keras层，该层接受多个输入并将一个函数应用于输入批处理的不同子集。

应使用第二个输入变量找到应一起处理的特定样品。 以下代码应演示我想做什么：

from keras import layers
import keras.backend as K

class MyCustomLayer(layers.Layer):
    def call(self, inputs):
        t1, t2 = inputs
        # Now I switch to pseudo code, because I do not know how to do this in Keras
        for element in t2.unique():
              rows = where(t2 == element)
              output[rows] = function(t1[rows])
        return output

是否可以在Keras / Tensorflow中实现它？