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: Compute the loss of the explanation network for node classification
  
  Parameters
  ----------
  nodes : dict[str, Iterable[int]]
      A dict mapping node types (keys) to an iterable set of node ids (values).
  graph : DGLGraph
      Input 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.
  kwargs : dict
      Additional arguments passed to the GNN model.
  
  Returns
  -------
  Tensor
      A scalar tensor representing the loss.
docstr_tok: ["Compute","the","loss","of","the","explanation","network","for","node","classification","Parameters","--","--","--","--","--","nodes",":","dict","[","str",",","Iterable","[","int","]","]","A","dict","mapping","node","types","(","keys",")","to","an","iterable","set","of","node","ids","(","values",")",".","graph",":","DGLGraph","Input","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",".","kwargs",":","dict","Additional","arguments","passed","to","the","GNN","model",".","Returns","--","--","--","-","Tensor","A","scalar","tensor","representing","the","loss","."]
code: def train_step_node(self, nodes, graph, feat, temperature, **kwargs):
      r"""Compute the loss of the explanation network for node classification
  
      Parameters
      ----------
      nodes : dict[str, Iterable[int]]
          A dict mapping node types (keys) to an iterable set of node ids (values).
      graph : DGLGraph
          Input 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.
      kwargs : dict
          Additional arguments passed to the GNN model.
  
      Returns
      -------
      Tensor
          A scalar tensor representing the loss.
      """
      assert (
          not self.graph_explanation
      ), '"explain_graph" must be False when initializing the module.'
  
      self.model = self.model.to(graph.device)
      self.elayers = self.elayers.to(graph.device)
  
      prob, _, batched_graph, inverse_indices = self.explain_node(
          nodes, graph, feat, temperature, training=True, **kwargs
      )
  
      pred = self.model(
          batched_graph, self.batched_feats, embed=False, **kwargs
      )
      pred = {
          ntype: pred[ntype].argmax(-1).data for ntype in pred.keys()
      }
  
      loss = self.loss(
          torch.cat(
              [
                  prob[ntype][nid]
                  for ntype, nid in inverse_indices.items()
              ]
          ),
          torch.cat(
              [
                  pred[ntype][nid]
                  for ntype, nid in inverse_indices.items()
              ]
          ),
      )
      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 2 for calling the HeteroPGExplainer obj's underlying member method code with example usage: obj.train_step_node(nodes, graph, feat, temperature) and returns: loss