def _edge_mask_to_heterogeneous(self, edge_mask, homograph, heterograph):
        r"""Convert an edge mask from homogeneous mappings built through
        embeddings into heterogenous format by leveraging the context from
        the source DGLGraphs in homogenous and heterogeneous form.

        The `edge_mask` needs to have been built using the embedding of the
        homogenous graph format for the mappings to work correctly.

        Parameters
        ----------
        edge_mask : dict[str, Tensor]
            A dict mapping node types (keys) to a tensor of edge weights (values).
        homograph : DGLGraph
            The homogeneous form of the source graph.
        heterograph : DGLGraph
            The heterogeneous form of the source graph.

        Returns
        -------
        dict[str, Tensor]
            A dict mapping node types (keys) to tensors of node ids (values)
        """
        return {
            etype: edge_mask[
                (homograph.edata[ETYPE] == heterograph.get_etype_id(etype))
                .nonzero()
                .squeeze(1)
            ]
            for etype in heterograph.canonical_etypes
        }
Convert an edge mask from homogeneous mappings built through
        embeddings into heterogenous format by leveraging the context from
        the source DGLGraphs in homogenous and heterogeneous form.

        The `edge_mask` needs to have been built using the embedding of the
        homogenous graph format for the mappings to work correctly.

        Parameters
        ----------
        edge_mask : dict[str, Tensor]
            A dict mapping node types (keys) to a tensor of edge weights (values).
        homograph : DGLGraph
            The homogeneous form of the source graph.
        heterograph : DGLGraph
            The heterogeneous form of the source graph.

        Returns
        -------
        dict[str, Tensor]
            A dict mapping node types (keys) to tensors of node ids (values)
        
python/dgl/nn/pytorch/explain/pgexplainer.py