title: How to use the output of attention wrapper applied over LSTM as an input to the TimeDistributed layer, Keras?
body: <p>I have been trying to implement an attention wrapper over the output of the LSTM model shown in this <a href="https://machinelearningmastery.com/timedistributed-layer-for-long-short-term-memory-networks-in-python/" rel="nofollow noreferrer">machinelearningmastery</a> tutorial:</p>
  
  <pre><code>from numpy import array
  from keras.models import Sequential
  from keras.layers import Dense
  from keras.layers import TimeDistributed
  from keras.layers import LSTM
  # prepare sequence
  length = 5
  seq = array([i/float(length) for i in range(length)])
  X = seq.reshape(1, length, 1)
  y = seq.reshape(1, length, 1)
  # define LSTM configuration
  n_neurons = length
  n_batch = 1
  n_epoch = 1000
  # create LSTM
  model = Sequential()
  model.add(LSTM(n_neurons, input_shape=(length, 1), return_sequences=True))
  model.add(TimeDistributed(Dense(1)))
  model.compile(loss='mean_squared_error', optimizer='adam')
  print(model.summary())
  # train LSTM
  model.fit(X, y, epochs=n_epoch, batch_size=n_batch, verbose=2)
  # evaluate
  result = model.predict(X, batch_size=n_batch, verbose=0)
  for value in result[0,:,0]:
      print('%.1f' % value)
  </code></pre>
  
  <p>The output<code>(samples, steps, features)</code> of the LSTM with <code>return_sequences = true</code> is received by the attention wrapper (which outputs the shape (samples, features)). Here is my modified code:</p>
  
  <pre><code>model = Sequential()
  model.add(LSTM(n_neurons, input_shape=(length, 1), return_sequences=True))
  model.add(Attention())
  model.add(TimeDistributed(Dense(1)))
  </code></pre>
  
  <p>I have been using the attention wrapper as described <a href="https://gist.github.com/cbaziotis/6428df359af27d58078ca5ed9792bd6d" rel="nofollow noreferrer">here</a>:</p>
  
  <pre><code>    def dot_product(x, kernel):
  
  
          if K.backend() == 'tensorflow':
              # todo: check that this is correct
              return K.squeeze(K.dot(x, K.expand_dims(kernel)), axis=-1)
          else:
              return K.dot(x, kernel)
  
  
  class Attention(Layer):
      def __init__(self,
                   W_regularizer=None, b_regularizer=None,
                   W_constraint=None, b_constraint=None,
                   bias=True, **kwargs):
  
          self.supports_masking = True
          self.init = initializers.get('glorot_uniform')
  
          self.W_regularizer = regularizers.get(W_regularizer)
          self.b_regularizer = regularizers.get(b_regularizer)
  
          self.W_constraint = constraints.get(W_constraint)
          self.b_constraint = constraints.get(b_constraint)
  
          self.bias = bias
          super(Attention, self).__init__(**kwargs)
  
      def build(self, input_shape):
          assert len(input_shape) == 3
  
          self.W = self.add_weight((input_shape[-1],),
                                   initializer=self.init,
                                   name='{}_W'.format(self.name),
                                   regularizer=self.W_regularizer,
                                   constraint=self.W_constraint)
          if self.bias:
              self.b = self.add_weight((input_shape[1],),
                                       initializer='zero',
                                       name='{}_b'.format(self.name),
                                       regularizer=self.b_regularizer,
                                       constraint=self.b_constraint)
          else:
              self.b = None
  
          self.built = True
  
      def compute_mask(self, input, input_mask=None):
          # do not pass the mask to the next layers
          return None
  
      def call(self, x, mask=None):
          eij = dot_product(x, self.W)
  
          if self.bias:
              eij += self.b
  
          eij = K.tanh(eij)
  
          a = K.exp(eij)
  
          # apply mask after the exp. will be re-normalized next
          if mask is not None:
              # Cast the mask to floatX to avoid float64 upcasting in theano
              a *= K.cast(mask, K.floatx())
  
          # in some cases especially in the early stages of training the sum may be almost zero
          # and this results in NaN's. A workaround is to add a very small positive number ε to the sum.
          # a /= K.cast(K.sum(a, axis=1, keepdims=True), K.floatx())
          a /= K.cast(K.sum(a, axis=1, keepdims=True) + K.epsilon(), K.floatx())
  
          a = K.expand_dims(a)
          weighted_input = x * a
          return K.sum(weighted_input, axis=1)
  
      def get_output_shape_for(self, input_shape):
          return input_shape[0], input_shape[-1]
  </code></pre>
  
  <p>However, the error I get is: </p>
  
  <blockquote>
    <p>ValueError: Input 0 is incompatible with layer time_distributed_1: expected ndim=3, found ndim=2</p>
  </blockquote>
  
  <p>Is there any way out to readjust the shape of the attention wrapper's output so as to use the TimeDistributed layer? </p>