def run_inception_score(params,
                        batch_iter,
                        *,
                        num_samples=10000,
                        num_repeats=4):
  """Run the C3D model with pmap on a batch iterator and get inception score."""
  params = jax_utils.replicate(params)

  def run_fn(params, x, batch_mask):
    out = run_model(params, x)
    out['batch_mask'] = batch_mask
    # Gather outputs to the first device under multinode setting
    out = jax.lax.all_gather(out, axis_name='device')
    return out

  model_fn_p = jax.pmap(run_fn, axis_name='device')
  outputs = []
  sample_count = 0
  num_samples = num_samples * num_repeats
  num_checked = 5
  for batch in batch_iter:
    if num_checked > 0:
      metrics_lib.check_input_range(batch['inputs'])
      num_checked -= 1
    out_p = model_fn_p(params, batch['inputs'], batch['batch_mask'])
    # Get flattened output from the first device
    out = jax.tree_util.tree_map(lambda x: jax.device_get(x[0]), out_p)
    outputs.append(out)
    sample_count += out['batch_mask'].sum()
    if sample_count >= num_samples:
      break
  outputs = metrics_lib.gather_outputs_with_mask(
      outputs, num_samples=num_samples)
  scores = inception_score_from_logits(outputs['logits'], num_repeats)
  return scores
Run the C3D model with pmap on a batch iterator and get inception score.
videogvt/train_lib/inception_score.py