batch_size = 64
vocab_size = 42
embedding_dim = 256
rnn_units = 1024

model = tf.keras.Sequential([
tf.keras.layers.Embedding(vocab_size, embedding_dim,
                          batch_input_shape=[batch_size, None]),
rnn(rnn_units,
    return_sequences=True,
    recurrent_initializer='glorot_uniform',
    stateful=True),
tf.keras.layers.Dense(vocab_size)
])

[[[9  1  0 ...  0  0  0]
   [31  0  0 ...  0  0  0]
   [28  1  0 ...  1  0  0]
   ...
   [ 9  1  0 ...  0  0  0]
   [15  1  0 ...  0  1  0]
   [25  0  0 ...  0  0  0]]

  [[31  0  0 ...  0  0  0]
   [6  0  0 ...  0  0  0]
   [22  1  1 ...  0  0  0]
   ...
   [15  1  0 ...  0  1  0]
   [28  1  0 ...  1  0  0]
   [2  0  0 ...  0  0  0]]

  [[31  0  0 ...  0  0  0]
   [11  0  0 ...  0  0  0]
   [9  1  0 ...  0  0  0]
   ...
   [6  0  0 ...  0  0  0]
   [36  1  0 ...  0  0  0]
   [17  0  0 ...  0  0  0]]
   ...
  [[27  0  1 ...  0  0  0]
   [27  0  1 ...  0  0  0]
   [2  0  0 ...  0  0  0]
   ...
   [9  1  0 ...  0  0  0]
   [12  0  0 ...  0  0  0]
   [36  1  0 ...  0  0  0]]

  [[37  1  0 ...  0  0  0]
   [34  0  0 ...  0  0  0]
   [16  0  0 ...  0  0  1]
   ...
   [1  0  0 ...  0  0  0]
   [23  1  0 ...  0  0  0]
   [22  1  1 ...  0  0  0]]

  [[3  0  0 ...  0  0  0]
   [10  1  0 ...  0  0  0]
   [12  0  0 ...  0  0  0]
   ...
   [27  0  1 ...  0  0  0]
   [12  0  0 ...  0  0  0]
   [31  0  0 ...  0  0  0]]], shape=(64, 100, 9), dtype=int32)