import tensorflow as tf



# Benchmark related
def device_and_data_format():
  return ("/gpu:0",
          "channels_first") if tf.test.is_gpu_available() else ("/cpu:0",
                                                                "channels_last")


def random_batch(batch_size, config):
  shape = (batch_size,) + config.input_shape
  images = tf.random_uniform(shape)
  labels = tf.random_uniform(
      [batch_size], minval=0, maxval=config.n_classes, dtype=tf.int32)

  return images, labels


class MockIterator(object):

  def __init__(self, tensors):
    self._tensors = [tf.identity(x) for x in tensors]
