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: The loss function that is used to learn the edge
  distribution.
  
  Parameters
  ----------
  prob: Tensor
      Tensor contains a set of probabilities for each possible
      class label of some model for all the batched graphs,
      which is of shape :math:`(B, L)`, where :math:`L` is the
      different types of label in the dataset and :math:`B` is
      the batch size.
  ori_pred: Tensor
      Tensor of shape :math:`(B, 1)`, representing the original prediction
      for the graph, where :math:`B` is the batch size.
  
  Returns
  -------
  float
      The function that returns the sum of the three loss components,
      which is a scalar tensor representing the total loss.
docstr_tok: ["The","loss","function","that","is","used","to","learn","the","edge","distribution",".","Parameters","--","--","--","--","--","prob",":","Tensor","Tensor","contains","a","set","of","probabilities","for","each","possible","class","label","of","some","model","for","all","the","batched","graphs",",","which","is","of","shape",":","math",":","`","(","B",",","L",")","`",",","where",":","math",":","`","L","`","is","the","different","types","of","label","in","the","dataset","and",":","math",":","`","B","`","is","the","batch","size",".","ori_pred",":","Tensor","Tensor","of","shape",":","math",":","`","(","B",",","1",")","`",",","representing","the","original","prediction","for","the","graph",",","where",":","math",":","`","B","`","is","the","batch","size",".","Returns","--","--","--","-","float","The","function","that","returns","the","sum","of","the","three","loss","components",",","which","is","a","scalar","tensor","representing","the","total","loss","."]
code: def loss(self, prob, ori_pred):
      r"""The loss function that is used to learn the edge
      distribution.
  
      Parameters
      ----------
      prob: Tensor
          Tensor contains a set of probabilities for each possible
          class label of some model for all the batched graphs,
          which is of shape :math:`(B, L)`, where :math:`L` is the
          different types of label in the dataset and :math:`B` is
          the batch size.
      ori_pred: Tensor
          Tensor of shape :math:`(B, 1)`, representing the original prediction
          for the graph, where :math:`B` is the batch size.
  
      Returns
      -------
      float
          The function that returns the sum of the three loss components,
          which is a scalar tensor representing the total loss.
      """
      target_prob = prob.gather(-1, ori_pred.unsqueeze(-1))
      # 1e-6 added to prob to avoid taking the logarithm of zero
      target_prob += 1e-6
      # computing the log likelihood for a single prediction
      pred_loss = torch.mean(-torch.log(target_prob))
  
      # size
      edge_mask = self.sparse_mask_values
      if self.coff_budget <= 0:
          size_loss = self.coff_budget * torch.sum(edge_mask)
      else:
          size_loss = self.coff_budget * F.relu(
              torch.sum(edge_mask) - self.coff_budget
          )
  
      # entropy
      scale = 0.99
      edge_mask = self.edge_mask * (2 * scale - 1.0) + (1.0 - scale)
      mask_ent = -edge_mask * torch.log(edge_mask) - (
          1 - edge_mask
      ) * torch.log(1 - edge_mask)
      mask_ent_loss = self.coff_connect * torch.mean(mask_ent)
  
      loss = pred_loss + size_loss + mask_ent_loss
      return loss
  
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 5 for calling the PGExplainer obj's underlying member method code with example usage: obj.loss(prob, ori_pred) and returns: loss