title: Tensorflow Keras LSTM with sequence_loss
body: <p>I'm building a simple LSTM model with <code>tensorflow.keras</code> and would like to use <code>sequence_loss</code> but it's not working with my current setup.</p>
  
  <p>My input is always 4 words and the target is shifted by 1 word:
  <code>[w1, w2, w3, w4] -&gt; [w2, w3, w4, w5]</code></p>
  
  <p>I know that sequence_loss takes a weight argument, so I had to write a wrapper around it to make it work in keras, but I'm having problem aligning the shape of the matrices.</p>
  
  <pre><code>import tensorflow as tf
  from tensorflow.keras import Sequential
  from tensorflow.keras.layers import Embedding, Dense, LSTM
  from tensorflow.keras.optimizers import Adam
  
  model = Sequential()
  model.add(Embedding(input_dim=vocab_size, 
                      output_dim=embed_size,
                      weights=[embedding_matrix],
                      trainable=False))
  model.add(LSTM(hidden_size, return_sequences=True, activation="tanh"))
  model.add(Dense(vocab_size, activation="softmax"))
  
  def seq_loss_fn(logits, targets):
      all_ones = tf.ones([batch_size, 4])
      return tfa.seq2seq.sequence_loss(logits, targets, all_ones)
  
  model.compile(optimizer=Adam(lr=0.001), loss=seq_loss_fn, metrics=["categorical_accuracy"])
  
  
  Model: "sequential"
  _________________________________________________________________
  Layer (type)                 Output Shape              Param #   
  =================================================================
  embedding (Embedding)        (None, 4, 100)            55400     
  _________________________________________________________________
  lstm (LSTM)                  (None, 4, 100)            80400     
  _________________________________________________________________
  dense (Dense)                (None, 4, 554)            55954     
  =================================================================
  Total params: 191,754
  Trainable params: 136,354
  Non-trainable params: 55,400
  _________________________________________________________________
  </code></pre>
  
  <p>I have tried it without <code>softmax</code> activation as well but the following is the error:</p>
  
  <pre><code>    loss=seq_loss_fn)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
      result = method(self, *args, **kwargs)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 373, in compile
      self._compile_weights_loss_and_weighted_metrics()
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
      result = method(self, *args, **kwargs)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1653, in _compile_weights_loss_and_weighted_metrics
      self.total_loss = self._prepare_total_loss(masks)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1713, in _prepare_total_loss
      per_sample_losses = loss_fn.call(y_true, y_pred)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py", line 221, in call
      return self.fn(y_true, y_pred, **self._fn_kwargs)
    File "train_rnnlm_tom2.py", line 120, in seq_loss_fn
      return tfa.seq2seq.sequence_loss(logits, targets, all_ones)
    File "/Users/tom/anaconda3/lib/python3.7/site-packages/tensorflow_addons/seq2seq/loss.py", line 90, in sequence_loss
      raise ValueError("Logits must be a "
  ValueError: Logits must be a [batch_size x sequence_length x logits] tensor
  
  </code></pre>