Multinomial on convolutionnal model returns nans (equivalent code with Pyro)

This is the same when I’m doing more simply:

data <- multinomial(tf.nn.conv2d(input, filter))
parameter <- Dirichlet(…)
qparameter <- Dirichlet(…)
latent <- tf.nn.conv2d(parameter, filter)
w <- multinomial(latent)
inference({parameter:qparameter}, {w:data})

without the “complex” roll function, it is then a basic convolution.

I’m missing something ?