content: # -*- coding: utf-8 -*-
  # @Author: Peng Xiang
  
  import types
  import torch
  import torch.nn.functional as F
  import numpy as np
  from torch import nn, einsum
  from pointnet2_ops.pointnet2_utils import furthest_point_sample, \
      gather_operation, ball_query, three_nn, three_interpolate, grouping_operation
  
  class Conv1d(nn.Module):
      def __init__(self, in_channel, out_channel, kernel_size=1, stride=1, if_bn=True, activation_fn=torch.relu):
          super(Conv1d, self).__init__()
          self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=stride)
          self.if_bn = if_bn
          self.bn = nn.BatchNorm1d(out_channel)
          self.activation_fn = activation_fn
  
      def forward(self, input):
          out = self.conv(input)
          if self.if_bn:
              out = self.bn(out)
  
          if self.activation_fn is not None:
              out = self.activation_fn(out)
  
          return out
  
  class Conv2d(nn.Module):
      def __init__(self, in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1), if_bn=True, activation_fn=torch.relu):
          super(Conv2d, self).__init__()
          self.conv = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride)
          self.if_bn = if_bn
          self.bn = nn.BatchNorm2d(out_channel)
          self.activation_fn = activation_fn
  
      def forward(self, input):
          out = self.conv(input)
          if self.if_bn:
              out = self.bn(out)
  
          if self.activation_fn is not None:
              out = self.activation_fn(out)
  
          return out
  
  class MLP(nn.Module):
      def __init__(self, in_channel, layer_dims, bn=None):
          super(MLP, self).__init__()
          layers = []
          last_channel = in_channel
          for out_channel in layer_dims[:-1]:
              layers.append(nn.Linear(last_channel, out_channel))
              if bn:
                  layers.append(nn.BatchNorm1d(out_channel))
              layers.append(nn.ReLU())
              last_channel = out_channel
          layers.append(nn.Linear(last_channel, layer_dims[-1]))
          self.mlp = nn.Sequential(*layers)
  
      def forward(self, inputs):
          return self.mlp(inputs)
  
  class MLP_CONV(nn.Module):
      def __init__(self, in_channel, layer_dims, bn=None):
          super(MLP_CONV, self).__init__()
          layers = []
          last_channel = in_channel
          for out_channel in layer_dims[:-1]:
              layers.append(nn.Conv1d(last_channel, out_channel, 1))
              if bn:
                  layers.append(nn.BatchNorm1d(out_channel))
              layers.append(nn.ReLU())
              last_channel = out_channel
          layers.append(nn.Conv1d(last_channel, layer_dims[-1], 1))
          self.mlp = nn.Sequential(*layers)
  
      def forward(self, inputs):
          return self.mlp(inputs)
  
  class MLP_Res(nn.Module):
      def __init__(self, in_dim=128, hidden_dim=None, out_dim=128):
          super(MLP_Res, self).__init__()
          if hidden_dim is None:
              hidden_dim = in_dim
          self.conv_1 = nn.Conv1d(in_dim, hidden_dim, 1)
          self.conv_2 = nn.Conv1d(hidden_dim, out_dim, 1)
          self.conv_shortcut = nn.Conv1d(in_dim, out_dim, 1)
  
      def forward(self, x):
          """
          Args:
              x: (B, out_dim, n)
          """
          shortcut = self.conv_shortcut(x)
          out = self.conv_2(torch.relu(self.conv_1(x))) + shortcut
          return out
  
  
  def sample_and_group(xyz, points, npoint, nsample, radius, use_xyz=True):
      """
      Args:
          xyz: Tensor, (B, 3, N)
          points: Tensor, (B, f, N)
          npoint: int
          nsample: int
          radius: float
          use_xyz: boolean
  
      Returns:
          new_xyz: Tensor, (B, 3, npoint)
          new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
          idx_local: Tensor, (B, npoint, nsample)
          grouped_xyz: Tensor, (B, 3, npoint, nsample)
  
      """
      xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
      new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
  
      idx = ball_query(radius, nsample, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) # (B, npoint, nsample)
      grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
      grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, nsample)
  
      if points is not None:
          grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
          if use_xyz:
              new_points = torch.cat([grouped_xyz, grouped_points], 1)
          else:
              new_points = grouped_points
      else:
          new_points = grouped_xyz
  
      return new_xyz, new_points, idx, grouped_xyz
  
  
  def sample_and_group_all(xyz, points, use_xyz=True):
      """
      Args:
          xyz: Tensor, (B, 3, nsample)
          points: Tensor, (B, f, nsample)
          use_xyz: boolean
  
      Returns:
          new_xyz: Tensor, (B, 3, 1)
          new_points: Tensor, (B, f|f+3|3, 1, nsample)
          idx: Tensor, (B, 1, nsample)
          grouped_xyz: Tensor, (B, 3, 1, nsample)
      """
      b, _, nsample = xyz.shape
      device = xyz.device
      new_xyz = torch.zeros((1, 3, 1), dtype=torch.float, device=device).repeat(b, 1, 1)
      grouped_xyz = xyz.reshape((b, 3, 1, nsample))
      idx = torch.arange(nsample, device=device).reshape(1, 1, nsample).repeat(b, 1, 1)
      if points is not None:
          if use_xyz:
              new_points = torch.cat([xyz, points], 1)
          else:
              new_points = points
          new_points = new_points.unsqueeze(2)
      else:
          new_points = grouped_xyz
  
      return new_xyz, new_points, idx, grouped_xyz
  
  
  class PointNet_SA_Module(nn.Module):
      def __init__(self, npoint, nsample, radius, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True):
          """
          Args:
              npoint: int, number of points to sample
              nsample: int, number of points in each local region
              radius: float
              in_channel: int, input channel of features(points)
              mlp: list of int,
          """
          super(PointNet_SA_Module, self).__init__()
          self.npoint = npoint
          self.nsample = nsample
          self.radius = radius
          self.mlp = mlp
          self.group_all = group_all
          self.use_xyz = use_xyz
          if use_xyz:
              in_channel += 3
  
          last_channel = in_channel
          self.mlp_conv = []
          for out_channel in mlp:
              self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
              last_channel = out_channel
  
          self.mlp_conv = nn.Sequential(*self.mlp_conv)
  
      def forward(self, xyz, points):
          """
          Args:
              xyz: Tensor, (B, 3, N)
              points: Tensor, (B, f, N)
  
          Returns:
              new_xyz: Tensor, (B, 3, npoint)
              new_points: Tensor, (B, mlp[-1], npoint)
          """
          if self.group_all:
              new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
          else:
              new_xyz, new_points, idx, grouped_xyz = sample_and_group(xyz, points, self.npoint, self.nsample, self.radius, self.use_xyz)
  
          new_points = self.mlp_conv(new_points)
          new_points = torch.max(new_points, 3)[0]
  
          return new_xyz, new_points
  
  
  class PointNet_FP_Module(nn.Module):
      def __init__(self, in_channel, mlp, use_points1=False, in_channel_points1=None, if_bn=True):
          """
          Args:
              in_channel: int, input channel of points2
              mlp: list of int
              use_points1: boolean, if use points
              in_channel_points1: int, input channel of points1
          """
          super(PointNet_FP_Module, self).__init__()
          self.use_points1 = use_points1
  
          if use_points1:
              in_channel += in_channel_points1
  
          last_channel = in_channel
          self.mlp_conv = []
          for out_channel in mlp:
              self.mlp_conv.append(Conv1d(last_channel, out_channel, if_bn=if_bn))
              last_channel = out_channel
  
          self.mlp_conv = nn.Sequential(*self.mlp_conv)
  
      def forward(self, xyz1, xyz2, points1, points2):
          """
          Args:
              xyz1: Tensor, (B, 3, N)
              xyz2: Tensor, (B, 3, M)
              points1: Tensor, (B, in_channel, N)
              points2: Tensor, (B, in_channel, M)
  
          Returns:MLP_CONV
              new_points: Tensor, (B, mlp[-1], N)
          """
          dist, idx = three_nn(xyz1.permute(0, 2, 1).contiguous(), xyz2.permute(0, 2, 1).contiguous())
          dist = torch.clamp_min(dist, 1e-10)  # (B, N, 3)
          recip_dist = 1.0/dist
          norm = torch.sum(recip_dist, 2, keepdim=True).repeat((1, 1, 3))
          weight = recip_dist / norm
          interpolated_points = three_interpolate(points2, idx, weight) # B, in_channel, N
  
          if self.use_points1:
              new_points = torch.cat([interpolated_points, points1], 1)
          else:
              new_points = interpolated_points
  
          new_points = self.mlp_conv(new_points)
          return new_points
  
  
  def square_distance(src, dst):
      """
      Calculate Euclid distance between each two points.
  
      src^T * dst = xn * xm + yn * ym + zn * zm；
      sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
      sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
      dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
           = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
  
      Input:
          src: source points, [B, N, C]
          dst: target points, [B, M, C]
      Output:
          dist: per-point square distance, [B, N, M]
      """
      B, N, _ = src.shape
      _, M, _ = dst.shape
      dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))  # B, N, M
      dist += torch.sum(src ** 2, -1).view(B, N, 1)
      dist += torch.sum(dst ** 2, -1).view(B, 1, M)
      return dist
  
  
  def query_knn(nsample, xyz, new_xyz, include_self=True):
      """Find k-NN of new_xyz in xyz"""
      pad = 0 if include_self else 1
      sqrdists = square_distance(new_xyz, xyz)  # B, S, N
      idx = torch.argsort(sqrdists, dim=-1, descending=False)[:, :, pad: nsample+pad]
      return idx.int()
  
  
  def sample_and_group_knn(xyz, points, npoint, k, use_xyz=True, idx=None):
      """
      Args:
          xyz: Tensor, (B, 3, N)
          points: Tensor, (B, f, N)
          npoint: int
          nsample: int
          radius: float
          use_xyz: boolean
  
      Returns:
          new_xyz: Tensor, (B, 3, npoint)
          new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
          idx_local: Tensor, (B, npoint, nsample)
          grouped_xyz: Tensor, (B, 3, npoint, nsample)
  
      """
      xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
      new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
      if idx is None:
          idx = query_knn(k, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous())
      grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
      grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, k)
  
      if points is not None:
          grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
          if use_xyz:
              new_points = torch.cat([grouped_xyz, grouped_points], 1)
          else:
              new_points = grouped_points
      else:
          new_points = grouped_xyz
  
      return new_xyz, new_points, idx, grouped_xyz
  
  
  class PointNet_SA_Module_KNN(nn.Module):
      def __init__(self, npoint, nsample, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True, if_idx=False):
          """
          Args:
              npoint: int, number of points to sample
              nsample: int, number of points in each local region
              radius: float
              in_channel: int, input channel of features(points)
              mlp: list of int,
          """
          super(PointNet_SA_Module_KNN, self).__init__()
          self.npoint = npoint
          self.nsample = nsample
          self.mlp = mlp
          self.group_all = group_all
          self.use_xyz = use_xyz
          self.if_idx = if_idx
          if use_xyz:
              in_channel += 3
  
          last_channel = in_channel
          self.mlp_conv = []
          for out_channel in mlp[:-1]:
              self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
              last_channel = out_channel
          self.mlp_conv.append(Conv2d(last_channel, mlp[-1], if_bn=False, activation_fn=None))
          self.mlp_conv = nn.Sequential(*self.mlp_conv)
  
      def forward(self, xyz, points, idx=None):
          """
          Args:
              xyz: Tensor, (B, 3, N)
              points: Tensor, (B, f, N)
  
          Returns:
              new_xyz: Tensor, (B, 3, npoint)
              new_points: Tensor, (B, mlp[-1], npoint)
          """
          if self.group_all:
              new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
          else:
              new_xyz, new_points, idx, grouped_xyz = sample_and_group_knn(xyz, points, self.npoint, self.nsample, self.use_xyz, idx=idx)
  
          new_points = self.mlp_conv(new_points)
          new_points = torch.max(new_points, 3)[0]
  
          if self.if_idx:
              return new_xyz, new_points, idx
          else:
              return new_xyz, new_points
  
  
  def fps_subsample(pcd, n_points=2048):
      """
      Args
          pcd: (b, 16384, 3)
  
      returns
          new_pcd: (b, n_points, 3)
      """
      new_pcd = gather_operation(pcd.permute(0, 2, 1).contiguous(), furthest_point_sample(pcd, n_points))
      new_pcd = new_pcd.permute(0, 2, 1).contiguous()
      return new_pcd
  
  
  class Transformer(nn.Module):
      def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4):
          super(Transformer, self).__init__()
          self.n_knn = n_knn
          self.conv_key = nn.Conv1d(dim, dim, 1)
          self.conv_query = nn.Conv1d(dim, dim, 1)
          self.conv_value = nn.Conv1d(dim, dim, 1)
  
          self.pos_mlp = nn.Sequential(
              nn.Conv2d(3, pos_hidden_dim, 1),
              nn.BatchNorm2d(pos_hidden_dim),
              nn.ReLU(),
              nn.Conv2d(pos_hidden_dim, dim, 1)
          )
  
          self.attn_mlp = nn.Sequential(
              nn.Conv2d(dim, dim * attn_hidden_multiplier, 1),
              nn.BatchNorm2d(dim * attn_hidden_multiplier),
              nn.ReLU(),
              nn.Conv2d(dim * attn_hidden_multiplier, dim, 1)
          )
  
          self.linear_start = nn.Conv1d(in_channel, dim, 1)
          self.linear_end = nn.Conv1d(dim, in_channel, 1)
  
      def forward(self, x, pos):
          """feed forward of transformer
          Args:
              x: Tensor of features, (B, in_channel, n)
              pos: Tensor of positions, (B, 3, n)
  
          Returns:
              y: Tensor of features with attention, (B, in_channel, n)
          """
  
          identity = x
  
          x = self.linear_start(x)
          b, dim, n = x.shape
  
          pos_flipped = pos.permute(0, 2, 1).contiguous()
          idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped)
          key = self.conv_key(x)
          value = self.conv_value(x)
          query = self.conv_query(x)
  
          key = grouping_operation(key, idx_knn)  # b, dim, n, n_knn
          qk_rel = query.reshape((b, -1, n, 1)) - key
  
          pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn)  # b, 3, n, n_knn
          pos_embedding = self.pos_mlp(pos_rel)  # b, dim, n, n_knn
  
          attention = self.attn_mlp(qk_rel + pos_embedding)
          attention = torch.softmax(attention, -1)
  
          value = value.reshape((b, -1, n, 1)) + pos_embedding
  
          agg = einsum('b c i j, b c i j -> b c i', attention, value)  # b, dim, n
          y = self.linear_end(agg)
  
          return y+identity
  
  
  class CouplingLayer(nn.Module):
  
      def __init__(self, d, intermediate_dim, swap=False):
          nn.Module.__init__(self)
          self.d = d - (d // 2)
          self.swap = swap
          self.net_s_t = nn.Sequential(
              nn.Linear(self.d, intermediate_dim),
              nn.ReLU(inplace=True),
              nn.Linear(intermediate_dim, intermediate_dim),
              nn.ReLU(inplace=True),
              nn.Linear(intermediate_dim, (d - self.d) * 2),
          )
  
      def forward(self, x, logpx=None, reverse=False):
  
          if self.swap:
              x = torch.cat([x[:, self.d:], x[:, :self.d]], 1)
  
          in_dim = self.d
          out_dim = x.shape[1] - self.d
  
          s_t = self.net_s_t(x[:, :in_dim])
          scale = torch.sigmoid(s_t[:, :out_dim] + 2.)
          shift = s_t[:, out_dim:]
  
          logdetjac = torch.sum(torch.log(scale).view(scale.shape[0], -1), 1, keepdim=True)
  
          if not reverse:
              y1 = x[:, self.d:] * scale + shift
              delta_logp = -logdetjac
          else:
              y1 = (x[:, self.d:] - shift) / scale
              delta_logp = logdetjac
  
          y = torch.cat([x[:, :self.d], y1], 1) if not self.swap else torch.cat([y1, x[:, :self.d]], 1)
  
          if logpx is None:
              return y
          else:
              return y, logpx + delta_logp
  
  
  class SequentialFlow(nn.Module):
      """A generalized nn.Sequential container for normalizing flows.
      """
  
      def __init__(self, layersList):
          super(SequentialFlow, self).__init__()
          self.chain = nn.ModuleList(layersList)
  
      def forward(self, x, logpx=None, reverse=False, inds=None):
          if inds is None:
              if reverse:
                  inds = range(len(self.chain) - 1, -1, -1)
              else:
                  inds = range(len(self.chain))
  
          if logpx is None:
              for i in inds:
                  x = self.chain[i](x, reverse=reverse)
              return x
          else:
              for i in inds:
                  x, logpx = self.chain[i](x, logpx, reverse=reverse)
              return x, logpx
  
  
  def build_latent_flow(args):
      chain = []
      for i in range(args.latent_flow_depth):
          chain.append(CouplingLayer(args.latent_dim, args.latent_flow_hidden_dim, swap=(i % 2 == 0)))
      return SequentialFlow(chain)
  
  
  ##################
  ## SpectralNorm ##
  ##################
  
  POWER_ITERATION_FN = "spectral_norm_power_iteration"
  
  
  class SpectralNorm(object):
      def __init__(self, name='weight', dim=0, eps=1e-12):
          self.name = name
          self.dim = dim
          self.eps = eps
  
      def compute_weight(self, module, n_power_iterations):
          if n_power_iterations < 0:
              raise ValueError(
                  'Expected n_power_iterations to be non-negative, but '
                  'got n_power_iterations={}'.format(n_power_iterations)
              )
  
          weight = getattr(module, self.name + '_orig')
          u = getattr(module, self.name + '_u')
          v = getattr(module, self.name + '_v')
          weight_mat = weight
          if self.dim != 0:
              # permute dim to front
              weight_mat = weight_mat.permute(self.dim, * [d for d in range(weight_mat.dim()) if d != self.dim])
          height = weight_mat.size(0)
          weight_mat = weight_mat.reshape(height, -1)
          with torch.no_grad():
              for _ in range(n_power_iterations):
                  # Spectral norm of weight equals to `u^T W v`, where `u` and `v`
                  # are the first left and right singular vectors.
                  # This power iteration produces approximations of `u` and `v`.
                  v = F.normalize(torch.matmul(weight_mat.t(), u), dim=0, eps=self.eps)
                  u = F.normalize(torch.matmul(weight_mat, v), dim=0, eps=self.eps)
          setattr(module, self.name + '_u', u)
          setattr(module, self.name + '_v', v)
  
          sigma = torch.dot(u, torch.matmul(weight_mat, v))
          weight = weight / sigma
          setattr(module, self.name, weight)
  
      def remove(self, module):
          weight = getattr(module, self.name)
          delattr(module, self.name)
          delattr(module, self.name + '_u')
          delattr(module, self.name + '_orig')
          module.register_parameter(self.name, torch.nn.Parameter(weight))
  
      def get_update_method(self, module):
          def update_fn(module, n_power_iterations):
              self.compute_weight(module, n_power_iterations)
  
          return update_fn
  
      def __call__(self, module, unused_inputs):
          del unused_inputs
          self.compute_weight(module, n_power_iterations=0)
  
          # requires_grad might be either True or False during inference.
          if not module.training:
              r_g = getattr(module, self.name + '_orig').requires_grad
              setattr(module, self.name, getattr(module, self.name).detach().requires_grad_(r_g))
  
      @staticmethod
      def apply(module, name, dim, eps):
          fn = SpectralNorm(name, dim, eps)
          weight = module._parameters[name]
          height = weight.size(dim)
  
          u = F.normalize(weight.new_empty(height).normal_(0, 1), dim=0, eps=fn.eps)
          v = F.normalize(weight.new_empty(int(weight.numel() / height)).normal_(0, 1), dim=0, eps=fn.eps)
          delattr(module, fn.name)
          module.register_parameter(fn.name + "_orig", weight)
          # We still need to assign weight back as fn.name because all sorts of
          # things may assume that it exists, e.g., when initializing weights.
          # However, we can't directly assign as it could be an nn.Parameter and
          # gets added as a parameter. Instead, we register weight.data as a
          # buffer, which will cause weight to be included in the state dict
          # and also supports nn.init due to shared storage.
          module.register_buffer(fn.name, weight.data)
          module.register_buffer(fn.name + "_u", u)
          module.register_buffer(fn.name + "_v", v)
  
          setattr(module, POWER_ITERATION_FN, types.MethodType(fn.get_update_method(module), module))
  
          module.register_forward_pre_hook(fn)
          return fn
  
  
  def inplace_spectral_norm(module, name='weight', dim=None, eps=1e-12):
      r"""Applies spectral normalization to a parameter in the given module.
      .. math::
           \mathbf{W} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})} \\
           \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2}
      Spectral normalization stabilizes the training of discriminators (critics)
      in Generaive Adversarial Networks (GANs) by rescaling the weight tensor
      with spectral norm :math:`\sigma` of the weight matrix calculated using
      power iteration method. If the dimension of the weight tensor is greater
      than 2, it is reshaped to 2D in power iteration method to get spectral
      norm. This is implemented via a hook that calculates spectral norm and
      rescales weight before every :meth:`~Module.forward` call.
      See `Spectral Normalization for Generative Adversarial Networks`_ .
      .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957
      Args:
          module (nn.Module): containing module
          name (str, optional): name of weight parameter
          n_power_iterations (int, optional): number of power iterations to
              calculate spectal norm
          dim (int, optional): dimension corresponding to number of outputs,
              the default is 0, except for modules that are instances of
              ConvTranspose1/2/3d, when it is 1
          eps (float, optional): epsilon for numerical stability in
              calculating norms
      Returns:
          The original module with the spectal norm hook
      Example::
          >>> m = spectral_norm(nn.Linear(20, 40))
          Linear (20 -> 40)
          >>> m.weight_u.size()
          torch.Size([20])
      """
      if dim is None:
          if isinstance(module, (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):
              dim = 1
          else:
              dim = 0
      SpectralNorm.apply(module, name, dim=dim, eps=eps)
      return module
  
  
  def remove_spectral_norm(module, name='weight'):
      r"""Removes the spectral normalization reparameterization from a module.
      Args:
          module (nn.Module): containing module
          name (str, optional): name of weight parameter
      Example:
          >>> m = spectral_norm(nn.Linear(40, 10))
          >>> remove_spectral_norm(m)
      """
      for k, hook in module._forward_pre_hooks.items():
          if isinstance(hook, SpectralNorm) and hook.name == name:
              hook.remove(module)
              del module._forward_pre_hooks[k]
              return module
  
      raise ValueError("spectral_norm of '{}' not found in {}".format(name, module))
  
  
  def add_spectral_norm(model, logger=None):
      """Applies spectral norm to all modules within the scope of a CNF."""
  
      def apply_spectral_norm(module):
          if 'weight' in module._parameters:
              if logger: logger.info("Adding spectral norm to {}".format(module))
              inplace_spectral_norm(module, 'weight')
  
      def find_coupling_layer(module):
          if isinstance(module, CouplingLayer):
              module.apply(apply_spectral_norm)
          else:
              for child in module.children():
                  find_coupling_layer(child)
  
      find_coupling_layer(model)
  
  
  def spectral_norm_power_iteration(model, n_power_iterations=1):
  
      def recursive_power_iteration(module):
          if hasattr(module, POWER_ITERATION_FN):
              getattr(module, POWER_ITERATION_FN)(n_power_iterations)
  
      model.apply(recursive_power_iteration)
  
  def reparameterize_gaussian(mean, logvar):
      std = torch.exp(0.5 * logvar)
      eps = torch.randn(std.size()).to(mean)
      return mean + std * eps
  
  
  
  def gaussian_entropy(logvar):
      const = 0.5 * float(logvar.size(1)) * (1. + np.log(np.pi * 2))
      ent = 0.5 * logvar.sum(dim=1, keepdim=False) + const
      return ent
  
  
  def standard_normal_logprob(z):
      dim = z.size(-1)
      log_z = -0.5 * dim * np.log(2 * np.pi)
      return log_z - z.pow(2) / 2
  
  def truncated_normal_(tensor, mean=0, std=1, trunc_std=2):
      """
      Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
      """
      size = tensor.shape
      tmp = tensor.new_empty(size + (4,)).normal_()
      valid = (tmp < trunc_std) & (tmp > -trunc_std)
      ind = valid.max(-1, keepdim=True)[1]
      tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
      tensor.data.mul_(std).add_(mean)
      return tensor