class_docstr: An implementation of FlowX in
  `FlowX: Towards Explainable Graph Neural Networks via Message Flows <https://arxiv.org/abs/2206.12987>`_.
  Args:
      model (torch.nn.Module): The target model prepared to explain.
      epochs (int, optional): The training steps.
      lr (float, optional): The explainer training learning rate.
      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.
class_docstr_tok: ["An","implementation","of","FlowX","in","`","FlowX",":","Towards","Explainable","Graph","Neural","Networks","via","Message","Flows","<","https",":","\/\/arxiv.org\/abs\/2206.12987",">","`","_",".","Args",":","model","(","torch.nn.Module",")",":","The","target","model","prepared","to","explain",".","epochs","(","int",",","optional",")",":","The","training","steps",".","lr","(","float",",","optional",")",":","The","explainer","training","learning","rate",".","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","."]
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, list, tuple, torch.Tensor): 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: (None, list, list)
  .. note::
      (None, masks, related_predictions):
      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",",","list",",","tuple",",","torch.Tensor",")",":","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",":","(","None",",","list",",","list",")","..","note",":",":","(","None",",","masks",",","related_predictions",")",":","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
  ) -> Union[
      Tuple[None, List, List[Dict]], Tuple[Dict, List, List[Dict]]
  ]:
      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, list, tuple, torch.Tensor): 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: (None, list, list)
      .. note::
          (None, masks, related_predictions):
          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)
  
      # Explanation initial process
      self.model.eval()
  
      # Initial original prediction
      self.ori_logits_pred = self.model(x, edge_index).softmax(1)
      # print(f'#I#pred label: {torch.argmax(self.ori_logits_pred)}')
  
      # Edge Index with self loop
      edge_index, _ = remove_self_loops(edge_index)
      edge_index_with_loop, _ = add_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")
          self.node_idx = node_idx
          assert node_idx is not None
          _, _, _, self.hard_edge_mask = subgraph(
              node_idx,
              self.__num_hops__,
              edge_index_with_loop,
              relabel_nodes=True,
              num_nodes=None,
              flow=self.__flow__(),
          )
  
          # 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]
  
      import time
  
      start = time.time()
      self.time_list = []
      labels = tuple(i for i in range(kwargs.get("num_classes")))
      ex_labels = tuple(
          torch.tensor([label]).to(self.device) for label in labels
      )
  
      # Connect mask
      with self.connect_mask(self):
          (
              iter_weighted_change_walks_list,
              iter_changed_subsets_score_list,
              walk_sample_count,
          ) = self.flow_shap(
              x,
              edge_index,
              edge_index_with_loop,
              walk_indices_list,
              **kwargs,
          )
  
      walk_score_list = []
      for ex_label in ex_labels:
          # --- training ---
          self.train_mask(
              x,
              edge_index,
              ex_label,
              walk_indices_list,
              edge_index_with_loop,
              iter_weighted_change_walks_list,
              iter_changed_subsets_score_list,
              walk_sample_count,
          )
  
          walk_score_list.append(self.flow_mask.data)
  
      walks = {
          "ids": walk_indices_list,
          "score": torch.cat(walk_score_list, dim=1),
      }
      # print(f'Walk scores summation: {torch.cat(walk_score_list, dim=0).sum(0)}')
      # print(f'#D#WalkScore total time: {time.time() - start}\n'
      #       f'predict time: {sum(self.time_list)}')
  
      # specify to edge with self-loop mask prediction
      labels = tuple(i for i in range(kwargs.get("num_classes")))
      ex_labels = tuple(
          torch.tensor([label]).to(self.device) for label in labels
      )
      start = time.time()
      masks = []
      for ex_label in ex_labels:
          edge_attr = self.explain_edges_with_loop(x, walks, ex_label)
          mask = edge_attr
          mask = self.control_sparsity(
              mask, kwargs.get("sparsity")
          ).sigmoid()
          masks.append(mask.detach())
      # print(f'#D#Edge mask predict total time: {time.time() - start}')
  
      # Connect mask
      with self.connect_mask(self):
          related_preds = self.eval_related_pred(
              x, edge_index, masks, **kwargs
          )
  
      return walks, masks, related_preds
  
file_path: DIG/dig/xgraph/method/flowx.py
imports: from itertools import combinations
  from typing import List, Tuple, Union, Dict
  import torch
  import torch.nn
  from torch import Tensor
  from torch.nn.functional import cross_entropy
  from torch_geometric.utils.loop import add_self_loops, remove_self_loops
  from .base_explainer import WalkBase
  from ..models.utils import subgraph, gumbel_softmax
label_desc: Use image node_id 2 for calling the FlowX obj's underlying member method code with example usage: obj.forward(x, edge_index) and returns: walks,  masks,  related_preds