1bb8f15b13d1fb3713366228d431da29b2e5c43e,pyro/distributions/torch_distribution.py,MaskedDistribution,__init__,#MaskedDistribution#Any#Any#,254

Before Change

        if isinstance(mask, bool):
            self._mask = mask
        else:
            if broadcast_shape(mask.shape, base_dist.batch_shape) != base_dist.batch_shape:
                raise ValueError("Expected mask.shape to be broadcastable to base_dist.batch_shape, "
                                 "actual {} vs {}".format(mask.shape, base_dist.batch_shape))
            self._mask = mask.bool()
        super(MaskedDistribution, self).__init__(base_dist.batch_shape, base_dist.event_shape)

    def expand(self, batch_shape, _instance=None):

After Change

            self._mask = mask
        else:
            batch_shape = broadcast_shape(mask.shape, base_dist.batch_shape)
            if mask.shape != batch_shape:
                mask = mask.expand(batch_shape)
            if base_dist.batch_shape != batch_shape:
                base_dist = base_dist.expand(batch_shape)
            self._mask = mask.bool()
        self.base_dist = base_dist

Instances