def _do_custom_gradients(self, x):
    """Calls this layer for a forward pass, but with custom gradients."""

    def _f(state, rng, y, weights):
      old_weights, old_state, old_rng = self.weights, self.state, self._rng
      self.weights, self.state, self._rng = weights, state, rng
      res = self.forward(y)
      s = self.state
      self.weights, self.state, self._rng = old_weights, old_state, old_rng
      return res, s

    def _f_fwd(state, rng, y, weights):
      old_weights, old_state, old_rng = self.weights, self.state, self._rng
      self.weights, self.state, self._rng = weights, state, rng
      res = self.forward(y)
      s = self.state
      self.weights, self.state, self._rng = old_weights, old_state, old_rng
      return (res, s), (state, rng, y, res, weights, s)

    def _f_bwd(residual, grad):
      """Custom gradient function."""
      state, rng, y, output, weights, new_state = residual
      grad = grad[0]  # Ignore dummy gradient wrt state.
      out = self.backward(y, output, grad, weights, state, new_state, rng)
      return (None, None, *out)

    do_forward = fastmath.custom_vjp(_f, _f_fwd, _f_bwd, nondiff_argnums=(0, 1))

    output, state = do_forward(self.state, self._rng, x, self.weights)
    return output, state
Calls this layer for a forward pass, but with custom gradients.
trax/layers/base.py