def _predict_batch(flax_module,
                   params,
                   batch_stats,
                   batch,
                   output_activation_fn=None):
  """Compute predictions for a batch of data.

  NOTE: We assume that batch_stats has been sync'ed.

  Args:
    flax_module: the evaluation Flax nn.Module (with train=False).
    params: A dict of trainable model parameters. Passed as {'params': params}
      into flax_module.apply().
    batch_stats: A dict of non-trainable model state. Passed as
      {'batch_stats': batch_stats} into flax_module.apply().
    batch: A dictionary with keys 'inputs', 'targets', 'weights'.
    output_activation_fn: An output activation function from jax.nn.functions

  Returns:
    An array of shape [batch_size, num_classes] that contains all the logits.
  """
  variables = {'params': params, 'batch_stats': batch_stats}
  logits = flax_module.apply(variables, batch['inputs'], mutable=False)
  if output_activation_fn:
    return output_activation_fn(logits)
  return logits