Issue with using pooling layers in Spektral package for graph neural networks

I am currently working on implementing graph neural networks using the Spektral package in Python. I am trying to include pooling layers (such as TopKPool) in my model, but I am running into dimension issues.

I have created a sample code to reproduce the error. Here is the code:

from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
from spektral.layers import GCNConv, TopKPool

# Example adjacency matrix and feature matrix
A = np.array([[0, 1, 1, 0],
              [1, 0, 1, 0],
              [1, 1, 0, 1],
              [0, 0, 1, 0]], dtype=np.float32)
X = np.array([[0, 1],
              [2, 3],
              [4, 5],
              [6, 7]], dtype=np.float32)

# Add a batch dimension to X
X = np.expand_dims(X, axis=0)
A = np.expand_dims(A, axis=0)

# Build the graph convolutional network model
X_in = Input((X.shape[1], X.shape[2]))
A_in = Input((A.shape[1],A.shape[2]))

gc1 = GCNConv(16, activation='relu')([X_in, A_in])
gc2 = GCNConv(32, activation='relu')([gc1, A_in])
pool1, A_out = TopKPool(ratio=0.5)([gc2, A_in])
flatten = Dense(128, activation='relu')(pool1)

model = Model(inputs=[X_in, A_in], outputs=flatten)

When I run this code, I get the following error message:

ValueError: Dimensions [1,1) of input[shape=[?]] = [] must match dimensions [1,2) of updates[shape=[?,1]] = [1]: Shapes must be equal rank, but are 0 and 1 for '{{node top_k_pool/TensorScatterUpdate}} = TensorScatterUpdate[T=DT_FLOAT, Tindices=DT_INT32](top_k_pool/sub_1, top_k_pool/strided_slice_5, top_k_pool/strided_slice_1)' with input shapes: [?], [?,1], [?,1].

I have tried using different pooling layers (SAGPool etc.) and checked the dimensions according to the package documentation, but I keep encountering the same error.

Here are Input dimensions:

Node features of shape (batch, n_nodes_in, n_node_features); Adjacency matrix of shape (batch, n_nodes_in, n_nodes_in);

I would greatly appreciate any assistance in resolving this issue.

danielegrattarola / spektral

Issue with using pooling layers in Spektral package for graph neural networks #418