class_docstr: PGExplainer from `Parameterized Explainer for Graph Neural Network
  <https://arxiv.org/pdf/2011.04573>`__, adapted for heterogeneous graphs
  
  PGExplainer adopts a deep neural network (explanation network) to
  parameterize the generation process of explanations, which enables it to
  explain multiple instances collectively. PGExplainer models the underlying
  structure as edge distributions, from which the explanatory graph is
  sampled.
  
  Parameters
  ----------
  model : nn.Module
      The GNN model to explain that tackles multiclass graph classification
  
      * Its forward function must have the form
        :attr:`forward(self, graph, nfeat, embed, edge_weight)`.
      * The output of its forward function is the logits if embed=False else
        the intermediate node embeddings.
  num_features : int
      Node embedding size used by :attr:`model`.
  coff_budget : float, optional
      Size regularization to constrain the explanation size. Default: 0.01.
  coff_connect : float, optional
      Entropy regularization to constrain the connectivity of explanation. Default: 5e-4.
  sample_bias : float, optional
      Some members of a population are systematically more likely to be selected
      in a sample than others. Default: 0.0.
class_docstr_tok: ["PGExplainer","from","`","Parameterized","Explainer","for","Graph","Neural","Network","<","https",":","\/\/arxiv.org\/pdf\/2011.04573",">","`","__",",","adapted","for","heterogeneous","graphs","PGExplainer","adopts","a","deep","neural","network","(","explanation","network",")","to","parameterize","the","generation","process","of","explanations",",","which","enables","it","to","explain","multiple","instances","collectively",".","PGExplainer","models","the","underlying","structure","as","edge","distributions",",","from","which","the","explanatory","graph","is","sampled",".","Parameters","--","--","--","--","--","model",":","nn.Module","The","GNN","model","to","explain","that","tackles","multiclass","graph","classification","*","Its","forward","function","must","have","the","form",":","attr",":","`","forward","(","self",",","graph",",","nfeat",",","embed",",","edge_weight",")","`",".","*","The","output","of","its","forward","function","is","the","logits","if","embed=False","else","the","intermediate","node","embeddings",".","num_features",":","int","Node","embedding","size","used","by",":","attr",":","`","model","`",".","coff_budget",":","float",",","optional","Size","regularization","to","constrain","the","explanation","size",".","Default",":","0.01.","coff_connect",":","float",",","optional","Entropy","regularization","to","constrain","the","connectivity","of","explanation",".","Default",":","5e-4",".","sample_bias",":","float",",","optional","Some","members","of","a","population","are","systematically","more","likely","to","be","selected","in","a","sample","than","others",".","Default",":","0.0","."]
docstr: Learn and return an edge mask that plays a crucial role to
  explain the prediction made by the GNN for provided set of node IDs.
  Also, return the prediction made with the batched graph and edge mask.
  
  Parameters
  ----------
  nodes : dict[str, Iterable[int]]
      A dict mapping node types (keys) to an iterable set of node ids (values).
  graph : DGLGraph
      A heterogeneous graph.
  feat : dict[str, Tensor]
      A dict mapping node types (keys) to feature tensors (values).
      The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is
      the number of nodes for node type :math:`t`, and :math:`D_t` is the
      feature size for node type :math:`t`
  temperature : float
      The temperature parameter fed to the sampling procedure.
  training : bool
      Training the explanation network.
  kwargs : dict
      Additional arguments passed to the GNN model.
  
  Returns
  -------
  dict[str, Tensor]
      A dict mapping node types (keys) to classification probabilities
      for node labels (values). The values are tensors of shape
      :math:`(N_t, L)`, where :math:`L` is the different types of node
      labels in the dataset, and :math:`N_t` is the number of nodes in
      the graph for node type :math:`t`.
  dict[str, Tensor]
      A dict mapping edge types (keys) to edge tensors (values) of shape
      :math:`(E_t)`, where :math:`E_t` is the number of edges in the graph
      for edge type :math:`t`.  A higher weight suggests a larger
      contribution of the edge.
  DGLGraph
      The batched set of subgraphs induced on the k-hop in-neighborhood
      of the input center nodes.
  dict[str, Tensor]
      A dict mapping node types (keys) to a tensor of node IDs (values)
      which correspond to the subgraph center nodes.
  
  Examples
  --------
  
  >>> import dgl
  >>> import torch as th
  >>> import torch.nn as nn
  >>> import numpy as np
  
  >>> # Define the model
  >>> class Model(nn.Module):
  ...     def __init__(self, in_feats, hid_feats, out_feats, rel_names):
  ...         super().__init__()
  ...         self.conv = dgl.nn.HeteroGraphConv(
  ...             {rel: dgl.nn.GraphConv(in_feats, hid_feats) for rel in rel_names},
  ...             aggregate="sum",
  ...         )
  ...         self.fc = nn.Linear(hid_feats, out_feats)
  ...         nn.init.xavier_uniform_(self.fc.weight)
  ...
  ...     def forward(self, g, h, embed=False, edge_weight=None):
  ...         if edge_weight:
  ...             mod_kwargs = {
  ...                 etype: {"edge_weight": mask} for etype, mask in edge_weight.items()
  ...             }
  ...             h = self.conv(g, h, mod_kwargs=mod_kwargs)
  ...         else:
  ...             h = self.conv(g, h)
  ...
  ...         return h
  
  >>> # Load dataset
  >>> input_dim = 5
  >>> hidden_dim = 5
  >>> num_classes = 2
  >>> g = dgl.heterograph({("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1])})
  >>> g.nodes["user"].data["h"] = th.randn(g.num_nodes("user"), input_dim)
  >>> g.nodes["game"].data["h"] = th.randn(g.num_nodes("game"), input_dim)
  
  >>> transform = dgl.transforms.AddReverse()
  >>> g = transform(g)
  
  >>> # define and train the model
  >>> model = Model(input_dim, hidden_dim, num_classes, g.canonical_etypes)
  >>> optimizer = th.optim.Adam(model.parameters())
  >>> for epoch in range(10):
  ...     logits = model(g, g.ndata["h"])['user']
  ...     loss = th.nn.functional.cross_entropy(logits, th.tensor([1,1,1]))
  ...     optimizer.zero_grad()
  ...     loss.backward()
  ...     optimizer.step()
  
  >>> # Initialize the explainer
  >>> explainer = dgl.nn.HeteroPGExplainer(
  ...     model, hidden_dim, num_hops=2, explain_graph=False
  ... )
  
  >>> # Train the explainer
  >>> # Define explainer temperature parameter
  >>> init_tmp, final_tmp = 5.0, 1.0
  >>> optimizer_exp = th.optim.Adam(explainer.parameters(), lr=0.01)
  >>> for epoch in range(20):
  ...     tmp = float(init_tmp * np.power(final_tmp / init_tmp, epoch / 20))
  ...     loss = explainer.train_step_node(
  ...         { ntype: g.nodes(ntype) for ntype in g.ntypes },
  ...         g, g.ndata["h"], tmp
  ...     )
  ...     optimizer_exp.zero_grad()
  ...     loss.backward()
  ...     optimizer_exp.step()
  
  >>> # Explain the graph
  >>> feat = g.ndata.pop("h")
  >>> probs, edge_mask, bg, inverse_indices = explainer.explain_node(
  ...     { "user": [0] }, g, feat
  ... )
docstr_tok: ["Learn","and","return","an","edge","mask","that","plays","a","crucial","role","to","explain","the","prediction","made","by","the","GNN","for","provided","set","of","node","IDs",".","Also",",","return","the","prediction","made","with","the","batched","graph","and","edge","mask",".","Parameters","--","--","--","--","--","nodes",":","dict","[","str",",","Iterable","[","int","]","]","A","dict","mapping","node","types","(","keys",")","to","an","iterable","set","of","node","ids","(","values",")",".","graph",":","DGLGraph","A","heterogeneous","graph",".","feat",":","dict","[","str",",","Tensor","]","A","dict","mapping","node","types","(","keys",")","to","feature","tensors","(","values",")",".","The","input","features","are","of","shape",":","math",":","`","(","N_t",",","D_t",")","`",".",":","math",":","`","N_t","`","is","the","number","of","nodes","for","node","type",":","math",":","`","t","`",",","and",":","math",":","`","D_t","`","is","the","feature","size","for","node","type",":","math",":","`","t","`","temperature",":","float","The","temperature","parameter","fed","to","the","sampling","procedure",".","training",":","bool","Training","the","explanation","network",".","kwargs",":","dict","Additional","arguments","passed","to","the","GNN","model",".","Returns","--","--","--","-","dict","[","str",",","Tensor","]","A","dict","mapping","node","types","(","keys",")","to","classification","probabilities","for","node","labels","(","values",")",".","The","values","are","tensors","of","shape",":","math",":","`","(","N_t",",","L",")","`",",","where",":","math",":","`","L","`","is","the","different","types","of","node","labels","in","the","dataset",",","and",":","math",":","`","N_t","`","is","the","number","of","nodes","in","the","graph","for","node","type",":","math",":","`","t","`",".","dict","[","str",",","Tensor","]","A","dict","mapping","edge","types","(","keys",")","to","edge","tensors","(","values",")","of","shape",":","math",":","`","(","E_t",")","`",",","where",":","math",":","`","E_t","`","is","the","number","of","edges","in","the","graph","for","edge","type",":","math",":","`","t","`",".","A","higher","weight","suggests","a","larger","contribution","of","the","edge",".","DGLGraph","The","batched","set","of","subgraphs","induced","on","the","k-hop","in-neighborhood","of","the","input","center","nodes",".","dict","[","str",",","Tensor","]","A","dict","mapping","node","types","(","keys",")","to","a","tensor","of","node","IDs","(","values",")","which","correspond","to","the","subgraph","center","nodes",".","Examples","--","--","--","--",">",">",">","import","dgl",">",">",">","import","torch","as","th",">",">",">","import","torch.nn","as","nn",">",">",">","import","numpy","as","np",">",">",">","#","Define","the","model",">",">",">","class","Model","(","nn.Module",")",":","...","def","__init__","(","self",",","in_feats",",","hid_feats",",","out_feats",",","rel_names",")",":","...","super","(",")",".__init__","(",")","...","self.conv","=","dgl.nn.HeteroGraphConv","(","...","{","rel",":","dgl.nn.GraphConv","(","in_feats",",","hid_feats",")","for","rel","in","rel_names","}",",","...","aggregate=","''","sum","''",",","...",")","...","self.fc","=","nn.Linear","(","hid_feats",",","out_feats",")","...","nn.init.xavier_uniform_","(","self.fc.weight",")","...","...","def","forward","(","self",",","g",",","h",",","embed=False",",","edge_weight=None",")",":","...","if","edge_weight",":","...","mod_kwargs","=","{","...","etype",":","{","``","edge_weight","''",":","mask","}","for","etype",",","mask","in","edge_weight.items","(",")","...","}","...","h","=","self.conv","(","g",",","h",",","mod_kwargs=mod_kwargs",")","...","else",":","...","h","=","self.conv","(","g",",","h",")","...","...","return","h",">",">",">","#","Load","dataset",">",">",">","input_dim","=","5",">",">",">","hidden_dim","=","5",">",">",">","num_classes","=","2",">",">",">","g","=","dgl.heterograph","(","{","(","``","user","''",",","``","plays","''",",","``","game","''",")",":","(","[","0",",","1",",","1",",","2","]",",","[","0",",","0",",","1",",","1","]",")","}",")",">",">",">","g.nodes","[","``","user","''","]",".data","[","``","h","''","]","=","th.randn","(","g.num_nodes","(","``","user","''",")",",","input_dim",")",">",">",">","g.nodes","[","``","game","''","]",".data","[","``","h","''","]","=","th.randn","(","g.num_nodes","(","``","game","''",")",",","input_dim",")",">",">",">","transform","=","dgl.transforms.AddReverse","(",")",">",">",">","g","=","transform","(","g",")",">",">",">","#","define","and","train","the","model",">",">",">","model","=","Model","(","input_dim",",","hidden_dim",",","num_classes",",","g.canonical_etypes",")",">",">",">","optimizer","=","th.optim.Adam","(","model.parameters","(",")",")",">",">",">","for","epoch","in","range","(","10",")",":","...","logits","=","model","(","g",",","g.ndata","[","``","h","''","]",")","[","'user","'","]","...","loss","=","th.nn.functional.cross_entropy","(","logits",",","th.tensor","(","[","1,1,1","]",")",")","...","optimizer.zero_grad","(",")","...","loss.backward","(",")","...","optimizer.step","(",")",">",">",">","#","Initialize","the","explainer",">",">",">","explainer","=","dgl.nn.HeteroPGExplainer","(","...","model",",","hidden_dim",",","num_hops=2",",","explain_graph=False","...",")",">",">",">","#","Train","the","explainer",">",">",">","#","Define","explainer","temperature","parameter",">",">",">","init_tmp",",","final_tmp","=","5.0",",","1.0",">",">",">","optimizer_exp","=","th.optim.Adam","(","explainer.parameters","(",")",",","lr=0.01",")",">",">",">","for","epoch","in","range","(","20",")",":","...","tmp","=","float","(","init_tmp","*","np.power","(","final_tmp","\/","init_tmp",",","epoch","\/","20",")",")","...","loss","=","explainer.train_step_node","(","...","{","ntype",":","g.nodes","(","ntype",")","for","ntype","in","g.ntypes","}",",","...","g",",","g.ndata","[","``","h","''","]",",","tmp","...",")","...","optimizer_exp.zero_grad","(",")","...","loss.backward","(",")","...","optimizer_exp.step","(",")",">",">",">","#","Explain","the","graph",">",">",">","feat","=","g.ndata.pop","(","``","h","''",")",">",">",">","probs",",","edge_mask",",","bg",",","inverse_indices","=","explainer.explain_node","(","...","{","``","user","''",":","[","0","]","}",",","g",",","feat","...",")"]
code: def explain_node(
      self,
      nodes,
      graph,
      feat,
      temperature=1.0,
      training=False,
      **kwargs,
  ):
      r"""Learn and return an edge mask that plays a crucial role to
      explain the prediction made by the GNN for provided set of node IDs.
      Also, return the prediction made with the batched graph and edge mask.
  
      Parameters
      ----------
      nodes : dict[str, Iterable[int]]
          A dict mapping node types (keys) to an iterable set of node ids (values).
      graph : DGLGraph
          A heterogeneous graph.
      feat : dict[str, Tensor]
          A dict mapping node types (keys) to feature tensors (values).
          The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is
          the number of nodes for node type :math:`t`, and :math:`D_t` is the
          feature size for node type :math:`t`
      temperature : float
          The temperature parameter fed to the sampling procedure.
      training : bool
          Training the explanation network.
      kwargs : dict
          Additional arguments passed to the GNN model.
  
      Returns
      -------
      dict[str, Tensor]
          A dict mapping node types (keys) to classification probabilities
          for node labels (values). The values are tensors of shape
          :math:`(N_t, L)`, where :math:`L` is the different types of node
          labels in the dataset, and :math:`N_t` is the number of nodes in
          the graph for node type :math:`t`.
      dict[str, Tensor]
          A dict mapping edge types (keys) to edge tensors (values) of shape
          :math:`(E_t)`, where :math:`E_t` is the number of edges in the graph
          for edge type :math:`t`.  A higher weight suggests a larger
          contribution of the edge.
      DGLGraph
          The batched set of subgraphs induced on the k-hop in-neighborhood
          of the input center nodes.
      dict[str, Tensor]
          A dict mapping node types (keys) to a tensor of node IDs (values)
          which correspond to the subgraph center nodes.
  
      Examples
      --------
  
      >>> import dgl
      >>> import torch as th
      >>> import torch.nn as nn
      >>> import numpy as np
  
      >>> # Define the model
      >>> class Model(nn.Module):
      ...     def __init__(self, in_feats, hid_feats, out_feats, rel_names):
      ...         super().__init__()
      ...         self.conv = dgl.nn.HeteroGraphConv(
      ...             {rel: dgl.nn.GraphConv(in_feats, hid_feats) for rel in rel_names},
      ...             aggregate="sum",
      ...         )
      ...         self.fc = nn.Linear(hid_feats, out_feats)
      ...         nn.init.xavier_uniform_(self.fc.weight)
      ...
      ...     def forward(self, g, h, embed=False, edge_weight=None):
      ...         if edge_weight:
      ...             mod_kwargs = {
      ...                 etype: {"edge_weight": mask} for etype, mask in edge_weight.items()
      ...             }
      ...             h = self.conv(g, h, mod_kwargs=mod_kwargs)
      ...         else:
      ...             h = self.conv(g, h)
      ...
      ...         return h
  
      >>> # Load dataset
      >>> input_dim = 5
      >>> hidden_dim = 5
      >>> num_classes = 2
      >>> g = dgl.heterograph({("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1])})
      >>> g.nodes["user"].data["h"] = th.randn(g.num_nodes("user"), input_dim)
      >>> g.nodes["game"].data["h"] = th.randn(g.num_nodes("game"), input_dim)
  
      >>> transform = dgl.transforms.AddReverse()
      >>> g = transform(g)
  
      >>> # define and train the model
      >>> model = Model(input_dim, hidden_dim, num_classes, g.canonical_etypes)
      >>> optimizer = th.optim.Adam(model.parameters())
      >>> for epoch in range(10):
      ...     logits = model(g, g.ndata["h"])['user']
      ...     loss = th.nn.functional.cross_entropy(logits, th.tensor([1,1,1]))
      ...     optimizer.zero_grad()
      ...     loss.backward()
      ...     optimizer.step()
  
      >>> # Initialize the explainer
      >>> explainer = dgl.nn.HeteroPGExplainer(
      ...     model, hidden_dim, num_hops=2, explain_graph=False
      ... )
  
      >>> # Train the explainer
      >>> # Define explainer temperature parameter
      >>> init_tmp, final_tmp = 5.0, 1.0
      >>> optimizer_exp = th.optim.Adam(explainer.parameters(), lr=0.01)
      >>> for epoch in range(20):
      ...     tmp = float(init_tmp * np.power(final_tmp / init_tmp, epoch / 20))
      ...     loss = explainer.train_step_node(
      ...         { ntype: g.nodes(ntype) for ntype in g.ntypes },
      ...         g, g.ndata["h"], tmp
      ...     )
      ...     optimizer_exp.zero_grad()
      ...     loss.backward()
      ...     optimizer_exp.step()
  
      >>> # Explain the graph
      >>> feat = g.ndata.pop("h")
      >>> probs, edge_mask, bg, inverse_indices = explainer.explain_node(
      ...     { "user": [0] }, g, feat
      ... )
      """
      assert (
          not self.graph_explanation
      ), '"explain_graph" must be False when initializing the module.'
      assert (
          self.num_hops is not None
      ), '"num_hops" must be provided when initializing the module.'
  
      self.model = self.model.to(graph.device)
      self.elayers = self.elayers.to(graph.device)
  
      batched_embed = []
      batched_homo_graph = []
      batched_hetero_graph = []
      for target_ntype, target_nids in nodes.items():
          if isinstance(target_nids, torch.Tensor):
              target_nids = target_nids.tolist()
  
          for target_nid in target_nids:
              sg, inverse_indices = khop_in_subgraph(
                  graph, {target_ntype: target_nid}, self.num_hops
              )
  
              for sg_ntype in sg.ntypes:
                  sg_feat = feat[sg_ntype][
                      sg.ndata[NID][sg_ntype].long()
                  ]
                  train_mask = [
                      sg_ntype in inverse_indices
                      and node_id in inverse_indices[sg_ntype]
                      for node_id in sg.nodes(sg_ntype)
                  ]
  
                  sg.nodes[sg_ntype].data["feat"] = sg_feat
                  sg.nodes[sg_ntype].data["train"] = torch.tensor(
                      train_mask, device=sg.device
                  )
  
              embed = self.model(
                  sg, sg.ndata["feat"], embed=True, **kwargs
              )
              for ntype in embed.keys():
                  sg.nodes[ntype].data["emb"] = embed[ntype].data
  
              homo_sg = to_homogeneous(sg, ndata=["emb"])
              homo_sg_embed = homo_sg.ndata["emb"]
  
              col, row = homo_sg.edges()
              col_emb = homo_sg_embed[col.long()]
              row_emb = homo_sg_embed[row.long()]
              self_emb = homo_sg_embed[
                  inverse_indices[target_ntype][0]
              ].repeat(sg.num_edges(), 1)
              emb = torch.cat([col_emb, row_emb, self_emb], dim=-1)
              batched_embed.append(emb)
              batched_homo_graph.append(homo_sg)
              batched_hetero_graph.append(sg)
  
      batched_homo_graph = batch(batched_homo_graph)
      batched_hetero_graph = batch(batched_hetero_graph)
  
      batched_embed = torch.cat(batched_embed)
      batched_embed = self.elayers(batched_embed)
      values = batched_embed.reshape(-1)
  
      values = self.concrete_sample(
          values, beta=temperature, training=training
      )
      self.sparse_mask_values = values
  
      col, row = batched_homo_graph.edges()
      reverse_eids = batched_homo_graph.edge_ids(row, col).long()
      edge_mask = (values + values[reverse_eids]) / 2
  
      self.set_masks(batched_homo_graph, edge_mask)
  
      # Convert the edge mask back into heterogeneous format.
      hetero_edge_mask = self._edge_mask_to_heterogeneous(
          edge_mask=edge_mask,
          homograph=batched_homo_graph,
          heterograph=batched_hetero_graph,
      )
  
      batched_feats = {
          ntype: batched_hetero_graph.nodes[ntype].data["feat"]
          for ntype in batched_hetero_graph.ntypes
      }
  
      # The model prediction with the updated edge mask.
      logits = self.model(
          batched_hetero_graph,
          batched_feats,
          edge_weight=hetero_edge_mask,
          **kwargs,
      )
      probs = {
          ntype: F.softmax(logits[ntype], dim=-1)
          for ntype in logits.keys()
      }
  
      batched_inverse_indices = {
          ntype: batched_hetero_graph.nodes[ntype]
          .data["train"]
          .nonzero()
          .squeeze(1)
          for ntype in batched_hetero_graph.ntypes
      }
  
      if training:
          self.batched_feats = batched_feats
          probs = {ntype: probs[ntype].data for ntype in probs.keys()}
      else:
          self.clear_masks()
  
      return (
          probs,
          hetero_edge_mask,
          batched_hetero_graph,
          batched_inverse_indices,
      )
  
imports: import math
  import torch
  import torch.nn
  import torch.nn.functional
  from ....None import batch, ETYPE, khop_in_subgraph, NID, to_homogeneous
label_desc: Use image node_id 4 for calling the HeteroPGExplainer obj's underlying member method code with example usage: obj.explain_node(nodes, graph, feat, temperature, training) and returns: probs,  hetero_edge_mask,  batched_hetero_graph,  batched_inverse_indices