fe9540ff6e0ad1aebca9bbc951b46289b4a2f9b2,pyro/contrib/epidemiology/compartmental.py,CompartmentalModel,_sequential_model,#CompartmentalModel#,648

Before Change

        auxiliary, non_compartmental = self._sample_auxiliary()

        // Reshape to accommodate the time_plate below.
        if self.is_regional:
            auxiliary = auxiliary.squeeze(1)
            for name, value in non_compartmental.items():
                non_compartmental[name] = value.squeeze(1)
        assert auxiliary.shape == (num_samples, 1, C, T) + R_shape
        aux = [aux.unbind(2) for aux in auxiliary.unbind(2)]

        // Sequentially transition.

After Change

        // Reshape to accommodate the time_plate below.
        assert auxiliary.shape == (num_samples, C, T) + R_shape, auxiliary.shape
        aux = [aux.unbind(2) for aux in auxiliary.unsqueeze().unbind(2)]

        // Sequentially transition.
        curr = self.initialize(params)

Instances