(self, layer, stack, grad_stack, step, rng,
                                   optimizer, opt_params, reverse_and_fbo,
                                   old_state, new_state)
return stack, grad_stack, layer_stats
Run one reversible layer backwards.
Run one reversible layer backwards.
Run one reversible layer backwards .
def _run_backward_one_reversible(self, layer, stack, grad_stack, step, rng,
                                   optimizer, opt_params, reverse_and_fbo,
                                   old_state, new_state):
    """Run one reversible layer backwards."""
    # We are running backwards and reversing, so we get *outputs* from stack.
    outputs = cb.inputs_from_stack(stack, layer.n_out)
    grads = cb.inputs_from_stack(grad_stack, layer.n_out)
    slots = self._replicate(optimizer.slots)
    weights = self._replicate(layer.weights)  # cpu -> accelerator
    # Ensure all arguments are on accelerator.
    outputs = tl.on_accelerator(outputs)
    grads = tl.on_accelerator(grads)
    old_state = tl.on_accelerator(old_state)
    new_state = tl.on_accelerator(new_state)
    opt_params = tl.on_accelerator(opt_params)
    rng = tl.on_accelerator(rng)
    new_weights, new_slots, inputs, grads, layer_stats = reverse_and_fbo(
        outputs, weights, grads, old_state, new_state,
        slots, opt_params, rng, step)
    layer.weights = self._lazy_unreplicate(new_weights)  # accelerator -> cpu
    layer.state = self._unreplicate(new_state)
    optimizer.slots = self._unreplicate(new_slots)
    stack = cb.outputs_onto_stack(inputs, stack, layer.n_out)
    grad_stack = cb.outputs_onto_stack(grads, grad_stack, layer.n_out)
    return stack, grad_stack, layer_stats
def _run_backward_one_reversible ( self , layer , stack , grad_stack , step , rng , optimizer , opt_params , reverse_and_fbo , old_state , new_state ) : # We are running backwards and reversing, so we get *outputs* from stack. outputs = cb . inputs_from_stack ( stack , layer . n_out ) grads = cb . inputs_from_stack ( grad_stack , layer . n_out ) slots = self . _replicate ( optimizer . slots ) weights = self . _replicate ( layer . weights ) # cpu -> accelerator # Ensure all arguments are on accelerator. outputs = tl . on_accelerator ( outputs ) grads = tl . on_accelerator ( grads ) old_state = tl . on_accelerator ( old_state ) new_state = tl . on_accelerator ( new_state ) opt_params = tl . on_accelerator ( opt_params ) rng = tl . on_accelerator ( rng ) new_weights , new_slots , inputs , grads , layer_stats = reverse_and_fbo ( outputs , weights , grads , old_state , new_state , slots , opt_params , rng , step ) layer . weights = self . _lazy_unreplicate ( new_weights ) # accelerator -> cpu layer . state = self . _unreplicate ( new_state ) optimizer . slots = self . _unreplicate ( new_slots ) stack = cb . outputs_onto_stack ( inputs , stack , layer . n_out ) grad_stack = cb . outputs_onto_stack ( grads , grad_stack , layer . n_out ) return stack , grad_stack , layer_stats