def all_layers(self):
    """Returns all layers that compose the model and loss in this class."""
    layers = []
    for (std_layer, rev_layers) in self._blocks:
      layers.append(std_layer)
      layers.extend(rev_layers)
    layers.append(self._loss_layer)
    return layers
Returns all layers that compose the model and loss in this class.
trax/optimizers/trainer.py