import tensorflow as tf

from official.mnist import mnist_eager
from official.utils.misc import keras_utils


def device():
  return "/device:GPU:0" if tfe.num_gpus() else "/device:CPU:0"


def data_format():
  return "channels_first" if tfe.num_gpus() else "channels_last"


def random_dataset():
  batch_size = 64
  images = tf.random_normal([batch_size, 784])
  labels = tf.random_uniform([batch_size], minval=0, maxval=10, dtype=tf.int32)
  return tf.data.Dataset.from_tensors((images, labels))


def train(defun=False):
  model = mnist.create_model(data_format())
  if defun:
    model.call = tfe.defun(model.call)
  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
  dataset = random_dataset()
  with tf.device(device()):
    mnist_eager.train(model, optimizer, dataset,
                      step_counter=tf.train.get_or_create_global_step())


