spektral: GNNExplainer: tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes

Hi Daniele and all,

thanks for creating and maintaining this great library!

I have been trying to use GNNExplainer, but I keep seeing the below error message. Still don’t know if it’s a bug or something i am doing wrong on my side, but there is no much documentation or examples around it.

I am able to run smoothly the sample code at https://github.com/danielegrattarola/spektral/blob/master/examples/other/explain_node_predictions.py.

But when applying to my dataset (where I can run successfully a GCN model), i get the below:

dataset
Out[66]: SADataset(n_graphs=1)

dataset[0]
Out[67]: Graph(n_nodes=1653, n_node_features=42, n_edge_features=None, n_labels=10)

x_exp, a_exp = dataset[0].x, dataset[0].a

x_exp.shape
Out[69]: (1653, 42)

a_exp.shape
Out[70]: TensorShape([1653, 1653])

explainer = GNNExplainer(model, preprocess=gcn_filter, verbose=True)
n_hops was automatically inferred to be 2

node_idx = 0

adj_mask, feat_mask = explainer.explain_node(x=x_exp, a=a_exp, node_idx=node_idx)

pred_loss: 1.097847819328308, a_size_loss: 0.5874298214912415, a_entropy_loss: 0.0692998617887497, smoothness_loss: [[0.]], x_size_loss: 2.0829315185546875, x_entropy_loss: 0.06919442862272263
pred_loss: 1.0877137184143066, a_size_loss: 0.5852940678596497, a_entropy_loss: 0.06929884105920792, smoothness_loss: [[0.]], x_size_loss: 2.075951099395752, x_entropy_loss: 0.06918510049581528
[... output removed]
pred_loss: 0.6421844959259033, a_size_loss: 0.3796449303627014, a_entropy_loss: 0.05964722856879234, smoothness_loss: [[0.]], x_size_loss: 1.379091501235962, x_entropy_loss: 0.05970795825123787
pred_loss: 0.6415124535560608, a_size_loss: 0.37782761454582214, a_entropy_loss: 0.05948375537991524, smoothness_loss: [[0.]], x_size_loss: 1.372214436531067, x_entropy_loss: 0.059564121067523956

adj_mask.shape
Out[75]: TensorShape([2349])

adj_mask
Out[76]: 
<tf.Variable 'Variable:0' shape=(2349,) dtype=float32, numpy=
array([ 0.8150444 ,  0.77765435, -0.9916512 , ..., -1.0242233 ,
       -0.9629407 , -0.9988212 ], dtype=float32)>


feat_mask.shape
Out[77]: TensorShape([1, 42])

feat_mask
Out[78]: 
<tf.Variable 'Variable:0' shape=(1, 42) dtype=float32, numpy=
array([[ 0.58385307, -1.3217939 , -1.0627872 , -0.00148061, -1.0020486 ,
        -0.9942789 , -0.97092587, -0.9922697 ,  0.3853194 , -0.83190703,
        -1.1318972 , -0.99104863, -1.0001428 , -0.9827519 , -0.9750702 ,
        -0.96384555, -0.890569  , -1.0193573 ,  0.4747884 , -0.91873515,
         0.7341433 , -0.97718424, -0.86869913, -0.9699511 ,  0.37709397,
        -1.0660834 , -0.92709947, -0.89111555, -1.0546191 , -1.0837208 ,
        -1.0699799 , -1.0806109 ,  0.61809593, -0.9817147 , -1.0526807 ,
        -0.95195514, -1.0162035 , -1.181156  , -1.0657567 , -1.0472083 ,
        -0.85559815, -1.0388821 ]], dtype=float32)>

G = explainer.plot_subgraph(adj_mask, feat_mask, node_idx)
Traceback (most recent call last):
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-74-182e1ffafc94>", line 1, in <module>
    G = explainer.plot_subgraph(adj_mask, feat_mask, node_idx)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/spektral/models/gnn_explainer.py", line 276, in plot_subgraph
    adj_mtx, top_ftrs = self._explainer_cleaning(a_mask, x_mask, node_idx, a_thresh)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/spektral/models/gnn_explainer.py", line 243, in _explainer_cleaning
    tf.multiply, self.comp_graph, selected_adj_mask
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/ops/sparse_ops.py", line 2931, in map_values
    op(*inner_args, **inner_kwargs),
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/ops/math_ops.py", line 530, in multiply
    return gen_math_ops.mul(x, y, name)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/ops/gen_math_ops.py", line 6240, in mul
    _ops.raise_from_not_ok_status(e, name)
  File ".pyenv/versions/3.7.6/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 6897, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible shapes: [2349] vs. [2589] [Op:Mul]