title: How to use Keras' predict_on_batch in tf.data.Dataset.map()?
body: <p>I would like to find a way to use Keras' <code>predict_on_batch</code> inside <code>tf.data.Dataset.map()</code> in <code>TF2.0.</code></p>
  
  <p>Let's say I have a numpy dataset</p>
  
  <pre class="lang-py prettyprint-override"><code>n_data = 10**5
  my_data    = np.random.random((n_data,10,1))
  my_targets = np.random.randint(0,2,(n_data,1))
  
  data = ({'x_input':my_data}, {'target':my_targets})
  </code></pre>
  
  <p>and a <code>tf.keras</code> model </p>
  
  <pre class="lang-py prettyprint-override"><code>x_input = Input((None,1), name = 'x_input')
  RNN     = SimpleRNN(100,  name = 'RNN')(x_input)
  dense   = Dense(1, name = 'target')(RNN)
  
  my_model = Model(inputs = [x_input], outputs = [dense])
  my_model.compile(optimizer='SGD', loss = 'binary_crossentropy')
  </code></pre>
  
  <p>I can create a batched <code>dataset</code> with</p>
  
  <pre class="lang-py prettyprint-override"><code>dataset = tf.data.Dataset.from_tensor_slices(data)
  dataset = dataset.batch(10)
  prediction_dataset = dataset.map(transform_predictions)
  </code></pre>
  
  <p>where <code>transform_predictions</code> is a user defined function that gets the predictions from <code>predict_on_batch</code></p>
  
  <pre class="lang-py prettyprint-override"><code>def transform_predictions(inputs, outputs):
      predictions = my_model.predict_on_batch(inputs)
      # predictions = do_transformations_here(predictions)
      return predictions
  </code></pre>
  
  <p>This gives an error from <code>predict_on_batch</code>:</p>
  
  <p><code>AttributeError: 'Tensor' object has no attribute 'numpy'</code></p>
  
  <p>As far as I understand, <code>predict_on_batch</code> expects a numpy array, and it is getting a tensor object from the dataset. </p>
  
  <p>It seems like one possible solution is to wrap <code>predict_on_batch</code> in a `tf.py_function, though I have not been able to get that working either. </p>
  
  <p>Does anyone know how to do this?</p>