Inverse Autoregressive Spline Flow Implementation

I have to implement an algorithm for a university project, however I can not seem to wrap my head around it. The algorithm should be an inverse autoregressive normalizing flow using splines. It should work for 1D data (time series) and not for 2D images. It is inspired by this paper https://arxiv.org/pdf/1606.04934.pdf. In particular what is important for me is this algorithm:

I dropped the part with the EncoderNN and for the AutoregressiveNN i am using a pixelCNN (implementation follows). In addition I have a function that can calculate z and the log probability given the W, H and D. Also I am sampling the initial z in a differently, however, that is not what’s important.

class PixelCNN(nn.Module):
    def __init__(self, c_in, c_out ,n_layers, K=3, B=8 ,kernel_size=7,n_classes=3):
        super(PixelCNN,self).__init__()

        ''' PixelCNN part '''
        self.init_layer = MaskedConv1d(c_in,c_out, kernel_size,mask_type='A')
        self.layers = nn.ModuleList([])
        for _ in range(n_layers):
            self.layers += nn.ModuleList([MaskedConv1d(c_out,c_out, kernel_size//2,mask_type='B')])
            self.layers += nn.ModuleList([nn.Tanh()])


        ''' Splines part '''
        self.K = K
        self.B = B
        self.n_transforms = n_transforms
        self.network = GatedResNet(c_in,c_out,3 * K - 1)

        ## Make rational quadratic splines
    def forward(self,x):
        x = self.init_layer(x)
        for layer in self.layers:
            x = layer(x)
        return self.create_splines(x)

    def create_splines(self,x):
        out = self.network(x)
        W,H,D = torch.split(out,self.K,dim=-1)
        W = torch.softmax(W, dim=-1)
        H = torch.softmax(H, dim=-1)

        W = W * 2 * self.B
        H = H * 2 * self.B

        D = F.softplus(D)
        return W,H,D

The AutoregressiveNN in the algorithm above has to output the parameters m and s which are used to improve z. In my case if I am not mistaken those should be the splines, so the forward method in the pixelCNN returns the splines only (I am not sure if this is correct). Also I do not know what to do with the loop for for t <- 1 to T do. Since each loop outputs m,s then should I also output a set of splines in this manner? I’ve seen some implementation of IAF using MAF without splines and sometimes they just make a function forward and backward and then they swap their places in training and sampling. Is that also possible here and is it a better option?

And finally, how should I update z with the splines? More clarification might be needed for people to answer. I can provide that if I missed something.