def one_step(self, batch, rng, step=0, learning_rate=None):
    """Updates layers weights/state and optimizers slots by running one step.

    Args:
      batch: Batch of data to use for optimization.
      rng: Random number generator to use for running this step.
      step: Which step of the training are we running.
      learning_rate: Learning rate to use instead of the default one.

    Returns:
      Tuple (loss, stats) with new values from one step
      of training, where stats are all optimizer statistics.
    """
    # Update the learning rate if needed.
    if learning_rate is not None:
      self._replicated_loss_opt_params['learning_rate'] = self._replicate_cpu(
          learning_rate)
      for (std_op, rev_ops) in self._replicated_opt_params:
        std_op['learning_rate'] = self._replicate_cpu(learning_rate)
        for op in rev_ops:
          op['learning_rate'] = self._replicate_cpu(learning_rate)

    # Batch needs to be split across the local devices -- the difference
    # between _for_n_devices and _reshape_by_device is that the latter splits
    # the batch dim to batch // n_devices, vs _for_n_devices
    # broadcasts/replicates to n_devices dimension.
    step_int = step
    if self._n_devices > 1:
      batch = tl.reshape_by_device(batch, self._n_devices, pure_np=True)
      step = np.repeat(step, self._n_devices)

    # Create separate rng for each device and layer.
    if self._n_devices == 1:
      rngs = fastmath.random.split(rng, self._n_layers)
    else:
      # JIT the function and run it on CPU to avoid memory fragmentation.
      rngs = self._jit_per_device_rngs(tl.on_cpu(rng))
    # Group rngs by layer blocks.
    rng_blocks, rng_i = [], 0
    for _, rev_layers in self._blocks:
      l = len(rev_layers)
      rng_blocks.append((rngs[rng_i], rngs[rng_i + 1: rng_i + l + 1]))
      rng_i += l + 1

    # Run the layers forward upto the loss layer.
    if self._do_free:
      self._free_accelerators()
    process = psutil.Process(os.getpid())
    if isinstance(batch, (list, tuple)):
      batch_shapes = [x.shape for x in batch]
    else:
      batch_shapes = batch.shape
    logging.info('running step %d on shapes %s', step_int, str(batch_shapes))
    if step_int % self._n_steps_per_log == 1:
      logging.info('run fwd: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))

    stack = batch
    block_inputs_states = []
    for i, (std_layer, rev_layers) in enumerate(self._blocks):
      acc_std_layer_fn, acc_rev_layer_fns = self._accelerated_layer_fns[i]
      std_rng, rev_rngs = rng_blocks[i]
      # Run the standard layer.
      stack, std_inputs, std_state = self._run_forward_standard(
          stack, std_layer, acc_std_layer_fn, std_rng, step_int)

      # Run the reversible layers and collect old and new states.
      stack, rev_old_states, rev_new_states = self._run_forward_reversible(
          stack, rev_layers, acc_rev_layer_fns, rev_rngs, step_int)
      block_inputs_states.append(tl.on_cpu(
          ((std_inputs, std_state), (rev_old_states, rev_new_states))))

    # Run the loss layer forward and backward with optimizer update.
    if step_int % self._n_steps_per_log == 1:
      logging.info('run loss: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))
    loss_state = self._replicate(self._loss_layer.state)
    loss_inputs = cb.inputs_from_stack(stack, self._loss_layer.n_in)
    loss_stats, grad_stack = self._run_backward_standard(
        None, step, self._loss_layer, loss_inputs,
        loss_state, self._loss_fbo, rngs[-1], self._loss_opt,
        self._replicated_loss_opt_params)
    self._collect_weights(self._loss_layer)
    stats = [tl.on_cpu(loss_stats)]

    # De-fragment memory.
    if self._do_free:
      stack, grad_stack = tl.on_cpu(stack), tl.on_cpu(grad_stack)
      self._free_accelerators()

    # Run the layers backward and run optimizer updates.
    if step_int % self._n_steps_per_log == 1:
      logging.info('run bwd: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))
    for i in range(len(self._blocks) - 1, -1, -1):
      std_layer, rev_layers = self._blocks[i]
      (std_inputs, std_state), (rev_old_states,
                                rev_new_states) = block_inputs_states[i]
      std_fbo, rev_fbos = self._fbos[i]
      std_opt, rev_opts = self._optimizers[i]
      std_rng, rev_rngs = rng_blocks[i]
      repl_std_opt_params, repl_rev_opts_params = self._replicated_opt_params[i]

      # Run reversible layers backward with optimizer update.
      stack, grad_stack, new_stats = self._run_backward_reversible(
          stack, grad_stack, step, rev_layers, rev_fbos, rev_old_states,
          rev_new_states, rev_rngs, rev_opts, repl_rev_opts_params)
      stats.extend(tl.on_cpu(new_stats))

      # Run the standard layer forward-and-backward pass and optimizer update.
      std_layer_stats, grad_stack = self._run_backward_standard(
          grad_stack, step, std_layer, std_inputs, std_state, std_fbo, std_rng,
          std_opt, repl_std_opt_params)
      stack = cb.outputs_onto_stack(  # Put layer inputs on the stack.
          std_inputs, stack, std_layer.n_out)
      stats.append(tl.on_cpu(std_layer_stats))

      # Collect lazily unreplicated layer weights.
      for rev_layer_id in range(self._n_async_layers):
        self._collect_weights(rev_layers[rev_layer_id])
      self._collect_weights(std_layer)

    # Join stats from different optimizers into one.
    joint_stats = {}
    for i, stat in enumerate(reversed(stats)):
      for k, v in stat.items():
        joint_stats[f'layer{i}/' + k] = v
    return stats[0]['loss'], joint_stats
Updates layers weights/state and optimizers slots by running one step.

    Args:
      batch: Batch of data to use for optimization.
      rng: Random number generator to use for running this step.
      step: Which step of the training are we running.
      learning_rate: Learning rate to use instead of the default one.

    Returns:
      Tuple (loss, stats) with new values from one step
      of training, where stats are all optimizer statistics.
    
trax/optimizers/trainer.py