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","."]
code: def _unreplicate(self, x):
      if self._n_devices == 1:
          return tl.on_cpu(x)
      return tl.on_cpu(fastmath.nested_map(lambda x: x[0], x))
  
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 10 for calling the ReversibleSerialTrainer obj's underlying member method code with example usage: obj._unreplicate(x) and returns: tl, tl