class_docstr: The GNNExplainer model from the paper:
  
  > [GNNExplainer: Generating Explanations for Graph Neural Networks](https://arxiv.org/abs/1903.03894)<br>
  > Rex Ying, Dylan Bourgeois, Jiaxuan You, Marinka Zitnik and Jure Leskovec.
  
  The model can be used to explain the predictions for a single node or for an entire
  graph. In both cases, it returns the subgraph that mostly contributes to the
  prediction.
  
  **Arguments**
  
  - `model`: tf.keras.Model to explain;
  - `n_hops`: number of hops from which the GNN aggregates info. If `None`, then the
  number is inferred from the Conv and MessagePassing layers in the model.
  - `preprocess`: a preprocessing function to transform the adjacency matrix before
  giving it as input to the GNN; this is usually the same `preprocess` function of the
  Conv or MessagePassing layers used in the GNN (e.g., `GCNConv.preprocess`).
  - `graph_level`: if True, the GNN is assumed to be for graph-level prediction and
  the explanation is computed for the whole graph (and not just a node).
  - `verbose`: if True, print info during training;
  - `learning_rate`: learning rate when training the model;
  - `a_size_coef`: coefficient to control the number of edges of the subgraph that
  contributes to the prediction;
  - `x_size_coef`: coefficient to control the number of features of the subgraph
  that contributes to the prediction;
  - `a_entropy_coef`: coefficient to control the discretization of the adjacency
  mask;
  - `x_entropy_coef`: coefficient to control the discretization of the features
  mask;
  - `laplacian_coef`: coefficient to control the graph Laplacian loss;
class_docstr_tok: ["The","GNNExplainer","model","from","the","paper",":",">","[","GNNExplainer",":","Generating","Explanations","for","Graph","Neural","Networks","]","(","https",":","\/\/arxiv.org\/abs\/1903.03894",")","<","br",">",">","Rex","Ying",",","Dylan","Bourgeois",",","Jiaxuan","You",",","Marinka","Zitnik","and","Jure","Leskovec",".","The","model","can","be","used","to","explain","the","predictions","for","a","single","node","or","for","an","entire","graph",".","In","both","cases",",","it","returns","the","subgraph","that","mostly","contributes","to","the","prediction",".","*","*","Arguments","*","*","-","`","model","`",":","tf.keras.Model","to","explain",";","-","`","n_hops","`",":","number","of","hops","from","which","the","GNN","aggregates","info",".","If","`","None","`",",","then","the","number","is","inferred","from","the","Conv","and","MessagePassing","layers","in","the","model",".","-","`","preprocess","`",":","a","preprocessing","function","to","transform","the","adjacency","matrix","before","giving","it","as","input","to","the","GNN",";","this","is","usually","the","same","`","preprocess","`","function","of","the","Conv","or","MessagePassing","layers","used","in","the","GNN","(","e.g.",",","`","GCNConv.preprocess","`",")",".","-","`","graph_level","`",":","if","True",",","the","GNN","is","assumed","to","be","for","graph-level","prediction","and","the","explanation","is","computed","for","the","whole","graph","(","and","not","just","a","node",")",".","-","`","verbose","`",":","if","True",",","print","info","during","training",";","-","`","learning_rate","`",":","learning","rate","when","training","the","model",";","-","`","a_size_coef","`",":","coefficient","to","control","the","number","of","edges","of","the","subgraph","that","contributes","to","the","prediction",";","-","`","x_size_coef","`",":","coefficient","to","control","the","number","of","features","of","the","subgraph","that","contributes","to","the","prediction",";","-","`","a_entropy_coef","`",":","coefficient","to","control","the","discretization","of","the","adjacency","mask",";","-","`","x_entropy_coef","`",":","coefficient","to","control","the","discretization","of","the","features","mask",";","-","`","laplacian_coef","`",":","coefficient","to","control","the","graph","Laplacian","loss",";"]
docstr: Train the GNNExplainer to explain the given graph.
  
  :param x: feature matrix of shape `(n_nodes, n_node_features)`;
  :param a: sparse adjacency matrix of shape `(n_nodes, n_nodes)`;
  :param node_idx: index of the node to explain. If `self.graph_level=True`, this
  is ignored;
  :param epochs: number of epochs to train for.
  :return:
  - `a_mask`: mask for the adjacency matrix;
  - `x_mask`: mask for the node features.
docstr_tok: ["Train","the","GNNExplainer","to","explain","the","given","graph",".",":","param","x",":","feature","matrix","of","shape","`","(","n_nodes",",","n_node_features",")","`",";",":","param","a",":","sparse","adjacency","matrix","of","shape","`","(","n_nodes",",","n_nodes",")","`",";",":","param","node_idx",":","index","of","the","node","to","explain",".","If","`","self.graph_level=True","`",",","this","is","ignored",";",":","param","epochs",":","number","of","epochs","to","train","for",".",":","return",":","-","`","a_mask","`",":","mask","for","the","adjacency","matrix",";","-","`","x_mask","`",":","mask","for","the","node","features","."]
code: def explain_node(self, x, a, node_idx=None, epochs=100):
      """
      Train the GNNExplainer to explain the given graph.
  
      :param x: feature matrix of shape `(n_nodes, n_node_features)`;
      :param a: sparse adjacency matrix of shape `(n_nodes, n_nodes)`;
      :param node_idx: index of the node to explain. If `self.graph_level=True`, this
      is ignored;
      :param epochs: number of epochs to train for.
      :return:
      - `a_mask`: mask for the adjacency matrix;
      - `x_mask`: mask for the node features.
      """
      x = tf.cast(x, tf.float32)
      if node_idx is None:
          node_idx = 0
  
      # Get the computational graph
      if self.graph_level:
          self.comp_graph = tf.cast(a, tf.float32)
          self.i = tf.zeros(x.shape[0], dtype=tf.int32)
          self.y_pred = tf.argmax(
              self.model([x, a, self.i], training=False), axis=1
          )
      else:
          self.comp_graph = k_hop_sparse_subgraph(
              a, node_idx, self.n_hops, self.preprocess
          )
          self.y_pred = tf.argmax(
              self.model([x, a], training=False), axis=1
          )
  
      self.node_pred = self.y_pred[node_idx]
      self.y_pred = tf.cast(self.y_pred, tf.float32)
  
      # Optimizer for training
      self.opt = tf.keras.optimizers.Adam(self.learning_rate)
  
      # Init the trainable masks
      x_mask = tf.Variable(
          tf.random.normal((1, x.shape[1]), stddev=0.1),
          dtype=tf.float32,
          trainable=True,
      )
      a_mask = tf.Variable(
          tf.random.normal(
              self.comp_graph.values.shape,
              stddev=(2 / x.shape[0]) ** 0.5,
          ),
          dtype=tf.float32,
          trainable=True,
      )
  
      # Training loop
      for i in range(epochs):
          losses = self._train_step(x, a_mask, x_mask, node_idx)
          if self.verbose:
              print(
                  ", ".join(
                      [f"{key}: {val}" for key, val in losses.items()]
                  )
              )
  
      return a_mask, x_mask
  
file_path: spektral/spektral/models/gnn_explainer.py
imports: import networkx
  import numpy
  import tensorflow
  from scipy.sparse import csr_matrix
  from spektral.layers import MessagePassing
  from spektral.layers.convolutional.conv import Conv
  from spektral.layers.ops import dot
  from spektral.utils.sparse import sp_matrix_to_sp_tensor
label_desc: Use image node_id 2 for calling the GNNExplainer obj's underlying member method code with example usage: obj.explain_node(x, a, node_idx, epochs) and returns: a_mask,  x_mask