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: Flow shapley calculations.
docstr_tok: ["Flow","shapley","calculations","."]
code: def flow_shap(
      self,
      x,
      edge_index,
      edge_index_with_loop,
      walk_indices_list,
      **kwargs,
  ):
      r"""
      Flow shapley calculations.
      """
      # --- Kernel algorithm ---
  
      # --- row: walk index, column: total score
      walk_sample_count = torch.zeros(
          walk_indices_list.shape[0],
          dtype=torch.float,
          device=self.device,
      )
  
      # General setting
      iter_weighted_change_walks_list = []
      iter_changed_subsets_score_list = []
  
      # Random sample ns_iter * ns_per_iter times
      for iter_idx in range(self.ns_iter):
          # --- random index ---
          unmask_pool = torch.cat(
              [
                  walk_indices_list[:, layer].unique()
                  + layer * edge_index_with_loop.shape[1]
                  for layer in range(self.num_layers)
              ]
          )
          self.ns_per_iter = (
              unmask_pool.shape[0]
              if self.explain_graph or unmask_pool.shape[0] <= 100
              else 100
          )
          idx = torch.randperm(unmask_pool.nelement())
          unmask_pool = unmask_pool.view(-1)[idx].view(
              unmask_pool.size()
          )
  
          mask_per_sub = unmask_pool.shape[0] // self.ns_per_iter
          weighted_change_walks_list = []
          last_eliminated_walks = torch.zeros(
              walk_indices_list.shape[0],
              dtype=torch.bool,
              device=self.device,
          )
          layer_edge_mask_list = []
          for sub_idx in range(self.ns_per_iter):
              # --- sub random index ---
              mask_pool = unmask_pool[: mask_per_sub * (sub_idx + 1)]
  
              # --- obtain changed walk idx ---
              eliminated_layer_edges = unmask_pool[
                  mask_per_sub * sub_idx : mask_per_sub * (sub_idx + 1)
              ]
              walk_plain_indices_list = walk_indices_list + (
                  edge_index_with_loop.shape[1]
                  * torch.arange(self.num_layers, device=self.device)
              ).repeat(walk_indices_list.shape[0], 1)
              eliminated_walks = (
                  torch.stack(
                      [
                          walk_plain_indices_list == edge
                          for edge in eliminated_layer_edges
                      ],
                      dim=0,
                  )
                  .long()
                  .sum(0)
                  .sum(1)
                  .bool()
                  .long()
              )
              weighted_changed_walks = eliminated_walks.clone().float()
              weighted_changed_walks[
                  eliminated_walks == last_eliminated_walks
              ] = 0.0
              weighted_changed_walks /= (
                  weighted_changed_walks > 1e-20
              ).sum() + 1e-30
              weighted_change_walks_list.append(weighted_changed_walks)
              last_eliminated_walks = eliminated_walks
  
              # --- setting a subset mask ---
              layer_edge_masks = torch.ones(
                  (self.num_layers, edge_index_with_loop.shape[1]),
                  device=self.device,
              )
              layer_edge_masks.view(-1)[mask_pool] -= 2
              layer_edge_mask_list.append(layer_edge_masks)
  
          weighted_change_walks_list = torch.stack(
              weighted_change_walks_list, dim=0
          )
          iter_weighted_change_walks_list.append(
              weighted_change_walks_list.detach()
          )
          layer_edge_mask_list = torch.stack(
              layer_edge_mask_list, dim=0
          ) * float("inf")
  
          # --- compute subsets' outputs of current iteration ---
          for layer_idx in range(self.num_layers):
              # --- Attention self-loop will be put at last because of the model will do it ---
              self.edge_mask[layer_idx].data = torch.cat(
                  [
                      layer_edge_mask_list[
                          :, layer_idx, : self.num_edges
                      ].reshape(-1),
                      layer_edge_mask_list[
                          :, layer_idx, self.num_edges :
                      ].reshape(-1),
                  ]
              ).sigmoid()
  
          batch = self.batch_input(x, edge_index, self.ns_per_iter)
          subsets_output = self.model(data=batch).softmax(1).detach()
          if not self.explain_graph:
              subsets_output = subsets_output.view(
                  self.ns_per_iter, -1, kwargs.get("num_classes")
              )[:, self.node_idx]
              last_subsets_output = torch.cat(
                  [
                      self.ori_logits_pred[self.node_idx].unsqueeze(0),
                      subsets_output.clone()[:-1],
                  ],
                  dim=0,
              )
          else:
              last_subsets_output = torch.cat(
                  [self.ori_logits_pred, subsets_output.clone()[:-1]],
                  dim=0,
              )
  
          changed_subsets_score_list = (
              last_subsets_output - subsets_output
          ).detach()
          iter_changed_subsets_score_list.append(
              changed_subsets_score_list
          )
  
          walk_sample_count += (
              (weighted_change_walks_list > 1e-30).float().sum(0)
          )
  
      # iter x subset_idx x flow_idx
      iter_weighted_change_walks_list = torch.stack(
          iter_weighted_change_walks_list, dim=0
      )
      iter_changed_subsets_score_list = torch.stack(
          iter_changed_subsets_score_list, dim=0
      )
  
      return (
          iter_weighted_change_walks_list,
          iter_changed_subsets_score_list,
          walk_sample_count,
      )
  
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 5 for calling the FlowX obj's underlying member method code with example usage: obj.flow_shap(x, edge_index, edge_index_with_loop, walk_indices_list) and returns: iter_weighted_change_walks_list,  iter_changed_subsets_score_list,  walk_sample_count