(self, x)
return output, state
Calls this layer for a forward pass, but with custom gradients.
Calls this layer for a forward pass, but with custom gradients.
Calls this layer for a forward pass but with custom gradients .
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
def _do_custom_gradients ( self , x ) : 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