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))
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.
    
trax/optimizers/trainer.py