How to use the spektral.layers.ops.matmul_AT_B function in spektral

def test_matmul_ops_inv_mixed_mode():
    A = np.random.randn(batch_size, N, N)
    B = np.random.randn(N, N)

    # A * B
    expected_output = np.array([A[i].dot(B) for i in range(batch_size)])
    _check_op(matmul_A_B, [A, B], expected_output)

    # A.T * B * A
    expected_output = np.array([A[i].T.dot(B).dot(A[i]) for i in range(batch_size)])
    _check_op(matmul_AT_B_A, [A, B], expected_output)

    # A.T * B
    expected_output = np.array([A[i].T.dot(B) for i in range(batch_size)])
    _check_op(matmul_AT_B, [A, B], expected_output)

    # A * B.T
    expected_output = np.array([A[i].dot(B.T) for i in range(batch_size)])
    _check_op(matmul_A_BT, [A, B], expected_output)

def test_matmul_ops_single_mode():
    A = np.random.randn(N, N)
    B = np.random.randn(N, N)
    convert_to_sparse = [[True, False], [False, True]]
    _check_op(matmul_A_B, [A, B], A.dot(B), convert_to_sparse)
    _check_op(matmul_AT_B_A, [A, B], A.T.dot(B).dot(A), convert_to_sparse)
    _check_op(matmul_AT_B, [A, B], A.T.dot(B), convert_to_sparse)
    _check_op(matmul_A_BT, [A, B], A.dot(B.T), convert_to_sparse)

if batch_mode:
            cut_loss = K.mean(cut_loss)
        self.add_loss(cut_loss)

        # Orthogonality regularization
        SS = ops.matmul_AT_B(S, S)
        I_S = tf.eye(self.k)
        ortho_loss = tf.norm(
            SS / tf.norm(SS, axis=(-1, -2)) - I_S / tf.norm(I_S), axis=(-1, -2)
        )
        if batch_mode:
            ortho_loss = K.mean(cut_loss)
        self.add_loss(ortho_loss)

        # Pooling
        X_pooled = ops.matmul_AT_B(S, X)
        A_pooled = tf.linalg.set_diag(A_pooled, tf.zeros(K.shape(A_pooled)[:-1]))  # Remove diagonal
        A_pooled = ops.normalize_A(A_pooled)

        output = [X_pooled, A_pooled]

        if I is not None:
            I_mean = tf.segment_mean(I, I)
            I_pooled = ops.repeat(I_mean, tf.ones_like(I_mean) * self.k)
            output.append(I_pooled)

        if self.return_mask:
            output.append(S)

        return output

S = K.bias_add(S, self.bias_out)
        S = activations.softmax(S, axis=-1)  # Apply softmax to get cluster assignments

        # MinCut regularization
        A_pooled = ops.matmul_AT_B_A(S, A)
        num = tf.trace(A_pooled)

        D = ops.degree_matrix(A)
        den = tf.trace(ops.matmul_AT_B_A(S, D))
        cut_loss = -(num / den)
        if batch_mode:
            cut_loss = K.mean(cut_loss)
        self.add_loss(cut_loss)

        # Orthogonality regularization
        SS = ops.matmul_AT_B(S, S)
        I_S = tf.eye(self.k)
        ortho_loss = tf.norm(
            SS / tf.norm(SS, axis=(-1, -2)) - I_S / tf.norm(I_S), axis=(-1, -2)
        )
        if batch_mode:
            ortho_loss = K.mean(cut_loss)
        self.add_loss(ortho_loss)

        # Pooling
        X_pooled = ops.matmul_AT_B(S, X)
        A_pooled = tf.linalg.set_diag(A_pooled, tf.zeros(K.shape(A_pooled)[:-1]))  # Remove diagonal
        A_pooled = ops.normalize_A(A_pooled)

        output = [X_pooled, A_pooled]

        if I is not None:

else:
            LP_loss = A - S_gram
        LP_loss = tf.norm(LP_loss, axis=(-1, -2))
        if self.reduce_loss:
            LP_loss = K.mean(LP_loss)
        self.add_loss(LP_loss)

        # Entropy loss
        entr = tf.negative(tf.reduce_sum(tf.multiply(S, K.log(S + K.epsilon())), axis=-1))
        entr_loss = K.mean(entr, axis=-1)
        if self.reduce_loss:
            entr_loss = K.mean(entr_loss)
        self.add_loss(entr_loss)

        # Pooling
        X_pooled = ops.matmul_AT_B(S, Z)
        A_pooled = ops.matmul_AT_B_A(S, A)

        if K.ndim(A_pooled) == 3:
            self.mixed_mode = True

        output = [X_pooled, A_pooled]

        if I is not None:
            I_mean = tf.segment_mean(I, I)
            I_pooled = ops.repeat(I_mean, tf.ones_like(I_mean) * self.k)
            output.append(I_pooled)

        if self.return_mask:
            output.append(S)

        return output