def _m_step(self, G, doc_data):
        """_m_step implements word data sampling and stores word_z stats.
        mean_change is in the sense of full graph with rho=1.
        """
        G = G.clone()
        self._prepare_graph(G, doc_data)
        G.update_all(
            lambda edges: {"phi": EdgeData(edges.src, edges.dst).phi},
            dgl.function.sum("phi", "nphi"),
        )
        self._last_mean_change = self.word_data.update_from(
            G, self.mult["word"], self.rho
        )

        if self.verbose:
            print(f"m-step mean_change={self._last_mean_change:.4f}, ", end="")
            Bayesian_gap = np.mean(
                [part.Bayesian_gap.mean().tolist() for part in self.word_data]
            )
            print(f"Bayesian_gap={Bayesian_gap:.4f}")
_m_step implements word data sampling and stores word_z stats.
        mean_change is in the sense of full graph with rho=1.
        
examples/pytorch/lda/lda_model.py