text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class Flip(nn.Module):
      def forward(self, x, sldj, reverse=False, condition_embd=None):
          assert isinstance(x, tuple) and len(x) == 2
          return (x[1], x[0]), sldj
  
  
  def mean_dim(tensor, dim=None, keepdims=False):
      """Take the mean along multiple dimensions.
  
      Args:
          tensor (torch.Tensor): Tensor of values to average.
          dim (list): List of dimensions along which to take the mean.
          keepdims (bool): Keep dimensions rather than squeezing.
  
      Returns:
          mean (torch.Tensor): New tensor of mean value(s).
      """
      if dim is None:
          return tensor.mean()
      else:
          if isinstance(dim, int):
              dim = [dim]
          dim = sorted(dim)
          for d in dim:
              tensor = tensor.mean(dim=d, keepdim=True)
          if not keepdims:
              for i, d in enumerate(dim):
                  tensor.squeeze_(d-i)
          return tensor
  
  
  def checkerboard(x, reverse=False):
      """Split x in a checkerboard pattern. Collapse horizontally."""
      # Get dimensions
      if reverse:
          b, c, h, w = x[0].size()
          w *= 2
          device = x[0].device
      else:
          b, c, h, w = x.size()
          device = x.device
  
      # Get list of indices in alternating checkerboard pattern
      y_idx = []
      z_idx = []
      for i in range(h):
          for j in range(w):
              if (i % 2) == (j % 2):
                  y_idx.append(i * w + j)
              else:
                  z_idx.append(i * w + j)
      y_idx = torch.tensor(y_idx, dtype=torch.int64, device=device)
      z_idx = torch.tensor(z_idx, dtype=torch.int64, device=device)
  
      if reverse:
          y, z = (t.contiguous().view(b, c, h * w // 2) for t in x)
          x = torch.zeros(b, c, h * w, dtype=y.dtype, device=y.device)
          x[:, :, y_idx] += y
          x[:, :, z_idx] += z
          x = x.view(b, c, h, w)
  
          return x
      else:
          if w % 2 != 0:
              raise RuntimeError('Checkerboard got odd width input: {}'.format(w))
  
          x = x.view(b, c, h * w)
          y = x[:, :, y_idx].view(b, c, h, w // 2)
          z = x[:, :, z_idx].view(b, c, h, w // 2)
  
          return y, z
  
  
  def channelwise(x, reverse=False):
      """Split x channel-wise."""
      if reverse:
          x = torch.cat(x, dim=1)
          return x
      else:
          y, z = x.chunk(2, dim=1)
          return y, z
  
  
  def squeeze(x):
      """Trade spatial extent for channels. I.e., convert each
      1x4x4 volume of input into a 4x1x1 volume of output.
  
      Args:
          x (torch.Tensor): Input to squeeze.
  
      Returns:
          x (torch.Tensor): Squeezed or unsqueezed tensor.
      """
      b, c, h, w = x.size()
      x = x.view(b, c, h // 2, 2, w // 2, 2)
      x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
      x = x.view(b, c * 2 * 2, h // 2, w // 2)
  
      return x
  
  
  def unsqueeze(x):
      """Trade channels channels for spatial extent. I.e., convert each
      4x1x1 volume of input into a 1x4x4 volume of output.
  
      Args:
          x (torch.Tensor): Input to unsqueeze.
  
      Returns:
          x (torch.Tensor): Unsqueezed tensor.
      """
      b, c, h, w = x.size()
      x = x.view(b, c // 4, 2, 2, h, w)
      x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
      x = x.view(b, c // 4, h * 2, w * 2)
  
      return x
  
  
  def concat_elu(x):
      """Concatenated ReLU (http://arxiv.org/abs/1603.05201), but with ELU."""
      return F.elu(torch.cat((x, -x), dim=1))
  
  
  def safe_log(x):
      return torch.log(x.clamp(min=1e-22))
  