code: def train_step(
      model: snt.Module,
      optimizer: snt.Optimizer,
      images: tf.Tensor,
      labels: tf.Tensor,
  ) -> tf.Tensor:
      """Runs a single training step of the model on the given input."""
      with tf.GradientTape() as tape:
          logits = model(images)
          loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
              labels=labels, logits=logits
          )
          loss = tf.reduce_mean(loss)
      variables = model.trainable_variables
      gradients = tape.gradient(loss, variables)
      optimizer.apply(gradients, variables)
      return loss
  
file_path: sonnet/examples/simple_mnist.py
imports: from typing import Dict
  from absl import app
  import sonnet
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: train_step(model, optimizer, images, labels) and returns: loss