def METHOD_NAME(
    state: TrainingState,
    transitions: types.Transition,
) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]:
  loss_and_grad = jax.value_and_grad(
      loss_fn, argnums=1, has_aux=loss_has_aux)
  # Compute losses and their gradients.
  key, key_input = jax.random.split(state.key)
  loss_result, gradients = loss_and_grad(networks, state.policy_params,
                                         key_input, transitions)
  # Combine the gradient across all devices (by taking their mean).
  gradients = jax.lax.pmean(gradients, axis_name=_PMAP_AXIS_NAME)
  # Compute and combine metrics across all devices.
  metrics = _create_loss_metrics(loss_has_aux, loss_result, gradients)
  metrics = jax.lax.pmean(metrics, axis_name=_PMAP_AXIS_NAME)
  policy_update, optimizer_state = optimizer.update(gradients,
                                                    state.optimizer_state,
                                                    state.policy_params)
  policy_params = optax.apply_updates(state.policy_params, policy_update)
  new_state = TrainingState(
      optimizer_state=optimizer_state,
      policy_params=policy_params,
      key=key,
      steps=state.steps + 1,
  )
  return new_state, metrics