def extract_reversible_blocks(layers, loss_chunk_size=0):
  """Extracts blocks and loss layer for use with ReversibleSerialTrainer.

  Args:
    layers: a list of layers of a single layer to extract blocks from;
      should end with a loss, e.g., [model, loss] or tl.Serial(model, loss).
    loss_chunk_size: int, if > 0 creates a chunked loss layer to save memory
      in models with larger vocabulary; requires the last sublayers of loss
      are [Dense, LogSoftmax, _CrossEntropy, _WeightedMean] in that order.

  Returns:
    a pair (blocks, loss_layer) to use with ReversibleSerialTrainer.
  """
  def _flatten(l):
    """Flatten all Serial layers and sub(sub-...) layers into a list."""
    if isinstance(l, (list, tuple)):
      return [x for layer in l for x in _flatten(layer)]  # pylint: disable=g-complex-comprehension
    elif isinstance(l, tl.Serial):
      return _flatten(l.sublayers)
    else:
      return [l]

  # Extract standard and reversible layer blocks.
  blocks, std_layers, rev_layers = [], [], []
  for layer in _flatten(layers):
    if isinstance(layer, tl.ReversibleLayer):
      rev_layers.append(layer)
    elif not rev_layers:
      std_layers.append(layer)
    else:
      blocks.append((std_layers, rev_layers))
      std_layers, rev_layers = [], []
      std_layers.append(layer)
  if rev_layers:
    raise ValueError('The final layer must be a standard loss, not reversible.')
  if loss_chunk_size > 0:
    # For now we only do chunking of [Dense, LogSoftmax, CrossEntopy, Mean]
    # Let's check that these are the last 4 layers.
    border_layers = ['StripFromConcatenateWithPadding', 'Select']

    loss_start = None
    for index, layer in enumerate(std_layers):
      if layer.name in border_layers:
        loss_start = index + 1
    if loss_start is None:
      raise ValueError('Loss layer should be preceeded by one of {}; got {}'
                       .format(border_layers, [l.name for l in std_layers]))
    if len(std_layers) - loss_start < 4:
      raise ValueError('Too short loss layer for chunking')
    last_3_names = ' '.join([l.name for l in std_layers[-3:]])
    if last_3_names != 'LogSoftmax _CrossEntropy _WeightedMean':
      raise ValueError('Loss chunking only works with last layers being "'
                       'LogSoftmax, _CrossEntropy, _WeightedMean" but got: ' +
                       last_3_names)

    # Create chunked dense+logsoftmax+cross-entropy-loss.
    chunked_xent = tl.Chunk(tl.Serial(std_layers[loss_start:-1]),
                            loss_chunk_size)
    # The chunked loss should operate on a merged batch dimension, e.g.,
    # including both length and batch size. Need to merge and un-merge later.
    def _reshape_to_batch_and_copy_targets(preds, targets):
      batched_preds = jnp.reshape(preds, [-1, preds.shape[-1]])
      batched_targets = jnp.reshape(targets, [-1])
      return batched_preds, batched_targets, targets
    def _reshape_xent_back(xent, targets):
      return jnp.reshape(xent, targets.shape)
    batched_xent = tl.Serial(
        tl.Fn('pre_xent_rebatch', _reshape_to_batch_and_copy_targets, n_out=3),
        chunked_xent,
        tl.Fn('after_xent_rebatch', _reshape_xent_back)
    )
    loss_layer = tl.Serial(std_layers[:loss_start] + [batched_xent],
                           std_layers[-1])
  else:
    loss_layer = tl.Serial(std_layers)
  return blocks, loss_layer
Extracts blocks and loss layer for use with ReversibleSerialTrainer.

  Args:
    layers: a list of layers of a single layer to extract blocks from;
      should end with a loss, e.g., [model, loss] or tl.Serial(model, loss).
    loss_chunk_size: int, if > 0 creates a chunked loss layer to save memory
      in models with larger vocabulary; requires the last sublayers of loss
      are [Dense, LogSoftmax, _CrossEntropy, _WeightedMean] in that order.

  Returns:
    a pair (blocks, loss_layer) to use with ReversibleSerialTrainer.
  
trax/optimizers/trainer.py