title: How do I predict on more than one batch from a Tensorflow Dataset, using .predict_on_batch?
body: <p>As the question says, I can only predict from my model with model.predict_on_batch(). Keras tries to concatenate everything together if I use model.predict() and that doesn't work.<br>
  For my application (a sequence to sequence model) it is faster to do grouping on the fly. But even if I had done it in Pandas and then only used Dataset for the padded batch, .predict() still shouldn't work?</p>
  
  <p>If I can get predict_on_batch to work then that's what works. But I can only predict on the first batch of the Dataset. How do I get predictions for the rest? I can't loop over the Dataset, I can't consume it... </p>
  
  <p>Here's a smaller code example. The group is the same as the labels but in the real world they are obviously two different things. There are 3 classes, maximum of 2 values in a sequence, 2 row of data per batch. There's a lot of comments and I nicked parts of the windowing from somewhere on StackOverflow. I hope it is fairly legible to most. </p>
  
  <p>If you have any other suggestions on how to improve the code, please comment. But no, that's not what my model looks like at all. So suggestions for that part probably aren't helpful. </p>
  
  <p>EDIT: Tensorflow version 2.1.0</p>
  
  <pre class="lang-py prettyprint-override"><code>import tensorflow as tf
  from tensorflow.keras.models import Model
  from tensorflow.keras.layers import Bidirectional, Masking, Input, Dense, GRU
  import random
  import numpy as np
  
  random.seed(100)
  # input data
  feature = list(range(3, 14))
  # shuffle data
  random.shuffle(feature)
  # make label from feature data, +1 because we are padding with zero
  label = [feat // 5 +1 for feat in feature]
  group = label[:]
  # random.shuffle(group)
  max_group = 2
  batch_size = 2
  
  print('Data:')
  print(*zip(group, feature, label), sep='\n')
  
  # make dataset from data arrays
  ds = tf.data.Dataset.zip((tf.data.Dataset.from_tensor_slices({'group': group, 'feature': feature}), 
                            tf.data.Dataset.from_tensor_slices({'label': label})))
  # group by window
  ds = ds.apply(tf.data.experimental.group_by_window(
      # use feature as key (you may have to use tf.reshape(x['group'], []) instead of tf.cast)
      key_func=lambda x, y: tf.cast(x['group'], tf.int64),
      # convert each window to a batch
      reduce_func=lambda _, window: window.batch(max_group),
      # use batch size as window size
      window_size=max_group))
  
  # shuffle at most 100k rows, but commented out because we don't want to predict on shuffled data
  # ds = ds.shuffle(int(1e5)) 
  ds = ds.padded_batch(batch_size,
                       padded_shapes=({s: (None,) for s in ['group', 'feature']}, 
                                      {s: (None,) for s in ['label']}))
  # show dataset contents
  print('Result:')
  for element in ds:
      print(element)
  
  # Keras matches the name in the input to the tensor names in the first part of ds
  inp = Input(shape=(None,), name='feature')
  # RNNs require an additional rank, even if it is a degenerate dimension
  duck = tf.expand_dims(inp, axis=-1)
  rnn = GRU(32, return_sequences=True)(duck)
  # again Keras matches names
  out = Dense(max(label)+1, activation='softmax', name='label')(rnn)
  model = Model(inputs=inp, outputs=out)
  model.summary()
  model.compile(loss='sparse_categorical_crossentropy', metrics=['accuracy'])
  
  model.fit(ds, epochs=3)
  
  model.predict_on_batch(ds)
  </code></pre>