code: def main(_):
      flags.FLAGS.alsologtostderr = True
  
      # Make training dataset.
      train_data = dataset.load(
          tfds.Split.TRAIN,
          batch_size=TRAIN_BATCH_SIZE.value,
          sequence_length=SEQUENCE_LENGTH.value,
      )
  
      # Make evaluation dataset(s).
      eval_data = {  # pylint: disable=g-complex-comprehension
          split: dataset.load(
              split,
              batch_size=EVAL_BATCH_SIZE.value,
              sequence_length=SEQUENCE_LENGTH.value,
          )
          for split in [tfds.Split.TRAIN, tfds.Split.TEST]
      }
  
      # Make loss, sampler, and optimizer.
      params_init, loss_fn = hk.without_apply_rng(
          hk.transform(sequence_loss)
      )
      _, sample_fn = hk.without_apply_rng(hk.transform(sample))
      opt_init, _ = make_optimizer()
  
      loss_fn = jax.jit(loss_fn)
      sample_fn = jax.jit(sample_fn, static_argnums=[3])
  
      # Initialize training state.
      rng = hk.PRNGSequence(SEED.value)
      initial_params = params_init(next(rng), next(train_data))
      initial_opt_state = opt_init(initial_params)
      state = TrainingState(
          params=initial_params, opt_state=initial_opt_state
      )
  
      # Training loop.
      for step in range(TRAINING_STEPS.value + 1):
          # Do a batch of SGD.
          train_batch = next(train_data)
          state = update(state, train_batch)
  
          # Periodically generate samples.
          if step % SAMPLING_INTERVAL.value == 0:
              context = train_batch["input"][
                  :, 0
              ]  # First element of training batch.
              assert context.ndim == 1
              rng_key = next(rng)
              samples = sample_fn(
                  state.params, rng_key, context, SAMPLE_LENGTH.value
              )
  
              prompt = dataset.decode(context)
              continuation = dataset.decode(samples)
  
              logging.info("Prompt: %s", prompt)
              logging.info("Continuation: %s", continuation)
  
          # Periodically evaluate training and test loss.
          if step % EVALUATION_INTERVAL.value == 0:
              for split, ds in eval_data.items():
                  eval_batch = next(ds)
                  loss = loss_fn(state.params, eval_batch)
                  logging.info(
                      {
                          "step": step,
                          "loss": float(loss),
                          "split": split,
                      }
                  )
  
imports: from typing import Any, NamedTuple
  from absl import app
  from absl import flags
  from absl import logging
  import haiku
  from examples.rnn import dataset
  import jax
  from jax import lax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 6 for calling a global function with example usage: main(_) without return types