PairNorm gnn в тензорном потоке

Question

Я пытаюсь реализовать <PAIRNORM: TACKLING OVERSMOOTING IN GNNS> в тензорном потоке, используя spektral, вот мой код:

      import numpy as np
from tensorflow.keras.layers import Dropout, Input, Dense, LayerNormalization, BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras.regularizers import l2
import tensorflow as tf
from tensorflow.keras.layers import Layer
import tensorflow.keras.backend as K
from spektral.layers import GATConv

class GraphNorm_PairNorm(Layer):
    def __init__(self, **kwargs):
        super(GraphNorm_PairNorm, self).__init__(**kwargs)

    def call(self, inputs):
        '''
        # example：
        np.random.seed(1)
        x = np.random.randint(1, 20, [4, 3])
        y = x[:3, :3]
        y
        (y - y.mean(axis=0))/y.std(axis=0)
        y.mean(axis=0)
        y.std(axis=0)
        x = x.astype(np.float)
        # array([[13,  3, 14],
        #        [ 7, 10,  5],
        #        [15,  2, 17],
        #        [15,  9,  8]]
        a = np.eye(4)
        a[:3,:3]=1
        x = tf.constant(x)
        a = tf.constant(a)
        inputs = [x, a]


        '''

        x, a = inputs

        input_shape = x.shape
        ndims = len(input_shape)
        n_nodes = input_shape[-2]
        n_feas = input_shape[-1]

        x = K.expand_dims(x, -1)
        shape_x_tile = [1] * ndims + [n_nodes]
        x = K.tile(x, shape_x_tile)  # (n_graph, nodes, feas, nodes)

        a = K.expand_dims(a, -2)
        shape_a_tile = np.ones_like(a.shape)
        shape_a_tile[-2] = n_feas
        a = K.tile(a, shape_a_tile)  # (n_graph, nodes, feas, nodes)

        x_mask = x * a  # (n_graph, nodes, feas, nodes)
        x_len = tf.reduce_sum(a, -3, keepdims=True)  #
        x_mean = tf.reduce_sum(x_mask, -3, keepdims=True) / x_len
        x2 = tf.square(x_mask - x_mean)

        x2 = tf.where(tf.equal(a, 0), 0, x2)
        x_std = tf.sqrt(tf.reduce_sum(x2, -3, keepdims=True) / x_len)

        x_mean = tf.einsum("...ijk->...kj", x_mean)
        x_std = tf.einsum("...ijk->...kj", x_std)
        x_std = tf.where(tf.equal(x_std, 0), 1, x_std)
        opt = (inputs[0] - x_mean) / x_std

        return opt

# Parameters
a = np.ones([100, 100])
x = np.random.random([100, 100])
channels = 8  # Number of channels in each head of the first GAT layer
n_attn_heads = 8  # Number of attention heads in first GAT layer
dropout = 0.6  # Dropout rate for the features and adjacency matrix
l2_reg = 2.5e-4  # L2 regularization rate
learning_rate = 5e-3  # Learning rate
epochs = 20000  # Number of training epochs
patience = 100  # Patience for early stopping

N = x.shape[-2]  # Number of nodes in the graph
F = x.shape[-1]  # Original size of node features
n_out = 1  # 如果Number of classes

# Model definition
x_in = Input(shape=(F,))  # <tf.Tensor 'input_1:0' shape=(None, 1433) dtype=float32>
a_in = Input((None,), sparse=False)  #

do_1 = Dropout(dropout)(x_in)
gc_1 = GATConv(
    channels,  #
    attn_heads=n_attn_heads,
    concat_heads=True,
    dropout_rate=dropout,  #
    activation="elu",
    kernel_regularizer=l2(l2_reg),
    attn_kernel_regularizer=l2(l2_reg),
    bias_regularizer=l2(l2_reg),
)([do_1, a_in])
gc_1 = GraphNorm_PairNorm()((gc_1, a_in))
out1 = Dense(1)(gc_1)

# Build model
model1 = Model(inputs=[x_in, a_in], outputs=out1)
model1((x, a))  #

Я хочу, чтобы моя модель могла работать с другим графом с разным количеством узлов!!!

и ошибка TypeError: не удалось преобразовать элементы [1, 1, None] в Tensor. Рассмотрим приведение элементов к поддерживаемому типу, как это исправить?