I think I know exactly the paper you are trying to implement, and the z distrubution is already supposed to be activated with Softmax. The annealing is just a way to rebalance the final distribution for each pixel to avoid the desaturated effect of taking the mean color of a distribution.
And for the record, it’s completely expected to have negative values in feature maps, otherwise ReLU would have no effect.