class_docstr: An implementation of GNN-LRP in
  `Higher-Order Explanations of Graph Neural Networks via Relevant Walks <https://arxiv.org/abs/2006.03589>`_.
  Args:
      model (torch.nn.Module): The target model prepared to explain.
      explain_graph (bool, optional): Whether to explain graph classification model.
          (default: :obj:`False`)
  .. note::
          For node classification model, the :attr:`explain_graph` flag is False.
          GNN-LRP is very model dependent. Please be sure you know how to modify it for different models.
          For an example, see `benchmarks/xgraph
          <https://github.com/divelab/DIG/tree/dig/benchmarks/xgraph>`_.
class_docstr_tok: ["An","implementation","of","GNN-LRP","in","`","Higher-Order","Explanations","of","Graph","Neural","Networks","via","Relevant","Walks","<","https",":","\/\/arxiv.org\/abs\/2006.03589",">","`","_",".","Args",":","model","(","torch.nn.Module",")",":","The","target","model","prepared","to","explain",".","explain_graph","(","bool",",","optional",")",":","Whether","to","explain","graph","classification","model",".","(","default",":",":","obj",":","`","False","`",")","..","note",":",":","For","node","classification","model",",","the",":","attr",":","`","explain_graph","`","flag","is","False",".","GNN-LRP","is","very","model","dependent",".","Please","be","sure","you","know","how","to","modify","it","for","different","models",".","For","an","example",",","see","`","benchmarks\/xgraph","<","https",":","\/\/github.com\/divelab\/DIG\/tree\/dig\/benchmarks\/xgraph",">","`","_","."]
docstr: Run the explainer for a specific graph instance.
  Args:
      x (torch.Tensor): The graph instance's input node features.
      edge_index (torch.Tensor): The graph instance's edge index.
      **kwargs (dict):
          :obj:`node_idx` （int): The index of node that is pending to be explained.
          (for node classification)
          :obj:`sparsity` (float): The Sparsity we need to control to transform a
          soft mask to a hard mask. (Default: :obj:`0.7`)
          :obj:`num_classes` (int): The number of task's classes.
  :rtype:
      (walks, edge_masks, related_predictions),
      walks is a dictionary including walks' edge indices and corresponding explained scores;
      edge_masks is a list of edge-level explanation for each class;
      related_predictions is a list of dictionary for each class
      where each dictionary includes 4 type predicted probabilities.
docstr_tok: ["Run","the","explainer","for","a","specific","graph","instance",".","Args",":","x","(","torch.Tensor",")",":","The","graph","instance","'s","input","node","features",".","edge_index","(","torch.Tensor",")",":","The","graph","instance","'s","edge","index",".","*","*","kwargs","(","dict",")",":",":","obj",":","`","node_idx","`","\uff08int",")",":","The","index","of","node","that","is","pending","to","be","explained",".","(","for","node","classification",")",":","obj",":","`","sparsity","`","(","float",")",":","The","Sparsity","we","need","to","control","to","transform","a","soft","mask","to","a","hard","mask",".","(","Default",":",":","obj",":","`","0.7","`",")",":","obj",":","`","num_classes","`","(","int",")",":","The","number","of","task","'s","classes",".",":","rtype",":","(","walks",",","edge_masks",",","related_predictions",")",",","walks","is","a","dictionary","including","walks","'","edge","indices","and","corresponding","explained","scores",";","edge_masks","is","a","list","of","edge-level","explanation","for","each","class",";","related_predictions","is","a","list","of","dictionary","for","each","class","where","each","dictionary","includes","4","type","predicted","probabilities","."]
code: def forward(self, x: Tensor, edge_index: Tensor, **kwargs):
      r"""
      Run the explainer for a specific graph instance.
      Args:
          x (torch.Tensor): The graph instance's input node features.
          edge_index (torch.Tensor): The graph instance's edge index.
          **kwargs (dict):
              :obj:`node_idx` （int): The index of node that is pending to be explained.
              (for node classification)
              :obj:`sparsity` (float): The Sparsity we need to control to transform a
              soft mask to a hard mask. (Default: :obj:`0.7`)
              :obj:`num_classes` (int): The number of task's classes.
      :rtype:
          (walks, edge_masks, related_predictions),
          walks is a dictionary including walks' edge indices and corresponding explained scores;
          edge_masks is a list of edge-level explanation for each class;
          related_predictions is a list of dictionary for each class
          where each dictionary includes 4 type predicted probabilities.
      """
      super().forward(x, edge_index, **kwargs)
      labels = tuple(i for i in range(kwargs.get("num_classes")))
      self.model.eval()
  
      walk_steps, fc_steps = self.extract_step(
          x, edge_index, detach=False, split_fc=True
      )
  
      edge_index_with_loop, _ = add_remaining_self_loops(
          edge_index, num_nodes=self.num_nodes
      )
  
      walk_indices_list = torch.tensor(
          self.walks_pick(
              edge_index_with_loop.cpu(),
              list(range(edge_index_with_loop.shape[1])),
              num_layers=self.num_layers,
          ),
          device=self.device,
      )
      if not self.explain_graph:
          node_idx = kwargs.get("node_idx")
          node_idx = node_idx.reshape([1]).to(self.device)
          assert node_idx is not None
          self.subset, _, _, self.hard_edge_mask = subgraph(
              node_idx,
              self.__num_hops__,
              edge_index_with_loop,
              relabel_nodes=True,
              num_nodes=None,
              flow=self.__flow__(),
          )
          self.new_node_idx = torch.where(self.subset == node_idx)[0]
  
          # walk indices list mask
          edge2node_idx = edge_index_with_loop[1] == node_idx
          walk_indices_list_mask = edge2node_idx[
              walk_indices_list[:, -1]
          ]
          walk_indices_list = walk_indices_list[walk_indices_list_mask]
  
      if kwargs.get("walks"):
          walks = kwargs.pop("walks")
  
      else:
  
          def compute_walk_score():
              # hyper-parameter gamma
              epsilon = 1e-30  # prevent from zero division
              gamma = [2, 1, 1]
  
              # --- record original weights of GNN ---
              ori_gnn_weights = []
              gnn_gamma_modules = []
              clear_probe = x
              for i, walk_step in enumerate(walk_steps):
                  modules = walk_step["module"]
                  gamma_ = gamma[i] if i <= 1 else 1
                  if hasattr(modules[0], "nn"):
                      clear_probe = modules[0](
                          clear_probe, edge_index, probe=False
                      )
                      # clear nodes that are not created by user
                  gamma_module = copy.deepcopy(modules[0])
                  if hasattr(modules[0], "nn"):
                      for j, fc_step in enumerate(
                          gamma_module.fc_steps
                      ):
                          fc_modules = fc_step["module"]
                          if hasattr(fc_modules[0], "weight"):
                              ori_fc_weight = fc_modules[0].weight.data
                              fc_modules[0].weight.data = (
                                  ori_fc_weight + gamma_ * ori_fc_weight
                              )
                  else:
                      ori_gnn_weights.append(modules[0].weight.data)
                      gamma_module.weight.data = (
                          ori_gnn_weights[i]
                          + gamma_ * ori_gnn_weights[i].relu()
                      )
                  gnn_gamma_modules.append(gamma_module)
  
              # --- record original weights of fc layer ---
              ori_fc_weights = []
              fc_gamma_modules = []
              for i, fc_step in enumerate(fc_steps):
                  modules = fc_step["module"]
                  gamma_module = copy.deepcopy(modules[0])
                  if hasattr(modules[0], "weight"):
                      ori_fc_weights.append(modules[0].weight.data)
                      gamma_ = 1
                      gamma_module.weight.data = (
                          ori_fc_weights[i]
                          + gamma_ * ori_fc_weights[i].relu()
                      )
                  else:
                      ori_fc_weights.append(None)
                  fc_gamma_modules.append(gamma_module)
  
              # --- GNN_LRP implementation ---
              for walk_indices in walk_indices_list:
                  walk_node_indices = [
                      edge_index_with_loop[0, walk_indices[0]]
                  ]
                  for walk_idx in walk_indices:
                      walk_node_indices.append(
                          edge_index_with_loop[1, walk_idx]
                      )
  
                  h = x.requires_grad_(True)
                  for i, walk_step in enumerate(walk_steps):
                      modules = walk_step["module"]
                      if hasattr(modules[0], "nn"):
                          # for the specific 2-layer nn GINs.
                          gin = modules[0]
                          run1 = gin(h, edge_index, probe=True)
                          std_h1 = gin.fc_steps[0]["output"]
                          gamma_run1 = gnn_gamma_modules[i](
                              h, edge_index, probe=True
                          )
                          p1 = gnn_gamma_modules[i].fc_steps[0][
                              "output"
                          ]
                          q1 = (p1 + epsilon) * (
                              std_h1 / (p1 + epsilon)
                          ).detach()
  
                          std_h2 = GraphSequential(
                              *gin.fc_steps[1]["module"]
                          )(q1)
                          p2 = GraphSequential(
                              *gnn_gamma_modules[i].fc_steps[1][
                                  "module"
                              ]
                          )(q1)
                          q2 = (p2 + epsilon) * (
                              std_h2 / (p2 + epsilon)
                          ).detach()
                          q = q2
                      else:
                          std_h = GraphSequential(*modules)(
                              h, edge_index
                          )
  
                          # --- LRP-gamma ---
                          p = gnn_gamma_modules[i](h, edge_index)
                          q = (p + epsilon) * (
                              std_h / (p + epsilon)
                          ).detach()
  
                      # --- pick a path ---
                      mk = torch.zeros(
                          (h.shape[0], 1), device=self.device
                      )
                      k = walk_node_indices[i + 1]
                      mk[k] = 1
                      ht = q * mk + q.detach() * (1 - mk)
                      h = ht
  
                  # --- FC LRP_gamma ---
                  # debug that torch.zeros(h.shape[0], dtype=torch.long, device=self.device)
                  # should be an edge_index with [num_edge, 2]
                  for i, fc_step in enumerate(fc_steps):
                      modules = fc_step["module"]
                      std_h = (
                          nn.Sequential(*modules)(h)
                          if i != 0
                          else GraphSequential(*modules)(
                              h,
                              torch.zeros(
                                  h.shape[0],
                                  dtype=torch.long,
                                  device=self.device,
                              ),
                          )
                      )
  
                      # --- gamma ---
                      s = (
                          fc_gamma_modules[i](h)
                          if i != 0
                          else fc_gamma_modules[i](
                              h,
                              torch.zeros(
                                  h.shape[0],
                                  dtype=torch.long,
                                  device=self.device,
                              ),
                          )
                      )
                      ht = (s + epsilon) * (
                          std_h / (s + epsilon)
                      ).detach()
                      h = ht
  
                  if not self.explain_graph:
                      f = h[node_idx, label]
                  else:
                      f = h[0, label]
                  x_grads = torch.autograd.grad(outputs=f, inputs=x)[0]
                  I = walk_node_indices[0]
                  r = x_grads[I, :] @ x[I].T
                  walk_scores.append(r)
  
          walk_scores_tensor_list = [None for i in labels]
          for label in labels:
              walk_scores = []
  
              compute_walk_score()
              walk_scores_tensor_list[label] = torch.stack(
                  walk_scores, dim=0
              ).view(-1, 1)
  
          walks = {
              "ids": walk_indices_list,
              "score": torch.cat(walk_scores_tensor_list, dim=1),
          }
  
      # --- Apply edge mask evaluation ---
      with torch.no_grad():
          with self.connect_mask(self):
              ex_labels = tuple(
                  torch.tensor([label]).to(self.device)
                  for label in labels
              )
              edge_masks = []
              hard_edge_masks = []
              for ex_label in ex_labels:
                  edge_attr = self.explain_edges_with_loop(
                      x, walks, ex_label
                  )
                  edge_mask = edge_attr.detach()
                  valid_mask = edge_mask != -math.inf
                  edge_mask[edge_mask == -math.inf] = (
                      edge_mask[valid_mask].min() - 1
                  )  # replace the negative inf
  
                  edge_masks.append(edge_mask)
                  hard_edge_masks.append(
                      self.control_sparsity(
                          edge_attr, kwargs.get("sparsity")
                      ).sigmoid()
                  )
  
              related_preds = self.eval_related_pred(
                  x, edge_index, hard_edge_masks, **kwargs
              )
  
      return walks, edge_masks, related_preds
  
file_path: DIG/dig/xgraph/method/gnn_lrp.py
imports: import math
  import torch
  from torch import Tensor
  import torch.nn
  import copy
  from torch_geometric.utils.loop import add_remaining_self_loops
  from ..models.utils import subgraph
  from ..models.models import GraphSequential
  from .base_explainer import WalkBase
label_desc: Use image node_id 2 for calling the GNN_LRP obj's underlying member method code with example usage: obj.forward(x, edge_index) and returns: walks,  edge_masks,  related_preds