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: Set the edge mask that plays a crucial role to explain the
  prediction made by the GNN for a graph. Initialize learnable edge
  mask if it is None.
  
  Parameters
  ----------
  graph : DGLGraph
      A homogeneous graph.
  edge_mask : Tensor, optional
      Learned importance mask of the edges in the graph, which is a tensor
      of shape :math:`(E)`, where :math:`E` is the number of edges in the
      graph. The values are within range :math:`(0, 1)`. The higher,
      the more important. Default: None.
docstr_tok: ["Set","the","edge","mask","that","plays","a","crucial","role","to","explain","the","prediction","made","by","the","GNN","for","a","graph",".","Initialize","learnable","edge","mask","if","it","is","None",".","Parameters","--","--","--","--","--","graph",":","DGLGraph","A","homogeneous","graph",".","edge_mask",":","Tensor",",","optional","Learned","importance","mask","of","the","edges","in","the","graph",",","which","is","a","tensor","of","shape",":","math",":","`","(","E",")","`",",","where",":","math",":","`","E","`","is","the","number","of","edges","in","the","graph",".","The","values","are","within","range",":","math",":","`","(","0",",","1",")","`",".","The","higher",",","the","more","important",".","Default",":","None","."]
code: def set_masks(self, graph, edge_mask=None):
      r"""Set the edge mask that plays a crucial role to explain the
      prediction made by the GNN for a graph. Initialize learnable edge
      mask if it is None.
  
      Parameters
      ----------
      graph : DGLGraph
          A homogeneous graph.
      edge_mask : Tensor, optional
          Learned importance mask of the edges in the graph, which is a tensor
          of shape :math:`(E)`, where :math:`E` is the number of edges in the
          graph. The values are within range :math:`(0, 1)`. The higher,
          the more important. Default: None.
      """
      if edge_mask is None:
          num_nodes = graph.num_nodes()
          num_edges = graph.num_edges()
  
          init_bias = self.init_bias
          std = nn.init.calculate_gain("relu") * math.sqrt(
              2.0 / (2 * num_nodes)
          )
          self.edge_mask = torch.randn(num_edges) * std + init_bias
      else:
          self.edge_mask = edge_mask
  
      self.edge_mask = self.edge_mask.to(graph.device)
  
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 2 for calling the PGExplainer obj's underlying member method code with example usage: obj.set_masks(graph, edge_mask) without return types