def _pure_forward(self, x, weights, state, rng):
    """Call self.forward in a pure way."""
    old_weights, old_state, old_rng = self.weights, self.state, self._rng
    self.weights, self.state, self._rng = weights, state, rng
    res = self.forward(x)
    self.weights, self.state, self._rng = old_weights, old_state, old_rng
    return res
Call self.forward in a pure way.
trax/layers/reversible.py