title: ELMo Embedding layer with Keras
body: <p>I have been using Keras default embedding layer with word embeddings in my architecture. Architecture looks like this - </p>
  
  <pre><code>left_input = Input(shape=(max_seq_length,), dtype='int32')
  right_input = Input(shape=(max_seq_length,), dtype='int32')
  
  embedding_layer = Embedding(len(embeddings), embedding_dim, weights=[embeddings], input_length=max_seq_length,
                              trainable=False)
  
  # Since this is a siamese network, both sides share the same LSTM
  shared_lstm = LSTM(n_hidden, name="lstm")
  
  left_output = shared_lstm(encoded_left)
  right_output = shared_lstm(encoded_right)
  </code></pre>
  
  <p>I want to replace the embedding layer with ELMo embeddings. So I used a custom embedding layer - found in this repo - <a href="https://github.com/strongio/keras-elmo/blob/master/Elmo%20Keras.ipynb" rel="nofollow noreferrer">https://github.com/strongio/keras-elmo/blob/master/Elmo%20Keras.ipynb</a>. Embedding layer looks like this - </p>
  
  <pre><code>class ElmoEmbeddingLayer(Layer):
  def __init__(self, **kwargs):
      self.dimensions = 1024
      self.trainable=True
      super(ElmoEmbeddingLayer, self).__init__(**kwargs)
  
  def build(self, input_shape):
      self.elmo = hub.Module('https://tfhub.dev/google/elmo/2', trainable=self.trainable,
                             name="{}_module".format(self.name))
  
      self.trainable_weights += K.tf.trainable_variables(scope="^{}_module/.*".format(self.name))
      super(ElmoEmbeddingLayer, self).build(input_shape)
  
  def call(self, x, mask=None):
      result = self.elmo(K.squeeze(K.cast(x, tf.string), axis=1),
                    as_dict=True,
                    signature='default',
                    )['default']
      return result
  
  def compute_mask(self, inputs, mask=None):
      return K.not_equal(inputs, '--PAD--')
  
  def compute_output_shape(self, input_shape):
      return (input_shape[0], self.dimensions)
  </code></pre>
  
  <p>I changed the architecture for the new embedding layer. </p>
  
  <pre><code> # The visible layer
  left_input = Input(shape=(1,), dtype="string")
  right_input = Input(shape=(1,), dtype="string")
  
  embedding_layer = ElmoEmbeddingLayer()
  
  # Embedded version of the inputs
  encoded_left = embedding_layer(left_input)
  encoded_right = embedding_layer(right_input)
  
  # Since this is a siamese network, both sides share the same LSTM
  shared_lstm = LSTM(n_hidden, name="lstm")
  
  left_output = shared_gru(encoded_left)
  right_output = shared_gru(encoded_right)
  </code></pre>
  
  <p>But I am getting error - </p>
  
  <p>ValueError: Input 0 is incompatible with layer lstm: expected ndim=3, found ndim=2</p>
  
  <p>What am I doing wrong here? </p>