title: Batch sequential data with tf.data
body: <p>Let's consider a toy dataset, ordered, with two features:</p>
  
  <ul>
  <li><code>value</code> (e.g. <code>1, 2, 3, 4, 5, 111, 222, 333, 444, 555</code>)</li>
  <li><code>sequence_id</code> (e.g. <code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code>)</li>
  </ul>
  
  <p>This data basically consists of two flattened sequences concatenated, <code>1, 2, 3, 4, 5</code> (sequence <code>0</code>), and <code>111, 222, 333, 444, 555</code> (sequence <code>1</code>).</p>
  
  <p>I would like to generate sequences of size <code>t</code> (say <code>3</code>) consisting of consecutive elements from the same sequence (<code>sequence_id</code>), I do not want a sequence to have elements belonging to different <code>sequence_id</code>.</p>
  
  <p>For instance, without any shuffling, I would like to get the following batches:</p>
  
  <ul>
  <li>1st batch: <code>1, 2, 3</code>,</li>
  <li>2nd batch: <code>2, 3, 4</code>,</li>
  <li>3rd batch: <code>3, 4, 5</code>,</li>
  <li>4th batch: <code>111, 222, 333</code>,</li>
  <li>5th batch: <code>222, 333, 444</code>,</li>
  <li>6th batch: <code>333, 444, 555</code>, </li>
  <li>7th batch: <code>1, 2, 3</code>,</li>
  <li>etc.</li>
  </ul>
  
  <p>I know how to generate sequence data using <code>tf.data.Dataset.window</code> or <code>tf.data.Dataset.batch</code>, but I do not know how to prevent a sequence from containing a mix of different <code>sequence_id</code> (e.g. the sequence <code>4, 5, 111</code> should not be valid as it mixes elements from sequence <code>0</code> and sequence <code>1</code>).</p>
  
  <p>Below is my failed attempt:</p>
  
  <pre><code>import tensorflow as tf
  
  data = tf.data.Dataset.from_tensor_slices(([1, 2, 3, 4, 5, 111, 222, 333, 444, 555], 
                                             [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]))\
                  .window(3, 1, drop_remainder=True)\
                  .repeat(-1)\
                  .flat_map(lambda x, y: x.batch(3))\
                  .batch(10)
  data_it = data.make_initializable_iterator()
  next_element = data_it.get_next()
  
  with tf.Session() as sess:
      sess.run(data_it.initializer)
      print(sess.run(next_element))
  </code></pre>
  
  <p>that outputs:</p>
  
  <pre><code>[[  1   2   3]   # good
   [  2   3   4]   # good
   [  3   4   5]   # good
   [  4   5 111]   # bad – mix of sequence 0 (4, 5) and sequence 1 (111)
   [  5 111 222]   # bad
   [111 222 333]   # good
   [222 333 444]   # good
   [333 444 555]   # good
   [  1   2   3]   # good
   [  2   3   4]]  # good
  </code></pre>