code: # Loosely derived from https://github.com/jzbontar/pixelcnn-pytorch/blob/master/main.py
  #  and moreso derived from https://github.com/rampage644/wavenet/blob/master/wavenet/models.py
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  import numpy as np
  import tqdm
  
  #owidth  = floor((width  + 2*padW - kW) / dW + 1)
  #oheight = floor((height + 2*padH - kH) / dH + 1)
  #dW is stride, assuming 1:
  # kW // 2 = padW
  def same_padding(kernel_size):
      # assumming stride 1
      if isinstance(kernel_size, int):
          return kernel_size // 2
      else:
          return (kernel_size[0] // 2, kernel_size[1] // 2)
  
  # PyTorch port of
  class MaskedConvolution2D(nn.Conv2d):
      def __init__(self, in_channels, out_channels, kernel_size, 
              *args, mask='B', vertical=False, mask_mode="noblind", **kwargs):
          if "padding" not in kwargs:
              assert "stride" not in kwargs
              kwargs["padding"] = same_padding(kernel_size)
          remove = {"conditional_features", "conditional_image_channels"}
          for feature in remove:
              if feature in kwargs:
                  del kwargs[feature]
          super(MaskedConvolution2D, self).__init__(in_channels,
                  out_channels, kernel_size, *args, **kwargs)
          Cout, Cin, kh, kw = self.weight.size()
          pre_mask = np.ones_like(self.weight.data.cpu().numpy()).astype(np.float32)
          yc, xc = kh // 2, kw // 2
  
          assert mask_mode in {"noblind", "turukin", "fig1-van-den-oord"}
          if mask_mode == "noblind":
              # context masking - subsequent pixels won't hav access
              # to next pixels (spatial dim)
              if vertical:
                  if mask == 'A':
                      # In the first layer, can ONLY access pixels above it
                      pre_mask[:, :, yc:, :] = 0.0
                  else:
                      # In the second layer, can access pixels above or even with it.
                      # Reason being that the pixels to the right or left of the current pixel
                      #  only have a receptive field of the layer above the current layer and up.
                      pre_mask[:, :, yc+1:, :] = 0.0
              else:
                  # All rows after center must be zero
                  pre_mask[:, :, yc+1:, :] = 0.0
                  ### All rows before center must be zero # XXX: not actually necessary
                  ##pre_mask[:, :, :yc, :] = 0.0
                  # All columns after center in center row must be zero
                  pre_mask[:, :, yc, xc+1:] = 0.0
  
              if mask == 'A':
                  # Center must be zero in first layer
                  pre_mask[:, :, yc, xc] = 0.0
              # same pixel masking - pixel won't access next color (conv filter dim)
              #def bmask(i_out, i_in):
              #    cout_idx = np.expand_dims(np.arange(Cout) % 3 == i_out, 1)
              #    cin_idx = np.expand_dims(np.arange(Cin) % 3 == i_in, 0)
              #    a1, a2 = np.broadcast_arrays(cout_idx, cin_idx)
              #    return a1 * a2
  
              #for j in range(3):
              #    pre_mask[bmask(j, j), yc, xc] = 0.0 if mask == 'A' else 1.0
  
              #pre_mask[bmask(0, 1), yc, xc] = 0.0
              #pre_mask[bmask(0, 2), yc, xc] = 0.0
              #pre_mask[bmask(1, 2), yc, xc] = 0.0
          elif mask_mode == "fig1-van-den-oord":
              if vertical:
                  pre_mask[:, :, yc:, :] = 0.0
              else:
                  # All rows after center must be zero
                  pre_mask[:, :, yc+1:, :] = 0.0
                  ### All rows before center must be zero # XXX: not actually necessary
                  ##pre_mask[:, :, :yc, :] = 0.0
                  # All columns after center in center row must be zero
                  pre_mask[:, :, yc, xc+1:] = 0.0
  
              if mask == 'A':
                  # Center must be zero in first layer
                  pre_mask[:, :, yc, xc] = 0.0
          elif mask_mode == "turukin":
              pre_mask[:, :, yc+1:, :] = 0.0
              pre_mask[:, :, yc, xc+1:] = 0.0
              if mask == 'A':
                  pre_mask[:, :, yc, xc] = 0.0
  
          print("%s %s MASKED CONV: %d x %d. Mask:" % (mask, "VERTICAL" if vertical else "HORIZONTAL", kh, kw))
          print(pre_mask[0, 0, :, :])
  
          self.register_buffer("mask", torch.from_numpy(pre_mask))
  
      def __call__(self, x):
          self.weight.data = self.weight.data * self.mask
          return super(MaskedConvolution2D, self).forward(x)
  
  class CroppedConvolution(nn.Conv2d):
      def __init__(self, in_channels, out_channels, kernel_size,
              *args, **kwargs):
          if "padding" not in kwargs:
              assert "stride" not in kwargs
              kwargs["padding"] = same_padding(kernel_size)
          super().__init__(in_channels, out_channels, kernel_size,
                  *args, **kwargs)
  
      def __call__(self, x):
          ret = super().__call__(x)
          _, _, kh, kw = self.weight.size()
          pad_h, pad_w = self.padding
          h_crop = -(kh + 1) if pad_h == kh else None
          w_crop = -(kw + 1) if pad_w == kw else None
          return ret[:, :, :h_crop, :w_crop]
  
  class PixelCNNGatedLayer(nn.Module):
      def __init__(self, primary, in_channels, out_channels, filter_size,
              mask='B', nobias=False, conditional_features=None,
              conditional_image_channels=None, residual_vertical=False,
              residual_horizontal=True, skips=False, gated=True,
              relu_out=False, horizontal_2d_convs=False, mask_mode="noblind"):
          super().__init__()
          self.primary = primary
          if primary:
              assert mask == 'A'
              assert not residual_vertical
              assert not residual_horizontal
          else:
              assert mask == 'B'
          self.out_channels = out_channels
          self.gated = gated
          gm = 2 if gated else 1
          self.vertical_conv = MaskedConvolution2D(
              in_channels, gm * out_channels, (filter_size, filter_size),
              mask=mask, vertical=True, mask_mode=mask_mode)
          self.v_to_h_conv = nn.Conv2d(gm * out_channels, gm * out_channels, 1)
  
          self.horizontal_conv = MaskedConvolution2D(
              in_channels, gm * out_channels,
              (filter_size if horizontal_2d_convs else 1, filter_size), # XXX: traditionally (1, filter_size),
              mask=mask, vertical=False, mask_mode=mask_mode)
  
          self.residual_vertical = None
          if residual_vertical:
              self.residual_vertical = nn.Conv2d(in_channels, gm * out_channels, 1)
  
          self.horizontal_output = nn.Conv2d(out_channels, out_channels, 1)
          self.horizontal_skip = None
          if skips:
              self.horizontal_skip = nn.Conv2d(out_channels, out_channels, 1)
          self.conditional_vector = conditional_features is not None
          self.conditional_image = conditional_image_channels is not None
          if self.conditional_image:
              self.cond_conv_h = nn.Conv2d(conditional_image_channels, gm * out_channels, 1, bias=False)
              self.cond_conv_v = nn.Conv2d(conditional_image_channels, gm * out_channels, 1, bias=False)
          if self.conditional_vector:
              self.cond_fc_h = nn.Linear(conditional_features, gm * out_channels, bias=False)
              self.cond_fc_v = nn.Linear(conditional_features, gm * out_channels, bias=False)
          self.residual_horizontal = residual_horizontal
          self.relu_out = relu_out
  
      @classmethod
      def primary(cls, in_channels, out_channels, filter_size,
              nobias=False, conditional_features=None,
              conditional_image_channels=None, 
              skips=False, gated=True,
              relu_out=False, horizontal_2d_convs=False, mask_mode="noblind"):
          return cls(True, in_channels, out_channels, filter_size, nobias=nobias,
                  mask='A', conditional_features=conditional_features,
                  conditional_image_channels=conditional_image_channels,
                  residual_vertical=False, residual_horizontal=False,
                  skips=skips, gated=gated,
                  relu_out=relu_out, horizontal_2d_convs=horizontal_2d_convs,
                  mask_mode=mask_mode)
      
      @classmethod
      def secondary(cls, in_channels, out_channels, filter_size,
              nobias=False, conditional_features=None,
              conditional_image_channels=None, residual_vertical=True,
              residual_horizontal=True, skips=False, gated=True,
              relu_out=False, horizontal_2d_convs=False, mask_mode="noblind"):
          return cls(False, in_channels, out_channels, filter_size, nobias=nobias,
                  mask='B', conditional_features=conditional_features,
                  conditional_image_channels=conditional_image_channels,
                  residual_vertical=residual_vertical, residual_horizontal=residual_horizontal,
                  skips=skips, gated=gated, relu_out=relu_out,
                  horizontal_2d_convs=horizontal_2d_convs, mask_mode=mask_mode)
  
      def _gate(self, x):
          if self.gated:
              return F.tanh(x[:,:self.out_channels]) * F.sigmoid(x[:,self.out_channels:])
          else:
              return x
  
      def __call__(self, v, h, conditional_image=None, conditional_vector=None):
          horizontal_preactivation = self.horizontal_conv(h) # 1xN
          vertical_preactivation = self.vertical_conv(v) # NxN
          v_to_h = self.v_to_h_conv(vertical_preactivation) # 1x1
          if self.residual_vertical is not None:
              vertical_preactivation = vertical_preactivation + self.residual_vertical(v) # 1x1 to residual
          horizontal_preactivation = horizontal_preactivation + v_to_h
          if self.conditional_image and conditional_image is not None:
              horizontal_preactivation = horizontal_preactivation + \
                      self.cond_conv_h(conditional_image)
              vertical_preactivation = vertical_preactivation + \
                      self.cond_conv_v(conditional_image)
          if self.conditional_vector and conditional_vector is not None:
              horizontal_preactivation = horizontal_preactivation + \
                      self.cond_fc_h(conditional_vector).unsqueeze(-1).unsqueeze(-1)
              vertical_preactivation = vertical_preactivation + \
                      self.cond_fc_v(conditional_vector).unsqueeze(-1).unsqueeze(-1)
          v_out = self._gate(vertical_preactivation)
          h_activated = self._gate(horizontal_preactivation)
          h_skip = None
          if self.horizontal_skip is not None:
              h_skip = self.horizontal_skip(h_activated)
          h_preres = self.horizontal_output(h_activated)
          if self.residual_horizontal:
              h_out = h + h_preres
          else:
              h_out = h_preres
          if self.relu_out:
              v_out = F.relu(v_out)
              h_out = F.relu(h_out)
              if h_skip is not None:
                  h_skip = F.relu(h_skip)
          return v_out, h_out, h_skip
  
  class PixelCNNGatedStack(nn.Module):
      def __init__(self, *args):
          super().__init__()
          layers = list(args)
          for i, layer in enumerate(layers):
              assert isinstance(layer, PixelCNNGatedLayer)
              if i == 0:
                  assert layer.primary
              else:
                  assert not layer.primary
          self.layers = nn.ModuleList(layers)
  
      def __call__(self, v, h, skips=None, conditional_image=None, conditional_vector=None):
          if skips is None:
              skips = []
          else:
              skips = [skips]
          for layer in self.layers:
              v, h, skip = layer(v, h, conditional_image=conditional_image, conditional_vector=conditional_vector)
              if skip is not None:
                  skips.append(skip)
          if len(skips) == 0:
              skips = None
          else:
              skips = torch.cat(skips, 1)
          return v, h, skips
  
  class PixelCNN(nn.Module):
      def _sample_from(self, probas):
          N, level_count = probas.size()
          val = torch.rand(N, 1)
          if probas.is_cuda:
              val = val.cuda()
          cutoffs = torch.cumsum(probas, dim=1)
          _, idx = torch.max(cutoffs > val, dim=1)
          out = idx.float() / (level_count - 1)
          return out 
  
      def sample_pixel(self, canvas, row, col, channel, **kwargs):
          probs = F.softmax(self(canvas, **kwargs)[:, :, channel, row, col], dim=1)
          return self._sample_from(probs.data)
  
      def generate_samples(self, height, width,
              channels, count, show_prog=False, **kwargs):
          samples = torch.zeros(count, channels, height, width).float()
          if next(self.parameters()).data.is_cuda:
              samples = samples.cuda()
          samples = Variable(samples, requires_grad=False)
          
          def generate(prog=None):
              for i in range(height):
                  for j in range(width):
                      for k in range(channels):
                          samples.data[:, k, i, j] = self.sample_pixel(samples,
                                  i, j, k, **kwargs)
                          if prog is not None:
                              prog.update()
          if show_prog:
              with tqdm.tqdm(total=height * width * channels) as prog:
                  generate(prog=prog)
          else:
              generate()
          return samples
  
  if __name__ == "__main__":
      x = PixelCNN(3, 128, 5, 16, 8, conditional_features=10,
              conditional_image_channels=8)
      print(x)
      inp = Variable(torch.randn(128, 3, 28, 28))
      cv = Variable(torch.randn(128, 10))
      ci = Variable(torch.randn(128, 8, 28, 28))
      print(x(inp, conditional_vector=cv, conditional_image=ci))
      print(x.generate_samples(28, 28, 3, 4, conditional_vector=cv, conditional_image=ci, show_prog=True))
  
apis: torch.nn.Linear torch.nn.functional.tanh torch.nn.functional.sigmoid torch.randn torch.rand torch.autograd.Variable torch.nn.functional.relu torch.cumsum torch.from_numpy torch.nn.Conv2d torch.nn.ModuleList torch.max torch.zeros torch.cat