code: def main(_):
      if _MICROBATCHES.value:
          raise NotImplementedError(
              "Microbatches < batch size not currently supported"
          )
  
      (
          train_images,
          train_labels,
          test_images,
          test_labels,
      ) = datasets.mnist()
      num_train = train_images.shape[0]
      num_complete_batches, leftover = divmod(
          num_train, _BATCH_SIZE.value
      )
      num_batches = num_complete_batches + bool(leftover)
      key = random.PRNGKey(_SEED.value)
  
      def data_stream():
          rng = npr.RandomState(_SEED.value)
          while True:
              perm = rng.permutation(num_train)
              for i in range(num_batches):
                  batch_idx = perm[
                      i
                      * _BATCH_SIZE.value : (i + 1)
                      * _BATCH_SIZE.value
                  ]
                  yield train_images[batch_idx], train_labels[batch_idx]
  
      batches = data_stream()
  
      opt_init, opt_update, get_params = optimizers.sgd(
          _LEARNING_RATE.value
      )
  
      @jit
      def update(_, i, opt_state, batch):
          params = get_params(opt_state)
          return opt_update(i, grad(loss)(params, batch), opt_state)
  
      @jit
      def private_update(rng, i, opt_state, batch):
          params = get_params(opt_state)
          rng = random.fold_in(
              rng, i
          )  # get new key for new random numbers
          return opt_update(
              i,
              private_grad(
                  params,
                  batch,
                  rng,
                  _L2_NORM_CLIP.value,
                  _NOISE_MULTIPLIER.value,
                  _BATCH_SIZE.value,
              ),
              opt_state,
          )
  
      _, init_params = init_random_params(key, (-1, 28, 28, 1))
      opt_state = opt_init(init_params)
      itercount = itertools.count()
  
      steps_per_epoch = 60000 // _BATCH_SIZE.value
      print("\nStarting training...")
      for epoch in range(1, _EPOCHS.value + 1):
          start_time = time.time()
          for _ in range(num_batches):
              if _DPSGD.value:
                  opt_state = private_update(
                      key,
                      next(itercount),
                      opt_state,
                      shape_as_image(*next(batches), dummy_dim=True),
                  )
              else:
                  opt_state = update(
                      key,
                      next(itercount),
                      opt_state,
                      shape_as_image(*next(batches)),
                  )
          epoch_time = time.time() - start_time
          print(f"Epoch {epoch} in {epoch_time:0.2f} sec")
  
          # evaluate test accuracy
          params = get_params(opt_state)
          test_acc = accuracy(
              params, shape_as_image(test_images, test_labels)
          )
          test_loss = loss(
              params, shape_as_image(test_images, test_labels)
          )
          print(
              "Test set loss, accuracy (%): ({:.2f}, {:.2f})".format(
                  test_loss, 100 * test_acc
              )
          )
  
          # determine privacy loss so far
          if _DPSGD.value:
              delta = 1e-5
              num_examples = 60000
              eps = compute_epsilon(
                  epoch * steps_per_epoch, num_examples, delta
              )
              print(
                  f"For delta={delta:.0e}, the current epsilon is: {eps:.2f}"
              )
          else:
              print("Trained with vanilla non-private SGD optimizer")
  
imports: import itertools
  import time
  import warnings
  from absl import app
  from absl import flags
  from jax import grad
  from jax import jit
  from jax import random
  from jax import vmap
  from jax.example_libraries import optimizers
  from jax.example_libraries import stax
  from jax.tree_util import tree_flatten, tree_unflatten
  import jax.numpy
  from examples import datasets
  import numpy.random
  from dp_accounting import dp_event
  from dp_accounting import rdp
label_desc: Use image node_id 7 for calling a global function with example usage: main(_) and returns: opt_update, opt_update