class_docstr: Runs an optimizer on a series of layers, reversible and not.
  
  We provide layers to this trainer in blocks, each block consisting of
  a list of standard layers and a list of reversible layers. They all run
  in turn (like one huge Serial block) but in a more memory-efficient way.
  
  The main motivation for this class is to save memory: it allows to train
  models that have more weights than the memory available on accelerators.
  This happens by caching the weights in CPU memory and transferring only
  the weights of one layer at a time. The reversible layers are used to make
  the backward pass without using additional memory for storing activations.
  
  Note: we do not allow sharing weights between blocks for now.
class_docstr_tok: ["Runs","an","optimizer","on","a","series","of","layers",",","reversible","and","not",".","We","provide","layers","to","this","trainer","in","blocks",",","each","block","consisting","of","a","list","of","standard","layers","and","a","list","of","reversible","layers",".","They","all","run","in","turn","(","like","one","huge","Serial","block",")","but","in","a","more","memory-efficient","way",".","The","main","motivation","for","this","class","is","to","save","memory",":","it","allows","to","train","models","that","have","more","weights","than","the","memory","available","on","accelerators",".","This","happens","by","caching","the","weights","in","CPU","memory","and","transferring","only","the","weights","of","one","layer","at","a","time",".","The","reversible","layers","are","used","to","make","the","backward","pass","without","using","additional","memory","for","storing","activations",".","Note",":","we","do","not","allow","sharing","weights","between","blocks","for","now","."]
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
imports: from concurrent import futures
  import functools
  import os
  import time
  from absl import logging
  import jax
  import numpy
  import psutil
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
label_desc: Use image node_id 15 for calling the ReversibleSerialTrainer obj's underlying member method code with example usage: obj.one_step(batch, rng, step, learning_rate) and returns: stats,  joint_stats