text: # -*- coding: utf-8 -*-
  from functools import partial
  import numpy as np
  import torch
  import torch.nn as nn
  import warnings
  from memcnn.models.additive import AdditiveCoupling
  from memcnn.models.affine import AffineCoupling
  from memcnn.models.utils import pytorch_version_one_and_above
  
  
  warnings.filterwarnings(action='ignore', category=UserWarning)
  
  
  def signal_hook(grad_output, valid_states, state_index):  # pragma: no cover
      state = valid_states[state_index]
      valid_states[state_index] = not state
  
  
  def backward_hook(grad_output, keep_input, compute_input_fn, compute_output_fn,
                    input_tensor, output_tensor, valid_states, state_index):  # pragma: no cover
      perform_action = valid_states[state_index]
      valid_states[state_index] = not perform_action
      if perform_action:
          # restore input
          if not keep_input:
              with torch.no_grad():
                  input_inverted = compute_input_fn(output_tensor)
                  if pytorch_version_one_and_above:
                      input_tensor.storage().resize_(int(np.prod(input_tensor.size())))
                      input_tensor.set_(input_inverted)
                  else:
                      input_tensor.set_(input_inverted)
  
          # compute gradients
          with torch.set_grad_enabled(True):
              temp_output = compute_output_fn(input_tensor)
              temp_output.backward(gradient=grad_output)
  
  
  class InvertibleModuleWrapper(nn.Module):
      def __init__(self, fn, keep_input=False, keep_input_inverse=False, disable=False):
          """
          The InvertibleModuleWrapper which enables memory savings during training by exploiting
          the invertible properties of the wrapped module.
  
          Parameters
          ----------
              fn : :obj:`torch.nn.Module`
                  A torch.nn.Module which has a forward and an inverse function implemented with
                  :math:`x == m.inverse(m.forward(x))`
  
              keep_input : :obj:`bool`, optional
                  Set to retain the input information on forward, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              keep_input_inverse : :obj:`bool`, optional
                  Set to retain the input information on inverse, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              disable : :obj:`bool`, optional
                  This will disable the detached graph approach with the backward hook.
                  Essentially this renders the function as `y = fn(x)` without any of the memory savings.
                  Setting this to true will also ignore the keep_input and keep_input_inverse properties.
  
          Attributes
          ----------
              keep_input : :obj:`bool`, optional
                  Set to retain the input information on forward, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              keep_input_inverse : :obj:`bool`, optional
                  Set to retain the input information on inverse, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
          Raises
          ------
          NotImplementedError
              If an unknown coupling or implementation is given.
  
          """
          super(InvertibleModuleWrapper, self).__init__()
          self.disable = disable
          self.keep_input = keep_input
          self.keep_input_inverse = keep_input_inverse
          self._fn = fn
          self._valid_states = []
          self._state_counter = 0
  
  
      def forward(self, xin):
          """Forward operation :math:`R(x) = y`
  
          Parameters
          ----------
              xin : :obj:`torch.Tensor`
                  Input torch tensor.
  
          Returns
          -------
              :obj:`torch.Tensor`
                  Output torch tensor y.
  
          """
          if not self.disable:
              x = xin.detach()  # Makes a detached copy which shares the storage
              y = self._fn(x)
              input_tensor = xin
              output_tensor = y
              # clears the referenced storage data linked to the input tensor as it can be reversed on the backward pass
              if not self.keep_input:
                  if not pytorch_version_one_and_above:
                      # PyTorch 0.4 way to clear storage
                      input_tensor.data.set_()
                  else:
                      # PyTorch 1.0+ way to clear storage
                      input_tensor.storage().resize_(0)
              if self.training:
                  self._valid_states.append(True)
                  xin.register_hook(hook=partial(signal_hook, valid_states=self._valid_states, state_index=self._state_counter))
                  y.register_hook(hook=partial(backward_hook, keep_input=self.keep_input,
                                               compute_input_fn=self._fn.inverse, compute_output_fn=self._fn.forward,
                                               valid_states=self._valid_states, state_index=self._state_counter,
                                               input_tensor=input_tensor, output_tensor=output_tensor))
                  self._state_counter += 1
  
              y.detach_()  # Detaches y in-place (inbetween computations can now be discarded)
              y.requires_grad = self.training
          else:
              y = self._fn(xin)
          return y
  
      def inverse(self, yin):
          """Inverse operation :math:`R^{-1}(y) = x`
  
          Parameters
          ----------
              yin : :obj:`torch.Tensor`
                  Input torch tensor.
  
          Returns
          -------
              :obj:`torch.Tensor`
                  Output torch tensor x.
  
          """
          if not self.disable:
              y = yin.detach()  # Makes a detached copy which shares the storage
              x = self._fn.inverse(y)
              input_tensor = yin
              output_tensor = x
              # clears the referenced storage data linked to the input tensor as it can be reversed on the backward pass
              if not self.keep_input_inverse:
                  if not pytorch_version_one_and_above:
                      # PyTorch 0.4 way to clear storage
                      input_tensor.data.set_()
                  else:
                      # PyTorch 1.0+ way to clear storage
                      input_tensor.storage().resize_(0)
              if self.training:
                  self._valid_states.append(True)
                  yin.register_hook(hook=partial(signal_hook, valid_states=self._valid_states, state_index=self._state_counter))
                  x.register_hook(hook=partial(backward_hook, keep_input=self.keep_input_inverse,
                                               compute_input_fn=self._fn.forward, compute_output_fn=self._fn.inverse,
                                               valid_states=self._valid_states, state_index=self._state_counter,
                                               input_tensor=input_tensor, output_tensor=output_tensor))
                  self._state_counter += 1
              x.detach_()  # Detaches x in-place (inbetween computations can now be discarded)
              x.requires_grad = self.training
          else:
              x = self._fn.inverse(yin)
          return x
  
  
  class ReversibleBlock(InvertibleModuleWrapper):
      def __init__(self, Fm, Gm=None, coupling='additive', keep_input=False, keep_input_inverse=False,
                   implementation_fwd=-1, implementation_bwd=-1, adapter=None):
          """The ReversibleBlock
  
          Warning
          -------
          This class has been deprecated. Use the more flexible InvertibleModuleWrapper class.
  
          Note
          ----
          The `implementation_fwd` and `implementation_bwd` parameters can be set to one of the following implementations:
  
          * -1 Naive implementation without reconstruction on the backward pass.
          * 0  Memory efficient implementation, compute gradients directly.
          * 1  Memory efficient implementation, similar to approach in Gomez et al. 2017.
  
  
          Parameters
          ----------
              Fm : :obj:`torch.nn.Module`
                  A torch.nn.Module encapsulating an arbitrary function
  
              Gm : :obj:`torch.nn.Module`, optional
                  A torch.nn.Module encapsulating an arbitrary function
                  (If not specified a deepcopy of Fm is used as a Module)
  
              coupling : :obj:`str`, optional
                  Type of coupling ['additive', 'affine']. Default = 'additive'
  
              keep_input : :obj:`bool`, optional
                  Set to retain the input information on forward, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              keep_input_inverse : :obj:`bool`, optional
                  Set to retain the input information on inverse, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              implementation_fwd : :obj:`int`, optional
                  Switch between different Operation implementations for forward training (Default = 1).
                  If using the naive implementation (-1) then `keep_input` should be True.
  
              implementation_bwd : :obj:`int`, optional
                  Switch between different Operation implementations for backward training (Default = 1).
                  If using the naive implementation (-1) then `keep_input_inverse` should be True.
  
              adapter : :obj:`class`, optional
                  Only relevant when using the 'affine' coupling.
                  Should be a class of type :obj:`torch.nn.Module` that serves as an
                  optional wrapper class A for Fm and Gm which must output
                  s, t = A(x) with shape(s) = shape(t) = shape(x).
                  s, t are respectively the scale and shift tensors for the affine coupling.
  
          Attributes
          ----------
              keep_input : :obj:`bool`, optional
                  Set to retain the input information on forward, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
              keep_input_inverse : :obj:`bool`, optional
                  Set to retain the input information on inverse, by default it can be discarded since it will be
                  reconstructed upon the backward pass.
  
          Raises
          ------
          NotImplementedError
              If an unknown coupling or implementation is given.
  
          """
          warnings.warn("This class has been deprecated. Use the more flexible InvertibleModuleWrapper class", DeprecationWarning)
          fn = create_coupling(Fm=Fm, Gm=Gm, coupling=coupling,
                               implementation_fwd=implementation_fwd, implementation_bwd=implementation_bwd,
                               adapter=adapter)
          super(ReversibleBlock, self).__init__(fn, keep_input=keep_input, keep_input_inverse=keep_input_inverse)
  
  
  def create_coupling(Fm, Gm=None, coupling='additive', implementation_fwd=-1, implementation_bwd=-1, adapter=None):
      if coupling == 'additive':
          fn = AdditiveCoupling(Fm, Gm,
                                implementation_fwd=implementation_fwd, implementation_bwd=implementation_bwd)
      elif coupling == 'affine':
          fn = AffineCoupling(Fm, Gm, adapter=adapter,
                              implementation_fwd=implementation_fwd, implementation_bwd=implementation_bwd)
      else:
          raise NotImplementedError('Unknown coupling method: %s' % coupling)
      return fn
  
  
  def is_invertible_module(module_in, test_input_shape, test_input_dtype=torch.float32, atol=1e-6):
      """Test if a :obj:`torch.nn.Module` is invertible
  
      Parameters
      ----------
      module_in : :obj:`torch.nn.Module`
          A torch.nn.Module to test.
      test_input_shape : :obj:`tuple`
          Dimensions of test tensor object to perform the test with.
      test_input_dtype : :obj:`torch.dtype`, optional
          Data type of test tensor object to perform the test with.
      atol : :obj:`float`, optional
          Tolerance value used for comparing the outputs
  
      Returns
      -------
          :obj:`bool`
              True if the input module is invertible, False otherwise.
  
      """
      test_input = torch.rand(test_input_shape, dtype=test_input_dtype)
      if not hasattr(module_in, "inverse"):
          return False
      with torch.no_grad():
          if not torch.allclose(module_in.inverse(module_in(test_input)), test_input, atol=atol):
              return False
          if not torch.allclose(module_in(module_in.inverse(test_input)), test_input, atol=atol):
              return False
          if test_input is module_in(test_input):
              return False
          if test_input is module_in.inverse(test_input):
              return False
      return True
  