title: reduce a special 4D tensor to 3D tensor in keras
body: <p>I have a 4D tensor <strong>x</strong> like: </p>
  
  <pre><code>(bs1, bs2, sent_len1, sentlen2)  # bs1 and bs2 are unknown and bs1 &gt;= bs2.
  </code></pre>
  
  <p>and an indices tensor <strong>ind</strong> like:</p>
  
  <pre><code>(bs1, 1)  
  </code></pre>
  
  <p>If I transform <strong>ind</strong> to a numpy array <strong>np_ind</strong> and do an operation:</p>
  
  <pre><code>cate_ind = to_categorical(np_ind, num_classes=None)
  </code></pre>
  
  <p>the shape of cate_ind is:</p>
  
  <pre><code>(bs1, bs2)
  </code></pre>
  
  <p>I want to use tensor <strong>ind</strong> to reduce the dimension of <strong>x</strong> from</p>
  
  <pre><code>(bs1, bs2, sent_len1, sentlen2) to (bs1, sent_len1, sentlen2)
  </code></pre>
  
  <p>the code example is:</p>
  
  <pre><code>from keras.layers import *
  from keras.models import Model
  from keras.utils import to_categorical
  from keras.layers.merge import *
  
  
  class CustomLayer(Layer):
      def __init__(self, **kwargs):
          super(CustomLayer, self).__init__(**kwargs)
  
      def build(self, input_shape):
          assert len(input_shape) == 3
          super(CustomLayer, self).build(input_shape)
  
      # TODO:
      def dimension_reduction(self, tensor, indices):
          pass
  
      def call(self, x, mask=None):
          tensor1, tensor2, ind = x[0], x[1], x[2]  # (bs1, sent_len1, wd) (bs2, sent_len2, wd) (bs1, 1)
          tensor2 = K.permute_dimensions(tensor2, [0, 2, 1])
          align = K.dot(tensor1, tensor2)
          align = K.permute_dimensions(align, [0, 2, 1, 3])  # (bs1, bs2, sent_len1, sent_len2)
          align = self.dimension_reduction(align, ind)  # (bs1, sen_len1, sent_len2)
          return align
  
      def compute_output_shape(self, input_shape):
          t1_shape, t2_shape = input_shape[0], input_shape[1]
          return t1_shape[0], t1_shape[1], t2_shape[1]
  
  
  # model example
  t1 = Input(shape=(7, 3))
  t2 = Input(shape=(6, 3))
  t3 = Input(shape=(1,))
  output = CustomLayer()([t1, t2, t3])
  model = Model([t1, t2, t3], output)
  
  # data example
  tensor1 = np.random.rand(10, 7, 3)  # (bs1, sent_len1, wd)
  tensor2 = np.random.rand(4, 6, 3)  # (bs2, sent_len2, wd)
  indices = np.array([0, 1, 3, 2, 0, 1, 2, 2, 3, 1])  # (bs1, 1)
  # indices = to_categorical(indices)  # (bs1, bs2)
  print(model.predict([tensor1, tensor2, indices]))  # (bs1, sen_len1, sent_len2)
  </code></pre>