def _run_forward_reversible(self, stack, rev_layers, accelerated_fns,
                              rngs, step):
    """Run reversible layers forward, collect states for backwards pass."""
    old_states, new_states = [], []
    for i, layer in enumerate(rev_layers):
      if step % self._n_steps_per_log == 1:
        logging.info('running forward reversible layer %s', str(layer))
      weights = self._replicate(layer.weights)  # also copies cpu -> accelerator
      state = self._replicate(layer.state)
      old_states.append(state)
      inputs = cb.inputs_from_stack(stack, layer.n_in)
      outputs, new_state = accelerated_fns[i](
          inputs, weights, state, rngs[i])
      stack = cb.outputs_onto_stack(outputs, stack, layer.n_in)
      new_states.append(new_state)
    return stack, old_states, new_states
Run reversible layers forward, collect states for backwards pass.
trax/optimizers/trainer.py