Masking events element-wise when calculating likelihood

[jubbens]

Issue Description

I have some categorical data of shape (B x E x C) where B is the batch size, E is the number of features, and C is the number of categories. However my model only produces logits for some elements in E (different features for each sample in the batch) and I don't care about the ones it is not producing for, so I would like to mask those ones out (represented as NaN in the logits) when calculating the likelihood.

I am aware that one can use obs_mask or poutine.mask when using sample to mask along the batch dimension, but how can one mask along other dimensions?

Apologies if there is already an obvious way to do this - I was going to post it in the forum, but there doesn't seem to be any way to create a new account (and github login is broken).

[fritzo]

You can use the .mask() method of some distributions to mask along a non-batch dimension.

But please prefer https://forum.pyro.ai for questions, as we try to reserve github for issues & contributions. Good luck!