def _build_parallel_train_step():
  """Builds an accelerated version of the train step function."""
  # We need to wrap and unwrap so that the final function can be called with
  # keyword arguments, but we still maintain the proper axes.

  @functools.partial(
      jax.pmap,
      axis_name="devices",
      in_axes=(0, 0, None, None, None, None),
      static_broadcasted_argnums=(2, 3))
  def wrapped(optimizer, batched_examples, static_batch_metadata, loss_fn,
              max_global_norm, optimizer_hyper_params):
    return _parallel_train_step(optimizer, batched_examples,
                                static_batch_metadata, loss_fn, max_global_norm,
                                **optimizer_hyper_params)

  @functools.wraps(_parallel_train_step)
  def wrapper(optimizer, batched_examples, static_batch_metadata, loss_fn,
              max_global_norm, **optimizer_hyper_params):
    return wrapped(optimizer, batched_examples, static_batch_metadata, loss_fn,
                   max_global_norm, optimizer_hyper_params)

  return wrapper