# coding=utf-8
# Copyright 2021 The Uncertainty Baselines Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Heteroscedastic ViT on JFT-300M."""

from functools import partial  # pylint: disable=g-importing-member so standard
import multiprocessing
import numbers
import os

from absl import app
from absl import flags
from absl import logging
from clu import parameter_overview
from clu import periodic_actions
import flax
import flax.jax_utils as flax_utils
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import robustness_metrics as rm
import tensorflow as tf
from tensorflow.io import gfile
import uncertainty_baselines as ub
import cifar10h_utils  # local file import


# TODO(dusenberrymw): Open-source remaining imports.
fewshot = None
input_pipeline = None
resformer = None
u = None
pp_builder = None
xm = None
xm_api = None


ml_collections.config_flags.DEFINE_config_file(
    'config', None, 'Training configuration.', lock_config=True)

flags.DEFINE_string('output_dir', default=None, help='Work unit directory.')
flags.DEFINE_integer(
    'num_cores', default=None, help='Unused. How many devices being used.')
flags.DEFINE_boolean(
    'use_gpu', default=None, help='Unused. Whether or not running on GPU.')
flags.DEFINE_string('tpu', None,
                    'Unused. Name of the TPU. Only used if use_gpu is False.')
flags.DEFINE_integer('seed', default=0, help='Random seed.')

FLAGS = flags.FLAGS

# Adds jax flags to the program.
jax.config.parse_flags_with_absl()


def main(argv):
  del argv

  config = FLAGS.config
  output_dir = FLAGS.output_dir

  if config.get('dataset_dir'):
    logging.info('data_dir=%s', config.dataset_dir)
  logging.info('Output dir: %s', output_dir)

  save_checkpoint_path = None
  if config.get('checkpoint_steps'):
    gfile.makedirs(output_dir)
    save_checkpoint_path = os.path.join(output_dir, 'checkpoint.npz')

  # The pool is used to perform misc operations such as logging in async way.
  pool = multiprocessing.pool.ThreadPool()

  # This seed makes the Jax part of things (like model init) deterministic.
  # However, full training still won't be deterministic, for example due to the
  # tf.data pipeline not being deterministic even if we would set TF seed.
  rng = jax.random.PRNGKey(config.get('seed', 0))

  xm_xp = None
  xm_wu = None
  def write_note(note):
    if jax.host_id() == 0:
      logging.info('NOTE: %s', note)
  write_note('Initializing...')

  fillin = lambda *_: None
  # Verify settings to make sure no checkpoints are accidentally missed.
  if config.get('keep_checkpoint_steps'):
    assert config.get('checkpoint_steps'), 'Specify `checkpoint_steps`.'
    assert config.keep_checkpoint_steps % config.checkpoint_steps == 0, (
        f'`keep_checkpoint_steps` ({config.checkpoint_steps}) should be'
        f'divisible by `checkpoint_steps ({config.checkpoint_steps}).`')

  batch_size = config.batch_size
  batch_size_eval = config.get('batch_size_eval', batch_size)
  if (batch_size % jax.device_count() != 0 or
      batch_size_eval % jax.device_count() != 0):
    raise ValueError(f'Batch sizes ({batch_size} and {batch_size_eval}) must '
                     f'be divisible by device number ({jax.device_count()})')

  local_batch_size = batch_size // jax.host_count()
  local_batch_size_eval = batch_size_eval // jax.host_count()
  logging.info(
      'Global batch size %d on %d hosts results in %d local batch size. '
      'With %d devices per host (%d devices total), that\'s a %d per-device '
      'batch size.',
      batch_size, jax.host_count(), local_batch_size,
      jax.local_device_count(), jax.device_count(),
      local_batch_size // jax.local_device_count())

  write_note('Initializing train dataset...')
  train_ds = input_pipeline.get_data(
      dataset=config.dataset,
      split=config.train_split,
      data_dir=fillin(config.get('dataset_dir')),
      batch_size=local_batch_size,
      preprocess_fn=pp_builder.get_preprocess_fn(config.pp_train),
      shuffle_buffer_size=config.shuffle_buffer_size,
      prefetch=config.get('prefetch_to_host', 2),
      cache=False)

  # Start prefetching already.
  train_iter = u.start_input_pipeline(
      train_ds, config.get('prefetch_to_device', 1), pad=local_batch_size)
  # We always pad to local_batch_size_eval even when less would be enough in
  # order to minimize memory fragmentation.

  write_note('Initializing val dataset(s)...')
  def _get_val_split(dataset, split, pp_eval, data_dir=None):
    # We do ceil rounding such that we include the last incomplete batch.
    nval_img = input_pipeline.get_num_examples(
        dataset, split, data_dir=fillin(data_dir))
    val_steps = int(np.ceil(nval_img / batch_size_eval))
    logging.info('Running validation for %d steps for %s, %s', val_steps,
                 dataset, split)

    val_it = input_pipeline.get_data(
        dataset=dataset,
        split=split,
        data_dir=fillin(data_dir),
        batch_size=local_batch_size_eval,
        preprocess_fn=pp_builder.get_preprocess_fn(pp_eval),
        cache=config.get('val_cache', 'batched'),
        repeat_after_batching=True,
        prefetch=0,  # Save memory since we cache.
        drop_remainder=False,
        shuffle_files=False)
    val_it = u.start_input_pipeline(
        val_it, config.get('prefetch_to_device', 1), pad=local_batch_size_eval)

    return (val_it, val_steps)

  if isinstance(config.val_split, str):
    val_ds = {'val': _get_val_split(config.dataset, config.val_split,
                                    config.pp_eval, config.get('dataset_dir'))}
  else:
    val_ds = {t[0]: _get_val_split(*t[1:]) for t in config.val_split}

  if config.get('eval_on_cifar_10h'):
    val_steps = int(np.ceil(10000 / batch_size_eval))

    cifar10h_dataset = cifar10h_utils.load_ds()

    val_it = input_pipeline.make_pipeline(
        data=cifar10h_dataset,
        batch_size=local_batch_size_eval,
        preprocess_fn=pp_builder.get_preprocess_fn(config.pp_eval_cifar_10h),
        cache=config.get('val_cache', 'batched'),
        repeats=None,
        repeat_after_batching=True,
        prefetch=0,
        drop_remainder=False,
        shuffle_buffer_size=None,
        ignore_errors=False,
        filter_fn=None)
    val_it = u.start_input_pipeline(
        val_it, config.get('prefetch_to_device', 1), pad=local_batch_size_eval)

    val_ds['cifar_10h'] = (val_it, val_steps)

  ood_ds = None
  if config.get('ood_dataset'):
    logging.info('loading OOD dataset = %s', config.get('ood_dataset'))
    if isinstance(config.ood_split, str):
      ood_ds = {
          'ind':
              _get_val_split(config.dataset, config.ood_split, config.pp_eval,
                             config.get('data_dir')),
          'ood':
              _get_val_split(config.ood_dataset, config.ood_split,
                             config.pp_eval, config.get('data_dir')),
      }
    else:
      raise NotImplementedError(
          'Only string type of val_split is supported! Got val_split=%s!' %
          str(config.ood_split))

  ntrain_img = input_pipeline.get_num_examples(
      config.dataset, config.train_split,
      data_dir=fillin(config.get('dataset_dir')))
  steps_per_epoch = ntrain_img / batch_size

  if config.get('num_epochs'):
    total_steps = int(config.num_epochs * steps_per_epoch)
    assert not config.get('total_steps'), 'Set either num_epochs or total_steps'
  else:
    total_steps = config.total_steps

  logging.info(
      'Running for %d steps, that means %f epochs and %f steps per epoch',
      total_steps, total_steps * batch_size / ntrain_img, steps_per_epoch)
  mw = u.BigVisionMetricWriter(xm_xp.id, xm_wu.id, steps_per_epoch)

  write_note('Initializing model...')
  logging.info('config.model = %s', config.get('model'))
  model = ub.models.het_vision_transformer(
      num_classes=config.num_classes, **config.get('model', {}))

  # We want all parameters to be created in host RAM, not on any device, they'll
  # be sent there later as needed, otherwise we already encountered two
  # situations where we allocate them twice.
  @partial(jax.jit, backend='cpu')
  def init(rng):
    image_size = tuple(train_ds.element_spec['image'].shape[1:])
    dummy_input = jnp.zeros((local_batch_size,) + image_size, jnp.float32)

    init_rngs = {'params': rng, 'diag_noise_samples': (rng + 1) * 7,
                 'standard_norm_noise_samples': (rng + 3) * 13}

    params = flax.core.unfreeze(model.init(init_rngs, dummy_input,
                                           train=False))['params']

    # Set bias in the head to a low value, such that loss is small initially.
    if 'head' in params:
      params['head']['loc_layer']['bias'] = jnp.full_like(
          params['head']['loc_layer']['bias'], config.get('init_head_bias', 0))

    return params

  rng, rng_init = jax.random.split(rng)
  params_cpu = init(rng_init)

  if jax.host_id() == 0:
    num_params = sum(p.size for p in jax.tree_flatten(params_cpu)[0])
    parameter_overview.log_parameter_overview(params_cpu)
    mw.measure('num_params', num_params)

  @partial(jax.pmap, axis_name='batch')
  def evaluation_fn(params, images, labels, mask):
    # Ignore the entries with all zero labels for evaluation.
    mask *= labels.max(axis=1)
    logits, _ = model.apply({'params': flax.core.freeze(params)},
                            images,
                            train=False,
                            rngs={
                                'dropout': rng,
                                'diag_noise_samples': (rng + 1) * 7,
                                'standard_norm_noise_samples': (rng + 3) * 13})

    losses = getattr(u, config.get('loss', 'sigmoid_xent'))(
        logits=logits, labels=labels, reduction=False)
    loss = jax.lax.psum(losses * mask, axis_name='batch')

    top1_idx = jnp.argmax(logits, axis=1)
    # Extracts the label at the highest logit index for each image.
    top1_correct = jnp.take_along_axis(labels, top1_idx[:, None], axis=1)[:, 0]
    ncorrect = jax.lax.psum(top1_correct * mask, axis_name='batch')
    n = jax.lax.psum(mask, axis_name='batch')

    metric_args = jax.lax.all_gather([logits, labels, mask], axis_name='batch')

    return ncorrect, loss, n, metric_args

  @partial(jax.pmap, axis_name='batch')
  def cifar_10h_evaluation_fn(params, images, labels, mask):
    logits, _ = model.apply({'params': flax.core.freeze(params)},
                            images,
                            train=False,
                            rngs={
                                'dropout': rng,
                                'diag_noise_samples': (rng + 1) * 7,
                                'standard_norm_noise_samples': (rng + 3) * 13})

    losses = getattr(u, config.get('loss', 'softmax_xent'))(
        logits=logits, labels=labels, reduction=False)
    loss = jax.lax.psum(losses, axis_name='batch')

    top1_idx = jnp.argmax(logits, axis=1)
    # Extracts the label at the highest logit index for each image.
    one_hot_labels = jnp.eye(10)[jnp.argmax(labels, axis=1)]

    top1_correct = jnp.take_along_axis(
        one_hot_labels, top1_idx[:, None], axis=1)[:, 0]
    ncorrect = jax.lax.psum(top1_correct, axis_name='batch')
    n = jax.lax.psum(one_hot_labels, axis_name='batch')

    metric_args = jax.lax.all_gather([logits, labels, mask],
                                     axis_name='batch')
    return ncorrect, loss, n, metric_args

  # Setup function for computing representation.
  @partial(jax.pmap, axis_name='batch')
  def representation_fn(params, images, labels, mask):
    _, outputs = model.apply({'params': flax.core.freeze(params)},
                             images,
                             train=False,
                             rngs={
                                 'dropout': rng,
                                 'diag_noise_samples': (rng + 1) * 7,
                                 'standard_norm_noise_samples': (rng + 3) * 13})
    representation = outputs[config.fewshot.representation_layer]
    representation = jax.lax.all_gather(representation, 'batch')
    labels = jax.lax.all_gather(labels, 'batch')
    mask = jax.lax.all_gather(mask, 'batch')
    return representation, labels, mask

  # Load the optimizer from flax.
  opt_name = config.get('optim_name')
  write_note(f'Initializing {opt_name} optimizer...')
  opt_def = getattr(flax.optim, opt_name)(**config.get('optim', {}))

  # We jit this, such that the arrays that are created are created on the same
  # device as the input is, in this case the CPU. Else they'd be on device[0].
  opt_cpu = jax.jit(opt_def.create)(params_cpu)

  @partial(jax.pmap, axis_name='batch', donate_argnums=(0,))
  def update_fn(opt, lr, images, labels, rng):
    """Update step."""

    measurements = {}

    if config.get('mixup') and config.mixup.p:
      rng, (images, labels), _ = u.mixup(rng, images, labels, **config.mixup)

    # Get device-specific loss rng.
    rng, rng_model = jax.random.split(rng, 2)
    rng_model_local = jax.random.fold_in(rng_model, jax.lax.axis_index('batch'))

    def loss_fn(params, images, labels):
      logits, _ = model.apply(
          {'params': flax.core.freeze(params)}, images,
          train=True, rngs={
              'dropout': rng_model_local,
              'diag_noise_samples': (rng_model_local + 1) * 7,
              'standard_norm_noise_samples': (rng_model_local + 3) * 13})
      return getattr(u, config.get('loss', 'sigmoid_xent'))(
          logits=logits, labels=labels)

    # Implementation considerations compared and summarized at
    # https://docs.google.com/document/d/1g3kMEvqu1DOawaflKNyUsIoQ4yIVEoyE5ZlIPkIl4Lc/edit?hl=en#
    l, g = u.accumulate_gradient(jax.value_and_grad(loss_fn), opt.target,
                                 images, labels,
                                 config.get('grad_accum_steps'))
    l, g = jax.lax.pmean((l, g), axis_name='batch')

    # Log the gradient norm only if we need to compute it anyways (clipping)
    # or if we don't use grad_accum_steps, as they interact badly.
    if config.get('grad_accum_steps', 1) == 1 or config.get('grad_clip_norm'):
      grads, _ = jax.tree_flatten(g)
      l2_g = jnp.sqrt(sum([jnp.vdot(p, p) for p in grads]))
      measurements['l2_grads'] = l2_g

    # Optionally resize the global gradient to a maximum norm. We found this
    # useful in some cases across optimizers, hence it's in the main loop.
    if config.get('grad_clip_norm'):
      g_factor = jnp.minimum(1.0, config.grad_clip_norm / l2_g)
      g = jax.tree_map(lambda p: g_factor * p, g)
    opt = opt.apply_gradient(g, learning_rate=lr)

    decay_rules = config.get('weight_decay', []) or []
    if isinstance(decay_rules, numbers.Number):
      decay_rules = [('.*kernel.*', decay_rules)]
    sched_m = lr/config.lr.base if config.get('weight_decay_decouple') else lr
    def decay_fn(v, wd):
      return (1.0 - sched_m * wd) * v
    opt = opt.replace(target=u.tree_map_with_regex(
        decay_fn, opt.target, decay_rules, name='weight decay'))

    params, _ = jax.tree_flatten(opt.target)
    measurements['l2_params'] = jnp.sqrt(sum([jnp.vdot(p, p) for p in params]))

    return opt, l, rng, measurements

  # Other things besides optimizer state to be stored.
  checkpoint_extra = dict(accum_train_time=0.0)

  # Decide how to initialize training. The order is important.
  # 1. Always resumes from the existing checkpoint, e.g. resumes a finetune job.
  # 2. Resume from a previous checkpoint, e.g. start a cooldown training job.
  # 3. Initialize model from something, e,g, start a fine-tuning job.
  # 4. Train from scratch.
  resume_checkpoint_path = None
  if save_checkpoint_path and gfile.exists(save_checkpoint_path):
    resume_checkpoint_path = save_checkpoint_path
  elif config.get('resume'):
    resume_checkpoint_path = fillin(config.resume)
  if resume_checkpoint_path:
    write_note('Resume training from checkpoint...')
    checkpoint = {'opt': opt_cpu, 'extra': checkpoint_extra}
    _, checkpoint_tree = jax.tree_flatten(checkpoint)
    loaded = u.load_checkpoint(checkpoint_tree, resume_checkpoint_path)
    # bfloat16 type gets lost when data is saved to disk, so we recover it.
    checkpoint = jax.tree_map(u.recover_dtype, loaded)
    opt_cpu, checkpoint_extra = checkpoint['opt'], checkpoint['extra']
  elif config.get('model_init'):
    write_note(f'Initialize model from {config.model_init}...')
    # TODO(dusenberrymw): Replace and test load function.
    reinit_params = ['head/scale_layer_homoscedastic/kernel',
                     'head/scale_layer_homoscedastic/bias',
                     'head/scale_layer_heteroscedastic/kernel',
                     'head/scale_layer_heteroscedastic/bias',
                     'head/loc_layer/kernel', 'head/diag_layer/kernel',
                     'head/loc_layer/bias', 'head/diag_layer/bias']
    for param in reinit_params:
      if param in params_cpu:
        del params_cpu[param]

    loaded = resformer.load(params_cpu, config.model_init, config.get('model'),
                            reinit_params)
    opt_cpu = opt_cpu.replace(target=loaded)
    if jax.host_id() == 0:
      logging.info('Restored parameter overview:')
      parameter_overview.log_parameter_overview(loaded)

  write_note('Kicking off misc stuff...')
  first_step = int(opt_cpu.state.step)  # Might be a DeviceArray type.
  chrono = u.Chrono(first_step, total_steps, batch_size,
                    checkpoint_extra['accum_train_time'])
  # Note: switch to ProfileAllHosts() if you need to profile all hosts.
  # (Xprof data become much larger and take longer to load for analysis)
  profiler = periodic_actions.Profile(
      # Create profile after every restart to analyze pre-emption related
      # problems and assure we get similar performance in every run.
      logdir=output_dir, first_profile=first_step + 10)

  # Prepare the learning-rate and pre-fetch it to device to avoid delays.
  lr_fn = u.create_learning_rate_schedule(
      batch_size, total_steps, steps_per_epoch, **config.get('lr', {}))
  # TODO(dusenberrymw): According to flax docs, prefetching shouldn't be
  # necessary for TPUs.
  lr_iter = u.prefetch_scalar(map(lr_fn, range(first_step, total_steps)),
                              config.get('prefetch_to_device', 1))

  write_note(f'Replicating...\n{chrono.note}')
  opt_repl = flax_utils.replicate(opt_cpu)

  write_note(f'Initializing few-shotters...\n{chrono.note}')
  if 'fewshot' in config:
    fewshotter = fewshot.FewShotEvaluator(
        representation_fn, config.fewshot,
        config.fewshot.get('batch_size') or batch_size_eval)

  rng, rng_loop = jax.random.split(rng, 2)
  rngs_loop = flax_utils.replicate(rng_loop)
  checkpoint_writer = None

  # Note: we return the train loss, val loss, and fewshot best l2s for use in
  # reproducibility unit tests.
  train_loss = -jnp.inf
  val_loss = -jnp.inf
  results = {'dummy': {(0, 1): -jnp.inf}}

  write_note(f'First step compilations...\n{chrono.note}')
  # Using a python integer for step here, because opt.state.step is allocated
  # on TPU during replication.
  for step, train_batch, lr_repl in zip(
      range(first_step + 1, total_steps + 1), train_iter, lr_iter):
    mw.step_start(step)

    with jax.profiler.TraceContext('train_step', step_num=step, _r=1):
      opt_repl, loss_value, rngs_loop, extra_measurements = update_fn(
          opt_repl,
          lr_repl,
          train_batch['image'],
          train_batch['labels'],
          rng=rngs_loop)

    if jax.host_id() == 0:
      profiler(step)

    # Checkpoint saving
    if u.itstime(step, config.get('checkpoint_steps'), total_steps, host=0):
      write_note('Checkpointing...')
      chrono.pause()
      u.checkpointing_timeout(checkpoint_writer,
                              config.get('checkpoint_timeout', 1))
      checkpoint_extra['accum_train_time'] = chrono.accum_train_time
      # We need to transfer the weights over now or else we risk keeping them
      # alive while they'll be updated in a future step, creating hard to debug
      # memory errors (see b/160593526). Also, takes device 0's params only.
      opt_cpu = jax.tree_map(lambda x: np.array(x[0]), opt_repl)

      # Check whether we want to keep a copy of the current checkpoint.
      copy_step = None
      if u.itstime(step, config.get('keep_checkpoint_steps'), total_steps):
        write_note('Keeping a checkpoint copy...')
        copy_step = step

      # Checkpoint should be a nested dictionary or FLAX datataclasses from
      # `flax.struct`. Both can be present in a checkpoint.
      checkpoint = {'opt': opt_cpu, 'extra': checkpoint_extra}
      checkpoint_writer = pool.apply_async(
          u.save_checkpoint, (checkpoint, save_checkpoint_path, copy_step))
      chrono.resume()

    # Report training progress
    if u.itstime(step, config.log_training_steps, total_steps, host=0):
      write_note('Reporting training progress...')
      train_loss = loss_value[0]  # Keep to return for reproducibility tests.
      mw.measure('learning_rate', lr_repl[0])
      mw.measure('training_loss', loss_value[0])
      for name, value in extra_measurements.items():
        mw.measure(name, value[0])
      chrono.tick(step, mw.measure, write_note)

    # Report validation performance
    if u.itstime(step, config.log_eval_steps, total_steps):
      write_note('Evaluating on the validation set...')
      chrono.pause()
      for val_name, (val_iter, val_steps) in val_ds.items():
        ncorrect, loss, nseen = 0, 0, 0
        ece_num_bins = config.get('ece_num_bins', 15)
        ece = rm.metrics.ExpectedCalibrationError(num_bins=ece_num_bins)
        label_diversity = tf.keras.metrics.Mean()
        sample_diversity = tf.keras.metrics.Mean()
        ged = tf.keras.metrics.Mean()
        for _, batch in zip(range(val_steps), val_iter):
          if val_name == 'cifar_10h':
            batch_ncorrect, batch_losses, batch_n, batch_metric_args = (
                cifar_10h_evaluation_fn(opt_repl.target, batch['image'],
                                        batch['labels'], batch['mask']))
          else:
            batch_ncorrect, batch_losses, batch_n, batch_metric_args = (
                evaluation_fn(opt_repl.target, batch['image'],
                              batch['labels'], batch['mask']))
          # All results are a replicated array shaped as follows:
          # (local_devices, per_device_batch_size, elem_shape...)
          # with each local device's entry being identical as they got psum'd.
          # So let's just take the first one to the host as numpy.
          ncorrect += np.sum(np.array(batch_ncorrect[0]))
          loss += np.sum(np.array(batch_losses[0]))
          nseen += np.sum(np.array(batch_n[0]))

          # Here we parse batch_metric_args to compute complicated metrics
          # such as ECE.
          logits, labels, masks = batch_metric_args
          masks = np.array(masks[0], dtype=np.bool)
          # From one-hot to integer labels, as required by ECE.
          int_labels = np.argmax(np.array(labels[0]), axis=-1)
          logits = np.array(logits[0])
          probs = jax.nn.softmax(logits)
          for p, l, m, label in zip(probs, int_labels, masks, labels[0]):
            ece.add_batch(p[m, :], label=l[m])

            if val_name == 'cifar_10h':
              batch_label_diversity, batch_sample_diversity, batch_ged = cifar10h_utils.generalized_energy_distance(
                  label[m], p[m, :], 10)
              label_diversity.update_state(batch_label_diversity)
              sample_diversity.update_state(batch_sample_diversity)
              ged.update_state(batch_ged)

        val_loss = loss / nseen  # Keep to return for reproducibility tests.
        mw.measure(f'{val_name}_prec@1', ncorrect / nseen)
        mw.measure(f'{val_name}_loss', val_loss)
        mw.measure(f'{val_name}_ece', float(ece.result()['ece']))
        if val_name == 'cifar_10h':
          mw.measure(
              f'{val_name}_label_diversity', float(label_diversity.result()))
          mw.measure(
              f'{val_name}_sample_diversity', float(sample_diversity.result()))
          mw.measure(f'{val_name}_ged', float(ged.result()))

      # OOD eval
      if ood_ds:
        ood_metrics = {
            'auroc':
                tf.keras.metrics.AUC(
                    curve='ROC', summation_method='interpolation'),
            'auprc':
                tf.keras.metrics.AUC(
                    curve='PR', summation_method='interpolation')
        }
        for metric in ood_metrics.values():
          metric.reset_states()
        for val_name, (val_iter, val_steps) in ood_ds.items():
          for _, batch in zip(range(val_steps), val_iter):
            batch_ncorrect, batch_losses, batch_n, batch_metric_args = evaluation_fn(
                opt_repl.target, batch['image'], batch['labels'], batch['mask'])
            # All results are a replicated array shaped as follows:
            # (local_devices, per_device_batch_size, elem_shape...)
            # with each local device's entry being identical as they got psum'd.
            # So let's just take the first one to the host as numpy.
            ncorrect += np.sum(np.array(batch_ncorrect[0]))
            loss += np.sum(np.array(batch_losses[0]))
            nseen += np.sum(np.array(batch_n[0]))

            # Here we parse batch_metric_args to compute
            # complicated metrics such as ECE and OOD AUROC
            logits, _, masks = batch_metric_args
            probs = jax.nn.softmax(logits[0], axis=-1)
            probs = probs[jnp.array(masks[0], dtype=bool)]
            confs = jnp.max(probs, axis=-1)
            ood_labels = np.ones_like(
                confs) if val_name == 'ind' else np.zeros_like(confs)
            for metric in ood_metrics.values():
              metric.update_state(ood_labels, confs)
          if val_name == 'ind':
            val_loss = loss / nseen  # Keep to return for reproducibility tests.
            mw.measure(f'{val_name}_prec@1', ncorrect / nseen)
            mw.measure(f'{val_name}_loss', val_loss)
        for name, value in ood_metrics.items():
          mw.measure(f'ood_{name}', value.result())
      chrono.resume()

    if 'fewshot' in config:
      # Compute few-shot on-the-fly evaluation.
      if u.itstime(step, config.fewshot.log_steps, total_steps):
        chrono.pause()
        write_note(f'Few-shot evaluation...\n{chrono.note}')
        # Keep `results` to return for reproducibility tests.
        results, best_l2 = fewshotter.run_all(opt_repl.target,
                                              config.fewshot.datasets)
        fewshotter.walk_results(mw.measure, results, best_l2)
        chrono.resume()
    mw.step_end()

  write_note(f'Done!\n{chrono.note}')
  pool.close()
  pool.join()
  mw.close()

  # Return final training loss, validation loss, and fewshot results for
  # reproducibility test cases.
  return train_loss, val_loss, results


if __name__ == '__main__':
  # TODO(dusenberrymw): Refactor `main` such that there is a `train_eval`
  # function that returns values for tests and does not directly access flags,
  # and then have `main` return None.

  def _main(argv):
    main(argv)

  app.run(_main)  # Ignore the returned values from `main`.

numpy.ones_like tensorflow.io.gfile.exists tensorflow.keras.metrics.AUC tensorflow.io.gfile.makedirs numpy.ceil numpy.zeros_like numpy.array tensorflow.keras.metrics.Mean
[(51, 'ml_collections.config_flags.DEFINE_config_file', 'ml_collections.config_flags.DEFINE_config_file', (['"""config"""', 'None', '"""Training configuration."""'], {'lock_config': '(True)'}), False, 'import ml_collections\n'), (54, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""output_dir"""'], {'default': 'None', 'help': '"""Work unit directory."""'}), False, 'from absl import flags\n'), (55, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""num_cores"""'], {'default': 'None', 'help': '"""Unused. How many devices being used."""'}), False, 'from absl import flags\n'), (57, 'absl.flags.DEFINE_boolean', 'flags.DEFINE_boolean', (['"""use_gpu"""'], {'default': 'None', 'help': '"""Unused. Whether or not running on GPU."""'}), False, 'from absl import flags\n'), (59, 'absl.flags.DEFINE_string', 'flags.DEFINE_string', (['"""tpu"""', 'None', '"""Unused. Name of the TPU. Only used if use_gpu is False."""'], {}), False, 'from absl import flags\n'), (61, 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', (['"""seed"""'], {'default': '(0)', 'help': '"""Random seed."""'}), False, 'from absl import flags\n'), (66, 'jax.config.parse_flags_with_absl', 'jax.config.parse_flags_with_absl', ([], {}), False, 'import jax\n'), (77, 'absl.logging.info', 'logging.info', (['"""Output dir: %s"""', 'output_dir'], {}), False, 'from absl import logging\n'), (85, 'multiprocessing.pool.ThreadPool', 'multiprocessing.pool.ThreadPool', ([], {}), False, 'import multiprocessing\n'), (222, 'absl.logging.info', 'logging.info', (['"""Running for %d steps, that means %f epochs and %f steps per epoch"""', 'total_steps', '(total_steps * batch_size / ntrain_img)', 'steps_per_epoch'], {}), False, 'from absl import logging\n'), (235, 'functools.partial', 'partial', (['jax.jit'], {'backend': '"""cpu"""'}), False, 'from functools import partial\n'), (253, 'jax.random.split', 'jax.random.split', (['rng'], {}), False, 'import jax\n'), (261, 'functools.partial', 'partial', (['jax.pmap'], {'axis_name': '"""batch"""'}), False, 'from functools import partial\n'), (287, 'functools.partial', 'partial', (['jax.pmap'], {'axis_name': '"""batch"""'}), False, 'from functools import partial\n'), (315, 'functools.partial', 'partial', (['jax.pmap'], {'axis_name': '"""batch"""'}), False, 'from functools import partial\n'), (339, 'functools.partial', 'partial', (['jax.pmap'], {'axis_name': '"""batch"""', 'donate_argnums': '(0,)'}), False, 'from functools import partial\n'), (444, 'clu.periodic_actions.Profile', 'periodic_actions.Profile', ([], {'logdir': 'output_dir', 'first_profile': '(first_step + 10)'}), False, 'from clu import periodic_actions\n'), (458, 'flax.jax_utils.replicate', 'flax_utils.replicate', (['opt_cpu'], {}), True, 'import flax.jax_utils as flax_utils\n'), (466, 'jax.random.split', 'jax.random.split', (['rng', '(2)'], {}), False, 'import jax\n'), (467, 'flax.jax_utils.replicate', 'flax_utils.replicate', (['rng_loop'], {}), True, 'import flax.jax_utils as flax_utils\n'), (658, 'absl.app.run', 'app.run', (['_main'], {}), False, 'from absl import app\n'), (76, 'absl.logging.info', 'logging.info', (['"""data_dir=%s"""', 'config.dataset_dir'], {}), False, 'from absl import logging\n'), (81, 'tensorflow.io.gfile.makedirs', 'gfile.makedirs', (['output_dir'], {}), False, 'from tensorflow.io import gfile\n'), (82, 'os.path.join', 'os.path.join', (['output_dir', '"""checkpoint.npz"""'], {}), False, 'import os\n'), (114, 'jax.host_count', 'jax.host_count', ([], {}), False, 'import jax\n'), (115, 'jax.host_count', 'jax.host_count', ([], {}), False, 'import jax\n'), (120, 'jax.host_count', 'jax.host_count', ([], {}), False, 'import jax\n'), (121, 'jax.local_device_count', 'jax.local_device_count', ([], {}), False, 'import jax\n'), (121, 'jax.device_count', 'jax.device_count', ([], {}), False, 'import jax\n'), (147, 'absl.logging.info', 'logging.info', (['"""Running validation for %d steps for %s, %s"""', 'val_steps', 'dataset', 'split'], {}), False, 'from absl import logging\n'), (175, 'cifar10h_utils.load_ds', 'cifar10h_utils.load_ds', ([], {}), False, 'import cifar10h_utils\n'), (238, 'jax.numpy.zeros', 'jnp.zeros', (['((local_batch_size,) + image_size)', 'jnp.float32'], {}), True, 'import jax.numpy as jnp\n'), (256, 'jax.host_id', 'jax.host_id', ([], {}), False, 'import jax\n'), (258, 'clu.parameter_overview.log_parameter_overview', 'parameter_overview.log_parameter_overview', (['params_cpu'], {}), False, 'from clu import parameter_overview\n'), (275, 'jax.lax.psum', 'jax.lax.psum', (['(losses * mask)'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (277, 'jax.numpy.argmax', 'jnp.argmax', (['logits'], {'axis': '(1)'}), True, 'import jax.numpy as jnp\n'), (280, 'jax.lax.psum', 'jax.lax.psum', (['(top1_correct * mask)'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (281, 'jax.lax.psum', 'jax.lax.psum', (['mask'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (283, 'jax.lax.all_gather', 'jax.lax.all_gather', (['[logits, labels, mask]'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (299, 'jax.lax.psum', 'jax.lax.psum', (['losses'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (301, 'jax.numpy.argmax', 'jnp.argmax', (['logits'], {'axis': '(1)'}), True, 'import jax.numpy as jnp\n'), (307, 'jax.lax.psum', 'jax.lax.psum', (['top1_correct'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (308, 'jax.lax.psum', 'jax.lax.psum', (['one_hot_labels'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (310, 'jax.lax.all_gather', 'jax.lax.all_gather', (['[logits, labels, mask]'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (325, 'jax.lax.all_gather', 'jax.lax.all_gather', (['representation', '"""batch"""'], {}), False, 'import jax\n'), (326, 'jax.lax.all_gather', 'jax.lax.all_gather', (['labels', '"""batch"""'], {}), False, 'import jax\n'), (327, 'jax.lax.all_gather', 'jax.lax.all_gather', (['mask', '"""batch"""'], {}), False, 'import jax\n'), (337, 'jax.jit', 'jax.jit', (['opt_def.create'], {}), False, 'import jax\n'), (349, 'jax.random.split', 'jax.random.split', (['rng', '(2)'], {}), False, 'import jax\n'), (367, 'jax.lax.pmean', 'jax.lax.pmean', (['(l, g)'], {'axis_name': '"""batch"""'}), False, 'import jax\n'), (392, 'jax.tree_flatten', 'jax.tree_flatten', (['opt.target'], {}), False, 'import jax\n'), (406, 'tensorflow.io.gfile.exists', 'gfile.exists', (['save_checkpoint_path'], {}), False, 'from tensorflow.io import gfile\n'), (413, 'jax.tree_flatten', 'jax.tree_flatten', (['checkpoint'], {}), False, 'import jax\n'), (416, 'jax.tree_map', 'jax.tree_map', (['u.recover_dtype', 'loaded'], {}), False, 'import jax\n'), (95, 'jax.host_id', 'jax.host_id', ([], {}), False, 'import jax\n'), (96, 'absl.logging.info', 'logging.info', (['"""NOTE: %s"""', 'note'], {}), False, 'from absl import logging\n'), (122, 'jax.local_device_count', 'jax.local_device_count', ([], {}), False, 'import jax\n'), (146, 'numpy.ceil', 'np.ceil', (['(nval_img / batch_size_eval)'], {}), True, 'import numpy as np\n'), (173, 'numpy.ceil', 'np.ceil', (['(10000 / batch_size_eval)'], {}), True, 'import numpy as np\n'), (279, 'jax.numpy.take_along_axis', 'jnp.take_along_axis', (['labels', 'top1_idx[:, (None)]'], {'axis': '(1)'}), True, 'import jax.numpy as jnp\n'), (303, 'jax.numpy.eye', 'jnp.eye', (['(10)'], {}), True, 'import jax.numpy as jnp\n'), (305, 'jax.numpy.take_along_axis', 'jnp.take_along_axis', (['one_hot_labels', 'top1_idx[:, (None)]'], {'axis': '(1)'}), True, 'import jax.numpy as jnp\n'), (350, 'jax.lax.axis_index', 'jax.lax.axis_index', (['"""batch"""'], {}), False, 'import jax\n'), (364, 'jax.value_and_grad', 'jax.value_and_grad', (['loss_fn'], {}), False, 'import jax\n'), (372, 'jax.tree_flatten', 'jax.tree_flatten', (['g'], {}), False, 'import jax\n'), (379, 'jax.numpy.minimum', 'jnp.minimum', (['(1.0)', '(config.grad_clip_norm / l2_g)'], {}), True, 'import jax.numpy as jnp\n'), (380, 'jax.tree_map', 'jax.tree_map', (['(lambda p: g_factor * p)', 'g'], {}), False, 'import jax\n'), (483, 'jax.profiler.TraceContext', 'jax.profiler.TraceContext', (['"""train_step"""'], {'step_num': 'step', '_r': '(1)'}), False, 'import jax\n'), (491, 'jax.host_id', 'jax.host_id', ([], {}), False, 'import jax\n'), (109, 'jax.device_count', 'jax.device_count', ([], {}), False, 'import jax\n'), (110, 'jax.device_count', 'jax.device_count', ([], {}), False, 'import jax\n'), (265, 'flax.core.freeze', 'flax.core.freeze', (['params'], {}), False, 'import flax\n'), (289, 'flax.core.freeze', 'flax.core.freeze', (['params'], {}), False, 'import flax\n'), (303, 'jax.numpy.argmax', 'jnp.argmax', (['labels'], {'axis': '(1)'}), True, 'import jax.numpy as jnp\n'), (317, 'flax.core.freeze', 'flax.core.freeze', (['params'], {}), False, 'import flax\n'), (434, 'jax.host_id', 'jax.host_id', ([], {}), False, 'import jax\n'), (435, 'absl.logging.info', 'logging.info', (['"""Restored parameter overview:"""'], {}), False, 'from absl import logging\n'), (436, 'clu.parameter_overview.log_parameter_overview', 'parameter_overview.log_parameter_overview', (['loaded'], {}), False, 'from clu import parameter_overview\n'), (536, 'robustness_metrics.metrics.ExpectedCalibrationError', 'rm.metrics.ExpectedCalibrationError', ([], {'num_bins': 'ece_num_bins'}), True, 'import robustness_metrics as rm\n'), (537, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {}), True, 'import tensorflow as tf\n'), (538, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {}), True, 'import tensorflow as tf\n'), (539, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {}), True, 'import tensorflow as tf\n'), (111, 'jax.device_count', 'jax.device_count', ([], {}), False, 'import jax\n'), (354, 'flax.core.freeze', 'flax.core.freeze', (['params'], {}), False, 'import flax\n'), (393, 'jax.numpy.vdot', 'jnp.vdot', (['p', 'p'], {}), True, 'import jax.numpy as jnp\n'), (504, 'numpy.array', 'np.array', (['x[0]'], {}), True, 'import numpy as np\n'), (560, 'numpy.array', 'np.array', (['masks[0]'], {'dtype': 'np.bool'}), True, 'import numpy as np\n'), (563, 'numpy.array', 'np.array', (['logits[0]'], {}), True, 'import numpy as np\n'), (564, 'jax.nn.softmax', 'jax.nn.softmax', (['logits'], {}), False, 'import jax\n'), (590, 'tensorflow.keras.metrics.AUC', 'tf.keras.metrics.AUC', ([], {'curve': '"""ROC"""', 'summation_method': '"""interpolation"""'}), True, 'import tensorflow as tf\n'), (593, 'tensorflow.keras.metrics.AUC', 'tf.keras.metrics.AUC', ([], {'curve': '"""PR"""', 'summation_method': '"""interpolation"""'}), True, 'import tensorflow as tf\n'), (257, 'jax.tree_flatten', 'jax.tree_flatten', (['params_cpu'], {}), False, 'import jax\n'), (373, 'jax.numpy.vdot', 'jnp.vdot', (['p', 'p'], {}), True, 'import jax.numpy as jnp\n'), (553, 'numpy.array', 'np.array', (['batch_ncorrect[0]'], {}), True, 'import numpy as np\n'), (554, 'numpy.array', 'np.array', (['batch_losses[0]'], {}), True, 'import numpy as np\n'), (555, 'numpy.array', 'np.array', (['batch_n[0]'], {}), True, 'import numpy as np\n'), (562, 'numpy.array', 'np.array', (['labels[0]'], {}), True, 'import numpy as np\n'), (613, 'jax.nn.softmax', 'jax.nn.softmax', (['logits[0]'], {'axis': '(-1)'}), False, 'import jax\n'), (615, 'jax.numpy.max', 'jnp.max', (['probs'], {'axis': '(-1)'}), True, 'import jax.numpy as jnp\n'), (569, 'cifar10h_utils.generalized_energy_distance', 'cifar10h_utils.generalized_energy_distance', (['label[m]', 'p[(m), :]', '(10)'], {}), False, 'import cifar10h_utils\n'), (606, 'numpy.array', 'np.array', (['batch_ncorrect[0]'], {}), True, 'import numpy as np\n'), (607, 'numpy.array', 'np.array', (['batch_losses[0]'], {}), True, 'import numpy as np\n'), (608, 'numpy.array', 'np.array', (['batch_n[0]'], {}), True, 'import numpy as np\n'), (616, 'numpy.ones_like', 'np.ones_like', (['confs'], {}), True, 'import numpy as np\n'), (617, 'numpy.zeros_like', 'np.zeros_like', (['confs'], {}), True, 'import numpy as np\n'), (614, 'jax.numpy.array', 'jnp.array', (['masks[0]'], {'dtype': 'bool'}), True, 'import jax.numpy as jnp\n')]