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 a graph. Also, return
  the prediction made with the edges chosen based on the edge mask.
  
  Parameters
  ----------
  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
  -------
  Tensor
      Classification probabilities given the masked graph. It is a tensor
      of shape :math:`(B, L)`, where :math:`L` is the different types of
      label in the dataset, and :math:`B` is the batch size.
  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.
  
  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)
  ...
  ...         if embed:
  ...             return h
  ...
  ...         with g.local_scope():
  ...             g.ndata["h"] = h
  ...             hg = 0
  ...             for ntype in g.ntypes:
  ...                 hg = hg + dgl.mean_nodes(g, "h", ntype=ntype)
  ...             return self.fc(hg)
  
  >>> # 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"])
  ...     loss = th.nn.functional.cross_entropy(logits, th.tensor([1]))
  ...     optimizer.zero_grad()
  ...     loss.backward()
  ...     optimizer.step()
  
  >>> # Initialize the explainer
  >>> explainer = dgl.nn.HeteroPGExplainer(model, hidden_dim)
  
  >>> # 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(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 = explainer.explain_graph(g, feat)
docstr_tok: ["Learn","and","return","an","edge","mask","that","plays","a","crucial","role","to","explain","the","prediction","made","by","the","GNN","for","a","graph",".","Also",",","return","the","prediction","made","with","the","edges","chosen","based","on","the","edge","mask",".","Parameters","--","--","--","--","--","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","--","--","--","-","Tensor","Classification","probabilities","given","the","masked","graph",".","It","is","a","tensor","of","shape",":","math",":","`","(","B",",","L",")","`",",","where",":","math",":","`","L","`","is","the","different","types","of","label","in","the","dataset",",","and",":","math",":","`","B","`","is","the","batch","size",".","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",".","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",")","...","...","if","embed",":","...","return","h","...","...","with","g.local_scope","(",")",":","...","g.ndata","[","``","h","''","]","=","h","...","hg","=","0","...","for","ntype","in","g.ntypes",":","...","hg","=","hg","+","dgl.mean_nodes","(","g",",","``","h","''",",","ntype=ntype",")","...","return","self.fc","(","hg",")",">",">",">","#","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","''","]",")","...","loss","=","th.nn.functional.cross_entropy","(","logits",",","th.tensor","(","[","1","]",")",")","...","optimizer.zero_grad","(",")","...","loss.backward","(",")","...","optimizer.step","(",")",">",">",">","#","Initialize","the","explainer",">",">",">","explainer","=","dgl.nn.HeteroPGExplainer","(","model",",","hidden_dim",")",">",">",">","#","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","(","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","=","explainer.explain_graph","(","g",",","feat",")"]
code: def explain_graph(
      self, 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 a graph. Also, return
      the prediction made with the edges chosen based on the edge mask.
  
      Parameters
      ----------
      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
      -------
      Tensor
          Classification probabilities given the masked graph. It is a tensor
          of shape :math:`(B, L)`, where :math:`L` is the different types of
          label in the dataset, and :math:`B` is the batch size.
      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.
  
      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)
      ...
      ...         if embed:
      ...             return h
      ...
      ...         with g.local_scope():
      ...             g.ndata["h"] = h
      ...             hg = 0
      ...             for ntype in g.ntypes:
      ...                 hg = hg + dgl.mean_nodes(g, "h", ntype=ntype)
      ...             return self.fc(hg)
  
      >>> # 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"])
      ...     loss = th.nn.functional.cross_entropy(logits, th.tensor([1]))
      ...     optimizer.zero_grad()
      ...     loss.backward()
      ...     optimizer.step()
  
      >>> # Initialize the explainer
      >>> explainer = dgl.nn.HeteroPGExplainer(model, hidden_dim)
  
      >>> # 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(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 = explainer.explain_graph(g, feat)
      """
      assert (
          self.graph_explanation
      ), '"explain_graph" must be True when initializing the module.'
  
      self.model = self.model.to(graph.device)
      self.elayers = self.elayers.to(graph.device)
  
      embed = self.model(graph, feat, embed=True, **kwargs)
      for ntype, emb in embed.items():
          graph.nodes[ntype].data["emb"] = emb.data
      homo_graph = to_homogeneous(graph, ndata=["emb"])
      homo_embed = homo_graph.ndata["emb"]
  
      col, row = homo_graph.edges()
      col_emb = homo_embed[col.long()]
      row_emb = homo_embed[row.long()]
      emb = torch.cat([col_emb, row_emb], dim=-1)
      emb = self.elayers(emb)
      values = emb.reshape(-1)
  
      values = self.concrete_sample(
          values, beta=temperature, training=training
      )
      self.sparse_mask_values = values
  
      reverse_eids = homo_graph.edge_ids(row, col).long()
      edge_mask = (values + values[reverse_eids]) / 2
  
      self.set_masks(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=homo_graph,
          heterograph=graph,
      )
  
      # the model prediction with the updated edge mask
      logits = self.model(
          graph, feat, edge_weight=hetero_edge_mask, **kwargs
      )
      probs = F.softmax(logits, dim=-1)
  
      if training:
          probs = probs.data
      else:
          self.clear_masks()
  
      return (probs, hetero_edge_mask)
  
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 3 for calling the HeteroPGExplainer obj's underlying member method code with example usage: obj.explain_graph(graph, feat, temperature, training) and returns: probs,  hetero_edge_mask