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: Creates a ReversibleSerialTrainer and the needed optimizers.
  
  This trainer performs updates equivalent to using the default Trainer on::
  
    tl.Serial(blocks + [loss_layer]).
  
  It is more memory-efficient though since weights are stored on CPU and only
  sent to accelerator layer-by-layer. Blocks are pairs consisting of a list
  of standard (arbitrary) layers and a list of reversible layers which help
  save memory thanks to being reversible.
  
  Args:
    blocks: A list of pairs of lists of standard and reversible layers.
    loss_layer: The final layer of the model; it can have trainable weights
      but should end with a loss: it is required to produce a scalar output.
    optimizer_fn: A function to create the optimizer, e.g., `optimizers.Adam`.
    n_devices: An optional integer, number of accelerator devices to use;
      by default, all available accelerators will be used.
    memoize_jit: Whether to memoize JITed functions; this significantly speeds
      up XLA compilation of larger models, but it uses `repr(layer)` as keys
      to memoize so it could fail if two layers with different functionality
      had the same string representaion. We have not encountered such case
      yet so this is turned on by default, but consider turning it off or
      reviewing your model if you use custom layers and encounter a problem.
    free_accelerators_on_step: If true, frees memory on accelerators when
      starting a step. All layers and arguments must be on host for that,
      otherwise it can lead to failures. Can prevent memory fragmentation.
    adasum: if True, use adaptive summation to gather multi-device gradients.
docstr_tok: ["Creates","a","ReversibleSerialTrainer","and","the","needed","optimizers",".","This","trainer","performs","updates","equivalent","to","using","the","default","Trainer","on",":",":","tl.Serial","(","blocks","+","[","loss_layer","]",")",".","It","is","more","memory-efficient","though","since","weights","are","stored","on","CPU","and","only","sent","to","accelerator","layer-by-layer",".","Blocks","are","pairs","consisting","of","a","list","of","standard","(","arbitrary",")","layers","and","a","list","of","reversible","layers","which","help","save","memory","thanks","to","being","reversible",".","Args",":","blocks",":","A","list","of","pairs","of","lists","of","standard","and","reversible","layers",".","loss_layer",":","The","final","layer","of","the","model",";","it","can","have","trainable","weights","but","should","end","with","a","loss",":","it","is","required","to","produce","a","scalar","output",".","optimizer_fn",":","A","function","to","create","the","optimizer",",","e.g.",",","`","optimizers.Adam","`",".","n_devices",":","An","optional","integer",",","number","of","accelerator","devices","to","use",";","by","default",",","all","available","accelerators","will","be","used",".","memoize_jit",":","Whether","to","memoize","JITed","functions",";","this","significantly","speeds","up","XLA","compilation","of","larger","models",",","but","it","uses","`","repr","(","layer",")","`","as","keys","to","memoize","so","it","could","fail","if","two","layers","with","different","functionality","had","the","same","string","representaion",".","We","have","not","encountered","such","case","yet","so","this","is","turned","on","by","default",",","but","consider","turning","it","off","or","reviewing","your","model","if","you","use","custom","layers","and","encounter","a","problem",".","free_accelerators_on_step",":","If","true",",","frees","memory","on","accelerators","when","starting","a","step",".","All","layers","and","arguments","must","be","on","host","for","that",",","otherwise","it","can","lead","to","failures",".","Can","prevent","memory","fragmentation",".","adasum",":","if","True",",","use","adaptive","summation","to","gather","multi-device","gradients","."]
code: def __init__(
      self,
      blocks,
      loss_layer,
      optimizer_fn,
      n_devices=None,
      memoize_jit=True,
      free_accelerators_on_step=False,
      adasum=False,
  ):
      """Creates a ReversibleSerialTrainer and the needed optimizers.
  
      This trainer performs updates equivalent to using the default Trainer on::
  
        tl.Serial(blocks + [loss_layer]).
  
      It is more memory-efficient though since weights are stored on CPU and only
      sent to accelerator layer-by-layer. Blocks are pairs consisting of a list
      of standard (arbitrary) layers and a list of reversible layers which help
      save memory thanks to being reversible.
  
      Args:
        blocks: A list of pairs of lists of standard and reversible layers.
        loss_layer: The final layer of the model; it can have trainable weights
          but should end with a loss: it is required to produce a scalar output.
        optimizer_fn: A function to create the optimizer, e.g., `optimizers.Adam`.
        n_devices: An optional integer, number of accelerator devices to use;
          by default, all available accelerators will be used.
        memoize_jit: Whether to memoize JITed functions; this significantly speeds
          up XLA compilation of larger models, but it uses `repr(layer)` as keys
          to memoize so it could fail if two layers with different functionality
          had the same string representaion. We have not encountered such case
          yet so this is turned on by default, but consider turning it off or
          reviewing your model if you use custom layers and encounter a problem.
        free_accelerators_on_step: If true, frees memory on accelerators when
          starting a step. All layers and arguments must be on host for that,
          otherwise it can lead to failures. Can prevent memory fragmentation.
        adasum: if True, use adaptive summation to gather multi-device gradients.
      """
      self._blocks = [(tl.Serial(std), rev) for (std, rev) in blocks]
      self._loss_layer = loss_layer
      self._optimizer_fn = optimizer_fn
      self._n_devices = n_devices or fastmath.local_device_count()
      self._adasum = adasum
      self._n_layers = 1 + sum(
          [len(revs) + 1 for (_, revs) in self._blocks]
      )
      self._n_steps_per_log = (
          100  # Log layers and stats every 100 steps.
      )
      self._n_async_layers = 1  # How many layers to run asynchronously.
      self._jit_memory = {} if memoize_jit else None
      self._do_free = free_accelerators_on_step
      self._jit_per_device_rngs = fastmath.jit(
          self._per_device_rngs, backend="cpu"
      )
  
      # Create accelerated versions of layers as pmaped/jited pure_fn.
      self._accelerated_layer_fns = fastmath.nested_map(
          lambda layer: self._pjit(layer.pure_fn, f"fwd {repr(layer)}"),
          self._blocks,
      )
  
      # Create per-layer optimizers and replicate opt_params.
      def _make_optimizer(layer):
          opt = optimizer_fn()
          opt.tree_init(layer.weights)
          opt.slots = tl.on_cpu(opt.slots)
          return opt
  
      self._optimizers = fastmath.nested_map(
          _make_optimizer, self._blocks
      )
      self._replicated_opt_params = fastmath.nested_map(
          lambda opt: self._replicate_cpu(opt.opt_params),
          self._optimizers,
      )
  
      self._loss_opt = _make_optimizer(loss_layer)
      self._replicated_loss_opt_params = self._replicate_cpu(
          self._loss_opt.opt_params
      )
  
      # Forward + backward + optimizer-update functions for all layers.
      # We call them in short FBO for "Forward + Backward + Optimizer update".
      # Reversible layers define a reverse_and_fbo function that also reverses.
  
      self._fbos = []
      for i, (std_layer, rev_layers) in enumerate(self._blocks):
          (std_opt, rev_opts) = self._optimizers[i]
          std_fbo = _fbo_with_layer_and_opt(
              std_layer, std_opt, self._n_devices, adasum=self._adasum
          )
          rev_and_fbos = []
          for layer, opt in zip(rev_layers, rev_opts):
              rev_and_fbo = _reverse_and_fbo_with_layer_and_opt(
                  layer, opt, self._n_devices, self._adasum
              )
              # The donated args are (outputs, weights, grads) and we can donate
              # them because weights and grads are immediately replaced and in
              # case of reversible layers, the outputs are never used again.
              rev_and_fbos.append(
                  self._pjit(
                      rev_and_fbo,
                      f"rev+bwd {repr(layer)}",
                      donate_argnums=(0, 1, 2),
                  )
              )
          # In standard layers, the inputs cannot be donated as they may be used
          # as outputs for the reversible block below, but weights and grads can.
          jit_std_fbo = self._pjit(
              std_fbo, f"bwd {repr(std_layer)}", donate_argnums=(1, 2)
          )
          self._fbos.append((jit_std_fbo, rev_and_fbos))
  
      loss_fbo = _fbo_with_layer_and_opt(
          self._loss_layer,
          self._loss_opt,
          self._n_devices,
          "loss",
          self._adasum,
      )
      self._loss_fbo = self._pjit(loss_fbo, donate_argnums=(1, 2))
  
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 1 to create a new ReversibleSerialTrainer object  with example: obj = ReversibleSerialTrainer(blocks, loss_layer, optimizer_fn, n_devices, memoize_jit, free_accelerators_on_step, adasum)