def METHOD_NAME(workload: spec.Workload,
                  current_param_container: spec.ParameterContainer,
                  current_params_types: spec.ParameterTypeTree,
                  model_state: spec.ModelAuxiliaryState,
                  hyperparameters: spec.Hyperparameters,
                  batch: Dict[str, spec.Tensor],
                  loss_type: spec.LossType,
                  optimizer_state: spec.OptimizerState,
                  eval_results: List[Tuple[int, float]],
                  global_step: int,
                  rng: spec.RandomState) -> spec.UpdateReturn:
  """Return (updated_optimizer_state, updated_params)."""
  del current_params_types
  del eval_results
  del loss_type
  lr = get_learning_rate(global_step, hyperparameters)
  optimizer_state, opt_update_fn = optimizer_state
  per_device_rngs = jax.random.split(rng, jax.local_device_count())
  outputs = pmapped_train_step(workload,
                               opt_update_fn,
                               model_state,
                               optimizer_state,
                               current_param_container,
                               hyperparameters,
                               batch,
                               per_device_rngs,
                               lr)
  new_model_state, new_optimizer_state, new_params, loss, grad_norm = outputs
  if global_step <= 1000 or global_step % 100 == 0:
    logging.info('%d) loss = %0.3f, grad_norm = %0.3f lr = %0.6f',
                 global_step,
                 loss.mean(),
                 grad_norm.mean(),
                 lr)
    if workload.metrics_logger is not None:
      workload.metrics_logger.append_scalar_metrics(
          {
              'train_step_ctc_loss': loss.mean(),
              'grad_norm': grad_norm.mean(),
              'learning_rate': lr,
          },
          global_step)
  return (new_optimizer_state, opt_update_fn), new_params, new_model_state