snippet.py

# See https://github.com/tensorflow/docs/blob/master/site/en/r2/guide/autograph.ipynb

def train_one_step(model, optimizer, x, y):
  with tf.GradientTape() as tape:
    logits = model(x)
    loss = compute_loss(y, logits)

  grads = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))

  compute_accuracy(y, logits)
  return loss

def train(model, optimizer):
  train_ds = mnist_dataset()
  step = 0
  loss = 0.0
  for x, y in train_ds:
    step += 1
    loss = train_one_step(model, optimizer, x, y)
    if tf.equal(step % 10, 0):
      tf.print('Step', step, ': loss',
                loss, '; accuracy', compute_accuracy.result())
  return step, loss, accuracy