title: Parallel threads with TensorFlow Dataset API and flat_map
body: <p>I'm changing my TensorFlow code from the old queue interface to the new <a href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset" rel="noreferrer">Dataset API</a>. With the old interface I could specify the <code>num_threads</code> argument to the <code>tf.train.shuffle_batch</code> queue. However, the only way to control the amount of threads in the Dataset API seems to be in the <code>map</code> function using the <code>num_parallel_calls</code> argument. However, I'm using the <code>flat_map</code> function instead, which doesn't have such an argument.</p>
  
  <p><strong>Question</strong>: Is there a way to control the number of threads/processes for the <code>flat_map</code> function? Or is there are way to use <code>map</code> in combination with <code>flat_map</code> and still specify the number of parallel calls?</p>
  
  <p>Note that it is of crucial importance to run multiple threads in parallel, as I intend to run heavy pre-processing on the CPU before data enters the queue.</p>
  
  <p>There are two (<a href="https://github.com/tensorflow/tensorflow/issues/7951#issuecomment-305796971" rel="noreferrer">here</a> and <a href="https://github.com/tensorflow/tensorflow/issues/7951#issuecomment-326098305" rel="noreferrer">here</a>) related posts on GitHub, but I don't think they answer this question.</p>
  
  <p>Here is a minimal code example of my use-case for illustration:</p>
  
  <pre><code>with tf.Graph().as_default():
      data = tf.ones(shape=(10, 512), dtype=tf.float32, name="data")
      input_tensors = (data,)
  
      def pre_processing_func(data_):
          # normally I would do data-augmentation here
          results = (tf.expand_dims(data_, axis=0),)
          return tf.data.Dataset.from_tensor_slices(results)
  
      dataset_source = tf.data.Dataset.from_tensor_slices(input_tensors)
      dataset = dataset_source.flat_map(pre_processing_func)
      # do something with 'dataset'
  </code></pre>