code: def main(argv):
      if len(argv) > 1:
          raise app.UsageError("Too many command-line arguments.")
  
      FLAGS.alsologtostderr = True
  
      train_split = dataset.Split.from_string(FLAGS.train_split)
      eval_split = dataset.Split.from_string(FLAGS.eval_split)
  
      # The total batch size is the batch size accross all hosts and devices. In a
      # multi-host training setup each host will only see a batch size of
      # `total_train_batch_size / jax.host_count()`.
      total_train_batch_size = (
          FLAGS.train_device_batch_size * jax.device_count()
      )
      num_train_steps = (
          train_split.num_examples * FLAGS.train_epochs
      ) // total_train_batch_size
  
      local_device_count = jax.local_device_count()
      train_dataset = dataset.load(
          train_split,
          is_training=True,
          batch_dims=[
              local_device_count,
              FLAGS.train_device_batch_size,
          ],
          dtype=get_policy().compute_dtype,
          transpose=FLAGS.dataset_transpose,
          zeros=FLAGS.dataset_zeros,
      )
  
      # Assign mixed precision policies to modules. Note that when training in f16
      # we keep BatchNorm in  full precision. When training with bf16 you can often
      # use bf16 for BatchNorm.
      mp_policy = get_policy()
      bn_policy = get_bn_policy().with_output_dtype(
          mp_policy.compute_dtype
      )
      # NOTE: The order we call `set_policy` doesn't matter, when a method on a
      # class is called the policy for that class will be applied, or it will
      # inherit the policy from its parent module.
      hk.mixed_precision.set_policy(hk.BatchNorm, bn_policy)
      hk.mixed_precision.set_policy(hk.nets.ResNet50, mp_policy)
  
      if jax.default_backend() == "gpu":
          # TODO(tomhennigan): This could be removed if XLA:GPU's allocator changes.
          train_dataset = dataset.double_buffer(train_dataset)
  
      # For initialization we need the same random key on each device.
      rng = jax.random.PRNGKey(FLAGS.train_init_random_seed)
      rng = jnp.broadcast_to(rng, (local_device_count,) + rng.shape)
      # Initialization requires an example input.
      batch = next(train_dataset)
      train_state = jax.pmap(initial_state)(rng, batch)
  
      # Print a useful summary of the execution of our module.
      summary = hk.experimental.tabulate(train_step)(train_state, batch)
      for line in summary.split("\n"):
          logging.info(line)
  
      eval_every = FLAGS.train_eval_every
      log_every = FLAGS.train_log_every
  
      with time_activity("train"):
          for step_num in range(num_train_steps):
              # Take a single training step.
              with jax.profiler.StepTraceAnnotation(
                  "train", step_num=step_num
              ):
                  batch = next(train_dataset)
                  train_state, train_scalars = train_step(
                      train_state, batch
                  )
  
              # By default we do not evaluate during training, but you can configure
              # this with a flag.
              if (
                  eval_every > 0
                  and step_num
                  and step_num % eval_every == 0
              ):
                  with time_activity("eval during train"):
                      eval_scalars = evaluate(
                          eval_split,
                          train_state.params,
                          train_state.state,
                      )
                  logging.info(
                      "[Eval %s/%s] %s",
                      step_num,
                      num_train_steps,
                      eval_scalars,
                  )
  
              # Log progress at fixed intervals.
              if step_num and step_num % log_every == 0:
                  train_scalars = jax.tree_util.tree_map(
                      lambda v: np.mean(v).item(),
                      jax.device_get(train_scalars),
                  )
                  logging.info(
                      "[Train %s/%s] %s",
                      step_num,
                      num_train_steps,
                      train_scalars,
                  )
  
      # Once training has finished we run eval one more time to get final results.
      with time_activity("final eval"):
          eval_scalars = evaluate(
              eval_split, train_state.params, train_state.state
          )
      logging.info("[Eval FINAL]: %s", eval_scalars)
  
imports: from collections.abc import Iterable, Mapping
  import contextlib
  import functools
  import timeit
  from typing import NamedTuple
  from absl import app
  from absl import flags
  from absl import logging
  import haiku
  from examples.imagenet import dataset
  import jax
  import jax.numpy
  import jmp
  import numpy
  import optax
  import tree
label_desc: Use image node_id 12 for calling a global function with example usage: main(argv) without return types