scvi.distributions.ZeroInflatedNegativeBinomial.enumerate_support

ZeroInflatedNegativeBinomial.enumerate_support(expand=True)

Returns tensor containing all values supported by a discrete distribution. The result will enumerate over dimension 0, so the shape of the result will be (cardinality,) + batch_shape + event_shape (where event_shape = () for univariate distributions).

Note that this enumerates over all batched tensors in lock-step [[0, 0], [1, 1], …]. With expand=False, enumeration happens along dim 0, but with the remaining batch dimensions being singleton dimensions, [[0], [1], ...

To iterate over the full Cartesian product use itertools.product(m.enumerate_support()).

Parameters
expand : bool

whether to expand the support over the batch dims to match the distribution’s batch_shape.

Returns

Tensor iterating over dimension 0.