A customized Lasagne layer whose output is a matrix of the first input layer indexed according to the second input layer

[HATEM EL-AZAB]

Hi guys,

I need a lasagne layer, which takes two layers as inputs, and return a tensor of the first layer indexed using second layer vector, which is the result after performing argmax operation on it. Here's an example: the first layer (layer0) shape is (64,9,19,21), while the second layer (layer1) shape is (64,8). so the operation i want to perform is: Argmax_layer: performing an argmax operation on dimension(1) of layer1, which will return a 1-d array. then take argmax result from the previouse step to index the 2nd dimension of layer0, i.e., layer0[np.arange(64),Argmax_layer,:,:], so the output shape is (64,19,21).

I need a lasagne layer, which takes two layers as inputs, and return a tensor of the first layer indexed using second layer vector, which is the result after performing argmax operation on it. Here's an example: the first layer (layer0) shape is (64,9,19,21), while the second layer (layer1) shape is (64,8). so the operation i want to perform is: Argmax_layer: performing an argmax operation on dimension(1) of layer1, which will return a 1-d array. then take argmax result from the previouse step to index the 2nd dimension of layer0, i.e., layer0[np.arange(64),Argmax_layer,:,:], so the output shape is (64,19,21).

Here's the implementation of the custom layer:

Here's the NN

Here's the NN summary:

when I try to build the NN i recieve the following error:

What is wrong with this implementation?