(layers, loss_chunk_size=0)
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.
Extracts blocks and loss layer for use with ReversibleSerialTrainer.
Extracts blocks and loss layer for use with ReversibleSerialTrainer .
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
def extract_reversible_blocks ( layers , loss_chunk_size = 0 ) : 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