How to slice a Keras tensor using the values from another tensor in Tensorflow 2.0?

Question

I have a Keras tensor of feature maps of size (None, 256, 256, 128) obtained after convolution on input image of size (None, 256, 256, 3). The another input tensor (index tensor) is of size (None, 128) that contains integers from 0 to 127. I want to select a subset of feature maps (channels) from the first Keras tensor using the indices from the index tensor. I am using the following code:

x = Input(shape=(256,256,3), name='input_image')
ind = Input(shape=(128,), name='index')
c1 = Conv2D(128, 3, padding = 'same', activation = 'relu')(x)
c1[ind[:,0]]

I get the following error after the implementation:

Exception encountered when calling layer "tf.__operators__.getitem_6" (type SlicingOpLambda).

Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar
tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor 'Placeholder_1:0' shape=(None,) 
dtype=float32>

Call arguments received:
• tensor=tf.Tensor(shape=(None, 256, 256, 128), dtype=float32)
• slice_spec=tf.Tensor(shape=(None,), dtype=float32)
• var=None

How can this be achieved in Keras/Tensorflow 2.0?

For reference, I am implementing the spatial attention module from the following paper: Hybrid Attention Module

Answer 1

Here ind[:,0] give out shape (None,), it's not a slice indicator but just a Placeholder since no real data feeding. And different from numpy slice API,c1[ind[:,0]] cannot map the structure in tensorflow, such as?

import numpy as np
import tensorflow as tf 

a = np.array([[1,2,3],[11,22,33]])
b = np.array([[[1,2,3,4,5,6],[11,22,33,44,55,66]],[[111,222,333,444,555,666],[1111,2222,3333,4444,5555,6666]]])

c = b[0]
print(c)
# [[ 1  2  3  4  5  6]
#  [11 22 33 44 55 66]]

c = b[0,1]
print(c)
# [11 22 33 44 55 66]

c = b[[[[[0,1]]]]] # b[tuple([[[[0,1]]]])]
print(c)
# [[[[[   1    2    3    4    5    6]
#     [  11   22   33   44   55   66]]
# 
#    [[ 111  222  333  444  555  666]
#     [1111 2222 3333 4444 5555 6666]]]]]

a = tf.constant([[1,2,3],[11,22,33]])
b = tf.constant([[[1,2,3,4,5,6],[11,22,33,44,55,66]],[[111,222,333,444,555,666],[1111,2222,3333,4444,5555,6666]]])

c = b[0]
print(c)
# tf.Tensor(
# [[ 1  2  3  4  5  6]
#  [11 22 33 44 55 66]], shape=(2, 6), dtype=int32)

c = b[0,1]
print(c)
# tf.Tensor([11 22 33 44 55 66], shape=(6,), dtype=int32)

c = b[[0,1]] # differ with numpy
print(c)
# tf.Tensor([11 22 33 44 55 66], shape=(6,), dtype=int32)

# c = b[[[0],[0]]] # TypeError

So here I have to infer your meaning on c1[ind[:,0]]. Do you mean the index will be different across batch dimension and so the ind should maintain its first dimension? It will be achieved easily in custom model coding, but may be hard to realize in sequential model coding. Here are the sample codes in custom model coding:

import tensorflow as tf 
import functools
class my_layer(tf.keras.layers.Layer):
    def __init__(self) -> None:
        super().__init__()
        self.c1 = tf.keras.layers.Conv2D(128, 3, padding = 'same', activation = 'relu')
    def _index_map_func(self,index,data):
        return data[...,index]
    def call(self,x,ind):
        y = self.c1(x)
        feature_index_map_func = functools.partial(self._index_map_func,data=y)
        return tf.stack(list(map(feature_index_map_func,ind[...,0])),axis=-1)
batch_size = 5
x = tf.random.normal(shape=(batch_size,256,256,3), name='input_image')
ind = tf.random.uniform(shape=[batch_size,128],maxval=128,dtype=tf.int32)
new_c1 = my_layer()

# and work in custom training loop
with tf.GradientTape() as tape:
    y = new_c1(x,ind) 
    print(y.shape)
    # (5, 256, 256, 5)
gd = tape.gradient(y,new_c1.trainable_variables)
print(len(gd))
# 2