class_docstr: PGExplainer from `Parameterized Explainer for Graph Neural Network
  <https://arxiv.org/pdf/2011.04573>`
  
  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`.
  num_hops : int, optional
      The number of hops for GNN information aggregation, which must match the
      number of message passing layers employed by the GNN to be explained.
  explain_graph : bool, optional
      Whether to initialize the model for graph-level or node-level predictions.
  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",">","`","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","`",".","num_hops",":","int",",","optional","The","number","of","hops","for","GNN","information","aggregation",",","which","must","match","the","number","of","message","passing","layers","employed","by","the","GNN","to","be","explained",".","explain_graph",":","bool",",","optional","Whether","to","initialize","the","model","for","graph-level","or","node-level","predictions",".","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 graph and edge mask.
  
  Parameters
  ----------
  nodes : int, iterable[int], tensor
      The nodes from the graph, which cannot have any duplicate value.
  graph : DGLGraph
      A homogeneous graph.
  feat : Tensor
      The input feature of shape :math:`(N, D)`. :math:`N` is the
      number of nodes, and :math:`D` is the feature size.
  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:`(N, L)`, where :math:`L` is the different types
      of node labels in the dataset, and :math:`N` is the number of nodes
      in the graph.
  Tensor
      Edge weights which is a tensor of shape :math:`(E)`, where :math:`E`
      is the number of edges in the graph. 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.
  Tensor
      The new IDs of the subgraph center nodes.
  
  Examples
  --------
  
  >>> import dgl
  >>> import numpy as np
  >>> import torch
  
  >>> # Define the model
  >>> class Model(torch.nn.Module):
  ...     def __init__(self, in_feats, out_feats):
  ...         super().__init__()
  ...         self.conv1 = dgl.nn.GraphConv(in_feats, out_feats)
  ...         self.conv2 = dgl.nn.GraphConv(out_feats, out_feats)
  ...
  ...     def forward(self, g, h, embed=False, edge_weight=None):
  ...         h = self.conv1(g, h, edge_weight=edge_weight)
  ...         if embed:
  ...             return h
  ...         return self.conv2(g, h)
  
  >>> # Load dataset
  >>> data = dgl.data.CoraGraphDataset(verbose=False)
  >>> g = data[0]
  >>> features = g.ndata["feat"]
  >>> labels = g.ndata["label"]
  
  >>> # Train the model
  >>> model = Model(features.shape[1], data.num_classes)
  >>> criterion = torch.nn.CrossEntropyLoss()
  >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
  >>> for epoch in range(20):
  ...     logits = model(g, features)
  ...     loss = criterion(logits, labels)
  ...     optimizer.zero_grad()
  ...     loss.backward()
  ...     optimizer.step()
  
  >>> # Initialize the explainer
  >>> explainer = dgl.nn.PGExplainer(
  ...     model, data.num_classes, num_hops=2, explain_graph=False
  ... )
  
  >>> # Train the explainer
  >>> # Define explainer temperature parameter
  >>> init_tmp, final_tmp = 5.0, 1.0
  >>> optimizer_exp = torch.optim.Adam(explainer.parameters(), lr=0.01)
  >>> epochs = 10
  >>> for epoch in range(epochs):
  ...     tmp = float(init_tmp * np.power(final_tmp / init_tmp, epoch / epochs))
  ...     loss = explainer.train_step_node(g.nodes(), g, features, tmp)
  ...     optimizer_exp.zero_grad()
  ...     loss.backward()
  ...     optimizer_exp.step()
  
  >>> # Explain the prediction for graph 0
  >>> probs, edge_weight, bg, inverse_indices = explainer.explain_node(
  ...     0, g, features
  ... )
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","graph","and","edge","mask",".","Parameters","--","--","--","--","--","nodes",":","int",",","iterable","[","int","]",",","tensor","The","nodes","from","the","graph",",","which","can","not","have","any","duplicate","value",".","graph",":","DGLGraph","A","homogeneous","graph",".","feat",":","Tensor","The","input","feature","of","shape",":","math",":","`","(","N",",","D",")","`",".",":","math",":","`","N","`","is","the","number","of","nodes",",","and",":","math",":","`","D","`","is","the","feature","size",".","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",":","`","(","N",",","L",")","`",",","where",":","math",":","`","L","`","is","the","different","types","of","node","labels","in","the","dataset",",","and",":","math",":","`","N","`","is","the","number","of","nodes","in","the","graph",".","Tensor","Edge","weights","which","is","a","tensor","of","shape",":","math",":","`","(","E",")","`",",","where",":","math",":","`","E","`","is","the","number","of","edges","in","the","graph",".","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",".","Tensor","The","new","IDs","of","the","subgraph","center","nodes",".","Examples","--","--","--","--",">",">",">","import","dgl",">",">",">","import","numpy","as","np",">",">",">","import","torch",">",">",">","#","Define","the","model",">",">",">","class","Model","(","torch.nn.Module",")",":","...","def","__init__","(","self",",","in_feats",",","out_feats",")",":","...","super","(",")",".__init__","(",")","...","self.conv1","=","dgl.nn.GraphConv","(","in_feats",",","out_feats",")","...","self.conv2","=","dgl.nn.GraphConv","(","out_feats",",","out_feats",")","...","...","def","forward","(","self",",","g",",","h",",","embed=False",",","edge_weight=None",")",":","...","h","=","self.conv1","(","g",",","h",",","edge_weight=edge_weight",")","...","if","embed",":","...","return","h","...","return","self.conv2","(","g",",","h",")",">",">",">","#","Load","dataset",">",">",">","data","=","dgl.data.CoraGraphDataset","(","verbose=False",")",">",">",">","g","=","data","[","0","]",">",">",">","features","=","g.ndata","[","``","feat","''","]",">",">",">","labels","=","g.ndata","[","``","label","''","]",">",">",">","#","Train","the","model",">",">",">","model","=","Model","(","features.shape","[","1","]",",","data.num_classes",")",">",">",">","criterion","=","torch.nn.CrossEntropyLoss","(",")",">",">",">","optimizer","=","torch.optim.Adam","(","model.parameters","(",")",",","lr=1e-2",")",">",">",">","for","epoch","in","range","(","20",")",":","...","logits","=","model","(","g",",","features",")","...","loss","=","criterion","(","logits",",","labels",")","...","optimizer.zero_grad","(",")","...","loss.backward","(",")","...","optimizer.step","(",")",">",">",">","#","Initialize","the","explainer",">",">",">","explainer","=","dgl.nn.PGExplainer","(","...","model",",","data.num_classes",",","num_hops=2",",","explain_graph=False","...",")",">",">",">","#","Train","the","explainer",">",">",">","#","Define","explainer","temperature","parameter",">",">",">","init_tmp",",","final_tmp","=","5.0",",","1.0",">",">",">","optimizer_exp","=","torch.optim.Adam","(","explainer.parameters","(",")",",","lr=0.01",")",">",">",">","epochs","=","10",">",">",">","for","epoch","in","range","(","epochs",")",":","...","tmp","=","float","(","init_tmp","*","np.power","(","final_tmp","\/","init_tmp",",","epoch","\/","epochs",")",")","...","loss","=","explainer.train_step_node","(","g.nodes","(",")",",","g",",","features",",","tmp",")","...","optimizer_exp.zero_grad","(",")","...","loss.backward","(",")","...","optimizer_exp.step","(",")",">",">",">","#","Explain","the","prediction","for","graph","0",">",">",">","probs",",","edge_weight",",","bg",",","inverse_indices","=","explainer.explain_node","(","...","0",",","g",",","features","...",")"]
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 graph and edge mask.
  
      Parameters
      ----------
      nodes : int, iterable[int], tensor
          The nodes from the graph, which cannot have any duplicate value.
      graph : DGLGraph
          A homogeneous graph.
      feat : Tensor
          The input feature of shape :math:`(N, D)`. :math:`N` is the
          number of nodes, and :math:`D` is the feature size.
      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:`(N, L)`, where :math:`L` is the different types
          of node labels in the dataset, and :math:`N` is the number of nodes
          in the graph.
      Tensor
          Edge weights which is a tensor of shape :math:`(E)`, where :math:`E`
          is the number of edges in the graph. 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.
      Tensor
          The new IDs of the subgraph center nodes.
  
      Examples
      --------
  
      >>> import dgl
      >>> import numpy as np
      >>> import torch
  
      >>> # Define the model
      >>> class Model(torch.nn.Module):
      ...     def __init__(self, in_feats, out_feats):
      ...         super().__init__()
      ...         self.conv1 = dgl.nn.GraphConv(in_feats, out_feats)
      ...         self.conv2 = dgl.nn.GraphConv(out_feats, out_feats)
      ...
      ...     def forward(self, g, h, embed=False, edge_weight=None):
      ...         h = self.conv1(g, h, edge_weight=edge_weight)
      ...         if embed:
      ...             return h
      ...         return self.conv2(g, h)
  
      >>> # Load dataset
      >>> data = dgl.data.CoraGraphDataset(verbose=False)
      >>> g = data[0]
      >>> features = g.ndata["feat"]
      >>> labels = g.ndata["label"]
  
      >>> # Train the model
      >>> model = Model(features.shape[1], data.num_classes)
      >>> criterion = torch.nn.CrossEntropyLoss()
      >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
      >>> for epoch in range(20):
      ...     logits = model(g, features)
      ...     loss = criterion(logits, labels)
      ...     optimizer.zero_grad()
      ...     loss.backward()
      ...     optimizer.step()
  
      >>> # Initialize the explainer
      >>> explainer = dgl.nn.PGExplainer(
      ...     model, data.num_classes, num_hops=2, explain_graph=False
      ... )
  
      >>> # Train the explainer
      >>> # Define explainer temperature parameter
      >>> init_tmp, final_tmp = 5.0, 1.0
      >>> optimizer_exp = torch.optim.Adam(explainer.parameters(), lr=0.01)
      >>> epochs = 10
      >>> for epoch in range(epochs):
      ...     tmp = float(init_tmp * np.power(final_tmp / init_tmp, epoch / epochs))
      ...     loss = explainer.train_step_node(g.nodes(), g, features, tmp)
      ...     optimizer_exp.zero_grad()
      ...     loss.backward()
      ...     optimizer_exp.step()
  
      >>> # Explain the prediction for graph 0
      >>> probs, edge_weight, bg, inverse_indices = explainer.explain_node(
      ...     0, g, features
      ... )
      """
      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.'
  
      if isinstance(nodes, torch.Tensor):
          nodes = nodes.tolist()
      if isinstance(nodes, int):
          nodes = [nodes]
  
      self.model = self.model.to(graph.device)
      self.elayers = self.elayers.to(graph.device)
  
      batched_graph = []
      batched_embed = []
      for node_id in nodes:
          sg, inverse_indices = khop_in_subgraph(
              graph, node_id, self.num_hops
          )
          sg.ndata["feat"] = feat[sg.ndata[NID].long()]
          sg.ndata["train"] = torch.tensor(
              [nid in inverse_indices for nid in sg.nodes()],
              device=sg.device,
          )
  
          embed = self.model(sg, sg.ndata["feat"], embed=True, **kwargs)
          embed = embed.data
  
          col, row = sg.edges()
          col_emb = embed[col.long()]
          row_emb = embed[row.long()]
          self_emb = embed[inverse_indices[0]].repeat(sg.num_edges(), 1)
          emb = torch.cat([col_emb, row_emb, self_emb], dim=-1)
          batched_embed.append(emb)
          batched_graph.append(sg)
  
      batched_graph = batch(batched_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_graph.edges()
      reverse_eids = batched_graph.edge_ids(row, col).long()
      edge_mask = (values + values[reverse_eids]) / 2
  
      self.set_masks(batched_graph, edge_mask)
  
      batched_feats = batched_graph.ndata["feat"]
      # the model prediction with the updated edge mask
      logits = self.model(
          batched_graph,
          batched_feats,
          edge_weight=self.edge_mask,
          **kwargs,
      )
      probs = F.softmax(logits, dim=-1)
  
      batched_inverse_indices = (
          batched_graph.ndata["train"].nonzero().squeeze(1)
      )
  
      if training:
          self.batched_feats = batched_feats
          probs = probs.data
      else:
          self.clear_masks()
  
      return (
          probs,
          edge_mask,
          batched_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 10 for calling the PGExplainer obj's underlying member method code with example usage: obj.explain_node(nodes, graph, feat, temperature, training) and returns: probs,  edge_mask,  batched_graph,  batched_inverse_indices