Ah, right, log_prob / local MLE simultaneous estimation works, it is just not too good with random nn initializations and SGD. May be ok for some tasks.
To directly answer your question:
    def log_prob(self, x):
        x = self._pad(x)
        log_prob_x = self.component_distribution.log_prob(x)  # [S, B, k]
        log_mix_prob = torch.log_softmax(self.mixture_distribution.logits,
                                         dim=-1)  # [B, k]
        return torch.logsumexp(log_prob_x + log_mix_prob, dim=-1)  # [S, B]

this is just transformed mixture density formula, where everything is differentiable (similarly to a weighted sum).