import tensorflow as tf


  def test_op_shape(self):
    with self.test_session():
      batcher = dynamic_batching._Batcher(minimum_batch_size=1,
                                          maximum_batch_size=1,
                                          timeout_ms=None)

      _, computation_id = batcher.get_inputs([tf.int32])

      self.assertEqual([], computation_id.shape)


class DynamicBatchingBenchmarks(tf.test.Benchmark):

  def benchmark_batching_small(self):
    with tf.Session() as session:
      @dynamic_batching.batch_fn
      def f(a, b):
        return a + b

      outputs = []
      for _ in xrange(1000):
        outputs.append(f(tf.ones([1, 10]), tf.ones([1, 10])))
      op_to_benchmark = tf.group(*outputs)

      tf.train.start_queue_runners()

      self.run_op_benchmark(
          name='batching_many_small',
          sess=session,
