Can I use a Tensor as a list index?

Question

I have this Custom Keras Layer that chooses between elements of a list, like a Dense layer, and I want it to return the element of the list it predicted directly. The list is a list of Keras.layers.Layer. I have this piece of code:

def call(self, inputs, context):   
    pred = tf.argmax(tf.matmul(context, self.kernel))
    return self.layers[pred](inputs)

It throws an error: TypeError: list indices must be integers or slices, not Tensor, which is understandable, but I can't find a way of making it work. The "pred" Tensor doesn't have a .numpy property, though I'm running the program eagerly, since this happens when the layer is being built.

I understand there may be no solutions, if so, submit ideas on how I could code this layer in another way.

Answer 1

There is a bigger problem.

This layer will not work, because you cannot get derivatives of argmax, the kernel will be impossible to train. And you will get an error message like "An operation has None for gradient"

As a workaround, I'd suggest you to:

• 1: calculate all layers (hope they have the same shape?)
• 2: stack their results in the second dimension: tf.stack([listf_of_outputs],axis=1)
• 2: take a softmax of the result of matmul
• 3: reshape the result of softmax to the same number of dimensions of the stacked result above: shape (-1, number_of_layers, _other_dims_if_exist, 1)
• 4: multiply (elementwise *) the stacked results by the reshaped softmax and sum the axis 1.