Pixel RNN Pytorch Implementation
I am trying to implement Pixel RNN in pytorch, but I cannot seem to find any documentation on this. The main parts of Pixel RNN are Row LSTM and BiDiagonal LSTM, so I am looking for some code of these algorithms to better understand what they are doing. Specifically, I am confused as to these algorithms calculate one row and diagonal at once, respectively. Any help would be much appreciated.

Summary

Here is an in progress partial implementation:

https://github.com/carpedm20/pixel-rnn-tensorflow

Here is a description of Row LSTM and BiDiagonal LSTM at google deepmind:

https://towardsdatascience.com/summary-of-pixelrnn-by-google-deepmind-7-min-read-938d9871d6d9



Row LSTM

From the linked deepmind blog:

The hidden state of a pixel, red in the image below, is based on the "memory" of the triangular three pixels before it. Because they are in a "row", we can compute in parallel, speeding up computation. We sacrifice some context information (using more history or memory) for the ability to do this parallel computation and speed up training.



The actual implementation relies on several other optimizations and is quite involved. From the original paper:


  The computation proceeds as follows. An LSTM layer has an
  input-to-state component and a recurrent state-to-state component that
  together determine the four gates inside the LSTM core. To enhance
  parallelization in the Row LSTM the input-to-state component is first
  computed for the entire two-dimensional input map; for this a k × 1
  convolution is used to follow the row-wise orientation of the LSTM
  itself. The convolution is masked to include only the valid context
  (see Section 3.4) and produces a tensor of size 4h × n × n,
  representing the four gate vectors for each position in the input map,
  where h is the number of output feature maps. To compute one step of
  the state-to-state component of the LSTM layer, one is given the
  previous hidden and cell states hi−1 and ci−1, each of size h × n × 1.
  The new hidden and cell states hi , ci are obtained as follows: 





  where xi of size h × n × 1 is row i of the input map, and ~ represents the convolution operation and  the elementwise
  multiplication. The weights Kss and Kis are the kernel weights for the
  state-to-state and the input-to-state components, where the latter is
  precomputed as described above. In the case of the output, forget and
  input gates oi , fi and ii , the activation σ is the logistic sigmoid
  function, whereas for the content gate gi , σ is the tanh function.
  Each step computes at once the new state for an entire row of the
  input map


Diagonal BLSTM

Diagonal BLSTM's were developed to leverage the speedup of parallelization without sacrificing as much context information. A node in a DBLSTM looks to its left and above it; since those nodes have also looked to the left and above, the conditional probability of a given node depends in some sense on all of its ancestors. Otherwise, the architectures are very similar. From the deepmind blog:


