def sync_model_state_across_replicas(state: TrainState) -> TrainState:
  """Sync the model_state (simple averaging) across replicas."""
  # TODO(): We simply do "mean" here and this doesn't work with
  #   statistics like variance. (check the discussion in Flax for fixing this).
  pmap_mean = jax.pmap(lambda x: jax.lax.pmean(x, 'x'), 'x')
  if jax.tree_util.tree_leaves(state.g_model_state):
    g_model_state = pmap_mean(state.g_model_state)
  else:
    g_model_state = state.g_model_state
  if jax.tree_util.tree_leaves(state.d_model_state):
    d_model_state = pmap_mean(state.d_model_state)
  else:
    d_model_state = state.d_model_state
  return state.replace(
      g_model_state=g_model_state,
      d_model_state=d_model_state)
Sync the model_state (simple averaging) across replicas.
videogvt/train_lib/vqgan_train_lib.py