import tensorflow as tf

        session.run(output)

      with self.assertRaisesRegexp(tf.errors.InvalidArgumentError,
                                   'Output shape must have a batch dimension'):
        coord.join()

  def test_output_must_have_same_batch_dimension_size_as_input(self):
    with self.test_session() as session:
      @dynamic_batching.batch_fn
      def f(_):
        return tf.constant([1, 2, 3, 4])

      output = f(tf.constant([1]))

      coord = tf.train.Coordinator()
      tf.train.start_queue_runners(coord=coord)

      with self.assertRaises(tf.errors.CancelledError):
        session.run(output)

      with self.assertRaisesRegexp(
          tf.errors.InvalidArgumentError,
          'Output shape must have the same batch dimension as the input batch '
          'size. Expected: 1 Observed: 4'):
        coord.join()

  def test_get_inputs_cancelled(self):
    with tf.Graph().as_default():

      @dynamic_batching.batch_fn
