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","."]
code: def train_mask(
      self,
      x: Tensor,
      edge_index: Tensor,
      ex_label: Tensor,
      walk_indices_list,
      edge_index_with_loop,
      iter_weighted_change_walks_list,
      iter_changed_subsets_score_list,
      walk_sample_count,
      t0=7.0,
      t1=0.5,
      **kwargs,
  ) -> None:
      # initialize a mask
      self.to(x.device)
  
      self.nec_suf_mask = nn.Parameter(
          1e-1
          * nn.init.uniform_(
              torch.empty(
                  (1, iter_weighted_change_walks_list.shape[1], 1),
                  device=self.device,
              )
          )
      )
  
      if self.no_mask:
          self.nec_suf_mask = nn.Parameter(
              100
              * torch.ones(
                  (1, iter_weighted_change_walks_list.shape[1], 1),
                  device=self.device,
              )
          )
      self.iter_weighted_change_walks_list = nn.Parameter(
          iter_weighted_change_walks_list.clone().detach()
      )
  
      # --- Training ---
      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)
  
      self.flow2layeredge_matrix = torch.stack(
          [
              (walk_plain_indices_list == i).float().sum(dim=1)
              for i in range(
                  self.num_layers * (self.num_edges + self.num_nodes)
              )
          ],
          dim=1,
      ).detach()
  
      # train to get the mask
      optimizer = torch.optim.Adam(
          [{"params": self.nec_suf_mask}],
          # {'params': self.iter_weighted_change_walks_list, 'lr': self.score_lr}],
          lr=self.lr,
      )
      # print('#I#begin this ex_label')
      for epoch in range(1, self.epochs + 1):
          masked_iter_weighted_change_walks_list = (
              self.iter_weighted_change_walks_list
              * self.nec_suf_mask.sigmoid()
          )
  
          walk_scores = (
              (
                  masked_iter_weighted_change_walks_list.unsqueeze(
                      3
                  ).repeat(
                      1, 1, 1, iter_changed_subsets_score_list.shape[2]
                  )
                  * iter_changed_subsets_score_list.unsqueeze(2)
              )
              .sum(1)
              .sum(0)
          )
          # EPS will affect the stability of training
          EPS = 1e-18
          shap_flow_score = walk_scores / (
              walk_sample_count.unsqueeze(1) + EPS
          )
  
          # --- score/mask transformer ---
          self.flow_mask = shap_flow_score[:, ex_label]
  
          # --- setting layer edge masks ---
          self.layer_edge_mask = (
              (self.flow_mask * self.flow2layeredge_matrix)
              .view(self.flow_mask.shape[0], self.num_layers, -1)
              .sum(0)
          )
          mask = self.layer_edge_mask.sum(0)
          mask = mask - mask.min()
          mask = mask / (mask.max() + EPS)
  
          climb = True
          if climb:
              # pass
              mask = mask**8
          else:
              end_epoch = 300
              temperature = (
                  float(t0 * ((t1 / t0) ** (epoch / end_epoch)))
                  if epoch < end_epoch
                  else t1
              )
              mask = gumbel_softmax(mask, temperature, training=True)
  
          mask = mask - mask.min()
          mask = mask / (mask.max() + EPS)
          cur_sparsity = (mask < 0.5).sum().float() / mask.shape[0]
          # if cur_sparsity < x_args.sparsity:
          #     # --- early stop ---
          #     break
          # if epoch % 20 == 0:
          #     print(f'Epoch: {epoch} --- training mask Sparsity: {cur_sparsity}')
          if self.fidelity_plus:
              mask = 1 - mask  # Fidelity +
          self.mask = mask
          # isig_mask = torch.log(self.mask / (1 - self.mask + EPS) + EPS)
  
          # --- temp update non-leaf edge_mask
          temp_edge_mask = []
          for layer_idx in range(self.num_layers):
              # --- Attention self-loop will be put at last because of the model will do it ---
              temp_edge_mask.append(mask)
  
          # debug:
          with self.temp_mask(self, temp_edge_mask):
              raw_preds = self.model(x, edge_index, **kwargs)
          loss = self.__loss__(raw_preds, ex_label)
  
          # if epoch % 20 == 0:
          #     print(f'#D#Loss:{loss.item()}')
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
      return
  
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 4 for calling the FlowX obj's underlying member method code with example usage: obj.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, t0, t1) without return types