k3-node 1.0.0__py3-none-any.whl

k3_node/models/sgformer.py

@@ -0,0 +1,195 @@
+from typing import Optional
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.layers.aggr.base import from_dense_batch, to_dense_batch
+
+from k3_node.layers.conv import GCNConv
+from k3_node.layers.attention import SGFormerAttention
+
+
+class GraphModule(keras.layers.Layer):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        num_layers: int = 2,
+        dropout: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.convs = []
+        self.fcs = [keras.layers.Dense(hidden_channels)]
+        self.bns = [keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)]
+
+        for _ in range(num_layers):
+            self.convs.append(GCNConv(hidden_channels, hidden_channels))
+            self.bns.append(keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5))
+
+        self.drop = keras.layers.Dropout(dropout)
+        self.dropout_rate = dropout
+
+    def call(self, x, edge_index, training=False):
+        x = self.fcs[0](x)
+        x = self.bns[0](x, training=training)
+        x = ops.relu(x)
+        x = self.drop(x, training=training)
+        last_x = x
+
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
+            x = self.bns[i + 1](x, training=training)
+            x = ops.relu(x)
+            x = self.drop(x, training=training)
+            x = x + last_x
+        return x
+
+
+class SGModule(keras.layers.Layer):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        num_layers: int = 2,
+        num_heads: int = 1,
+        dropout: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.attns = []
+        self.fcs = [keras.layers.Dense(hidden_channels)]
+        self.bns = [keras.layers.LayerNormalization()]
+
+        for _ in range(num_layers):
+            self.attns.append(
+                SGFormerAttention(hidden_channels, num_heads, hidden_channels)
+            )
+            self.bns.append(keras.layers.LayerNormalization())
+
+        self.drop = keras.layers.Dropout(dropout)
+        self.dropout_rate = dropout
+
+    def call(self, x, batch=None, training=False):
+        if batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int64")
+
+        # Group the nodes by graph (static shapes: works under jit)
+        batch = ops.cast(batch, "int32")
+        indices = ops.argsort(batch)
+        rev_perm = ops.argsort(indices)
+        batch_sorted = ops.take(batch, indices, axis=0)
+        x_dense, mask = to_dense_batch(ops.take(x, indices, axis=0), batch_sorted)
+
+        layer_ = []
+
+        # input MLP layer
+        x = self.fcs[0](x_dense)
+        x = self.bns[0](x, training=training)
+        x = ops.relu(x)
+        x = self.drop(x, training=training)
+
+        layer_.append(x)
+
+        for i, attn in enumerate(self.attns):
+            x = attn(x, mask)
+            x = (x + layer_[i]) / 2.0
+            x = self.bns[i + 1](x, training=training)
+            x = ops.relu(x)
+            x = self.drop(x, training=training)
+            layer_.append(x)
+
+        unsorted_x_mask = ops.take(from_dense_batch(x, batch_sorted), rev_perm, axis=0)
+        return unsorted_x_mask
+
+
+class SGFormer(keras.Model):
+    r"""The sgformer module from the
+    `"SGFormer: Simplifying and Empowering Transformers for
+    Large-Graph Representations"
+    <https://arxiv.org/abs/2306.10759>`_ paper.
+
+    Args:
+        in_channels (int): Input channels.
+        hidden_channels (int): Hidden channels.
+        out_channels (int): Output channels.
+        trans_num_layers (int, optional): The number of layers for all-pair attention.
+            (default: :obj:`2`)
+        trans_num_heads (int, optional): The number of heads for attention.
+            (default: :obj:`1`)
+        trans_dropout (float, optional): Global dropout rate.
+            (default: :obj:`0.5`)
+        gnn_num_layers (int, optional): The number of layers for GNN.
+            (default: :obj:`3`)
+        gnn_dropout (float, optional): GNN dropout rate.
+            (default: :obj:`0.5`)
+        graph_weight (float, optional): The weight balance global and gnn module.
+            (default: :obj:`0.5`)
+        aggregate (str, optional): Aggregate type (:obj:`'add'` or :obj:`'cat'`).
+            (default: :obj:`'add'`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import SGFormer
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        model = SGFormer(in_channels=8, hidden_channels=32, out_channels=4, trans_num_layers=1,
+                         trans_num_heads=1, gnn_num_layers=1)
+        print(tuple(model(x, edge_index).shape))  # (10, 4)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        trans_num_layers: int = 2,
+        trans_num_heads: int = 1,
+        trans_dropout: float = 0.5,
+        gnn_num_layers: int = 3,
+        gnn_dropout: float = 0.5,
+        graph_weight: float = 0.5,
+        aggregate: str = "add",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.graph_weight = graph_weight
+        self.aggregate = aggregate
+
+        self.trans_conv = SGModule(
+            in_channels,
+            hidden_channels,
+            trans_num_layers,
+            trans_num_heads,
+            trans_dropout,
+        )
+        self.graph_conv = GraphModule(
+            in_channels,
+            hidden_channels,
+            gnn_num_layers,
+            gnn_dropout,
+        )
+
+        if aggregate == "add":
+            self.fc = keras.layers.Dense(out_channels)
+        elif aggregate == "cat":
+            self.fc = keras.layers.Dense(out_channels)
+        else:
+            raise ValueError(f"Invalid aggregate type: {aggregate}")
+
+    def call(self, x, edge_index, batch: Optional[any] = None, training=False):
+        x1 = self.trans_conv(x, batch, training=training)
+        x2 = self.graph_conv(x, edge_index, training=training)
+        if self.aggregate == "add":
+            x = self.graph_weight * x2 + (1.0 - self.graph_weight) * x1
+        else:
+            x = ops.concatenate([x1, x2], axis=1)
+        x = self.fc(x)
+        return ops.log_softmax(x, axis=-1)
+

k3_node/models/signed_gcn.py

@@ -0,0 +1,185 @@
+from typing import Optional, Tuple
+
+import keras
+from keras import ops
+import numpy as np
+
+from k3_node.layers.conv import SignedConv
+
+
+class SignedGCN(keras.layers.Layer):
+    r"""The signed graph convolutional network model from the `"Signed Graph
+    Convolutional Network" <https://arxiv.org/abs/1808.06354>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        num_layers (int): Number of layers.
+        lamb (float, optional): Balances the contributions of the overall
+            objective. (default: :obj:`5`)
+        bias (bool, optional): If set to :obj:`False`, all layers will not
+            learn an additive bias. (default: :obj:`True`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import SignedGCN
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        pos_edge_index, neg_edge_index = edge_index[:, :15], edge_index[:, 15:]  # positive / negative links
+        model = SignedGCN(in_channels=8, hidden_channels=16, num_layers=2)
+        z = model(x, pos_edge_index, neg_edge_index)
+        print(tuple(z.shape))  # (10, 16)
+        print(tuple(model.discriminate(z, edge_index).shape))  # (30, 3): scores for (positive, negative, no link)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        num_layers: int,
+        lamb: float = 5.0,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.num_layers = num_layers
+        self.lamb = lamb
+
+        self.conv1 = SignedConv(in_channels, hidden_channels // 2, first_aggr=True, bias=bias)
+        self.convs = [
+            SignedConv(hidden_channels // 2, hidden_channels // 2, first_aggr=False, bias=bias)
+            for _ in range(num_layers - 1)
+        ]
+        self.lin = keras.layers.Dense(3, use_bias=bias)
+        self.lin.build((None, 2 * hidden_channels))  # scores a pair of node embeddings
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def reset_parameters(self):
+        self.conv1.reset_parameters()
+        for conv in self.convs:
+            conv.reset_parameters()
+        if self.lin.built:
+            self.lin.kernel.assign(keras.initializers.GlorotUniform()(self.lin.kernel.shape))
+            if self.lin.bias is not None:
+                self.lin.bias.assign(ops.zeros(self.lin.bias.shape))
+
+    @staticmethod
+    def split_edges(edge_index, test_ratio: float = 0.2):
+        num_edges = ops.shape(edge_index)[1]
+        perm = np.random.permutation(num_edges)
+        num_test = int(test_ratio * num_edges)
+        test_mask = np.zeros(num_edges, dtype=bool)
+        test_mask[perm[:num_test]] = True
+
+        edge_index_np = ops.convert_to_numpy(edge_index)
+        train_edge_index = ops.convert_to_tensor(edge_index_np[:, ~test_mask])
+        test_edge_index = ops.convert_to_tensor(edge_index_np[:, test_mask])
+        return train_edge_index, test_edge_index
+
+    def create_spectral_features(self, pos_edge_index, neg_edge_index, num_nodes: Optional[int] = None):
+        r"""Creates ``in_channels`` spectral node features from the positive and negative edges:
+        a truncated SVD of the signed, symmetric adjacency matrix."""
+        import scipy.sparse as sp
+        from sklearn.decomposition import TruncatedSVD
+
+        from k3_node.utils.graph import coalesce
+
+        pos = np.asarray(ops.convert_to_numpy(pos_edge_index)).astype(np.int64)
+        neg = np.asarray(ops.convert_to_numpy(neg_edge_index)).astype(np.int64)
+        edge_index = np.concatenate([pos, neg], axis=1)
+        N = int(edge_index.max()) + 1 if num_nodes is None else num_nodes
+        val = np.concatenate([np.full(pos.shape[1], 2.0), np.zeros(neg.shape[1])]).astype(np.float32)
+        edge_index = np.concatenate([edge_index, edge_index[::-1]], axis=1)
+        val = np.concatenate([val, val])
+        edge_index, val = coalesce(edge_index, val, num_nodes=N)
+        edge_index, val = np.asarray(ops.convert_to_numpy(edge_index)), np.asarray(ops.convert_to_numpy(val)) - 1
+
+        A = sp.coo_matrix((val, edge_index), shape=(N, N))
+        svd = TruncatedSVD(n_components=self.in_channels, n_iter=128)
+        svd.fit(A)
+        return ops.convert_to_tensor(svd.components_.T.astype(np.float32))
+
+    def call(self, x, pos_edge_index, neg_edge_index):
+        z = ops.relu(self.conv1(x, pos_edge_index, neg_edge_index))
+        for conv in self.convs:
+            z = ops.relu(conv(z, pos_edge_index, neg_edge_index))
+        return z
+
+    def discriminate(self, z, edge_index):
+        value = ops.concatenate([
+            ops.take(z, edge_index[0], axis=0),
+            ops.take(z, edge_index[1], axis=0),
+        ], axis=-1)
+        value = self.lin(value)
+        return ops.log_softmax(value, axis=-1)
+
+    def nll_loss(self, z, pos_edge_index, neg_edge_index):
+        from k3_node.models.utils import negative_sampling
+
+        edge_index = ops.concatenate([pos_edge_index, neg_edge_index], axis=1)
+        none_edge_index = negative_sampling(edge_index, ops.shape(z)[0])
+
+        nll_loss = 0.0
+        pos_pred = self.discriminate(z, pos_edge_index)
+        nll_loss = nll_loss - ops.mean(pos_pred[:, 0])
+
+        neg_pred = self.discriminate(z, neg_edge_index)
+        nll_loss = nll_loss - ops.mean(neg_pred[:, 1])
+
+        none_pred = self.discriminate(z, none_edge_index)
+        nll_loss = nll_loss - ops.mean(none_pred[:, 2])
+
+        return nll_loss / 3.0
+
+    def pos_embedding_loss(self, z, pos_edge_index):
+        from k3_node.models.utils import structured_negative_sampling
+
+        i, j, k = structured_negative_sampling(pos_edge_index, ops.shape(z)[0])
+        out = (
+            ops.sum((ops.take(z, i, axis=0) - ops.take(z, j, axis=0)) ** 2, axis=1)
+            - ops.sum((ops.take(z, i, axis=0) - ops.take(z, k, axis=0)) ** 2, axis=1)
+        )
+        return ops.mean(ops.maximum(out, 0.0))
+
+    def neg_embedding_loss(self, z, neg_edge_index):
+        from k3_node.models.utils import structured_negative_sampling
+
+        i, j, k = structured_negative_sampling(neg_edge_index, ops.shape(z)[0])
+        out = (
+            ops.sum((ops.take(z, i, axis=0) - ops.take(z, k, axis=0)) ** 2, axis=1)
+            - ops.sum((ops.take(z, i, axis=0) - ops.take(z, j, axis=0)) ** 2, axis=1)
+        )
+        return ops.mean(ops.maximum(out, 0.0))
+
+    def loss(self, z, pos_edge_index, neg_edge_index):
+        nll_loss = self.nll_loss(z, pos_edge_index, neg_edge_index)
+        loss_1 = self.pos_embedding_loss(z, pos_edge_index)
+        loss_2 = self.neg_embedding_loss(z, neg_edge_index)
+        return nll_loss + self.lamb * (loss_1 + loss_2)
+
+    def test(self, z, pos_edge_index, neg_edge_index):
+        from sklearn.metrics import f1_score, roc_auc_score
+
+        pos_p = ops.argmax(self.discriminate(z, pos_edge_index)[:, :2], axis=1)
+        neg_p = ops.argmax(self.discriminate(z, neg_edge_index)[:, :2], axis=1)
+        pred = 1 - ops.convert_to_numpy(ops.concatenate([pos_p, neg_p], axis=0))
+        y = np.concatenate([
+            np.ones(ops.shape(pos_p)[0], dtype=np.int64),
+            np.zeros(ops.shape(neg_p)[0], dtype=np.int64),
+        ])
+        auc = roc_auc_score(y, pred)
+        f1 = f1_score(y, pred, average="binary") if pred.sum() > 0 else 0.0
+        return auc, f1
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.hidden_channels}, num_layers={self.num_layers})')
+

k3_node/models/test_attentive_fp.py

@@ -0,0 +1,32 @@
+from keras import ops
+
+from k3_node.models import AttentiveFP
+
+
+def test_attentive_fp():
+    model = AttentiveFP(8, 16, 32, edge_dim=3, num_layers=2, num_timesteps=2)
+    assert str(model) == (
+        "AttentiveFP(in_channels=8, hidden_channels=16, "
+        "out_channels=32, edge_dim=3, num_layers=2, "
+        "num_timesteps=2)"
+    )
+
+    x = ops.convert_to_tensor([[1.0] * 8] * 4, dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]], dtype="int64")
+    edge_attr = ops.convert_to_tensor([[1.0] * 3] * 6, dtype="float32")
+    batch = ops.convert_to_tensor([0, 0, 0, 0], dtype="int64")
+
+    out = model(x, edge_index, edge_attr, batch)
+    assert ops.shape(out) == (1, 32)
+
+
+def test_attentive_fp_multi_graph_batch():
+    model = AttentiveFP(4, 8, 2, edge_dim=2, num_layers=2, num_timesteps=2)
+
+    x = ops.convert_to_tensor([[1.0] * 4] * 6, dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 1, 3, 4], [1, 0, 4, 3]], dtype="int64")
+    edge_attr = ops.convert_to_tensor([[1.0] * 2] * 4, dtype="float32")
+    batch = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int64")
+
+    out = model(x, edge_index, edge_attr, batch)
+    assert ops.shape(out) == (2, 2)

k3_node/models/test_attract_repel.py

@@ -0,0 +1,33 @@
+from keras import ops
+
+from k3_node.models import ARLinkPredictor
+
+
+def test_ar_link_predictor():
+    model = ARLinkPredictor(in_channels=16, hidden_channels=32, num_layers=2)
+    x = ops.convert_to_tensor([[0.1] * 16] * 4, dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 3]], dtype="int64")
+
+    pred = model(x, edge_index)
+    assert ops.shape(pred)[0] == ops.shape(edge_index)[1]
+    pred_np = ops.convert_to_numpy(pred)
+    assert (pred_np >= 0).all() and (pred_np <= 1).all()
+
+    attract_z, repel_z = model.encode(x)
+    assert ops.shape(attract_z) == (4, 16)
+    assert ops.shape(repel_z) == (4, 16)
+
+    raw_scores = model.decode(attract_z, repel_z, edge_index)
+    assert ops.shape(raw_scores)[0] == ops.shape(edge_index)[1]
+
+    r_fraction = model.calculate_r_fraction(attract_z, repel_z)
+    assert 0 <= r_fraction <= 1
+
+
+def test_ar_link_predictor_with_custom_ratio():
+    model = ARLinkPredictor(in_channels=8, hidden_channels=20, attract_ratio=0.7)
+    x = ops.convert_to_tensor([[0.1] * 8] * 5, dtype="float32")
+
+    attract_z, repel_z = model.encode(x)
+    assert ops.shape(attract_z) == (5, 14)
+    assert ops.shape(repel_z) == (5, 6)

k3_node/models/test_autoencoder.py

@@ -0,0 +1,119 @@
+import numpy as np
+from keras import ops
+
+from k3_node.models import ARGA, ARGVA, GAE, VGAE
+
+
+def test_gae():
+    model = GAE(encoder=lambda x: x)
+    model.reset_parameters()
+
+    x = ops.convert_to_tensor([[1.0, -1.0], [1.0, 2.0], [2.0, 1.0]], dtype="float32")
+    z = model.encode(x)
+    assert np.allclose(ops.convert_to_numpy(z), ops.convert_to_numpy(x))
+
+    adj = model.decoder.forward_all(z)
+    expected = ops.sigmoid(ops.convert_to_tensor([
+        [2.0, -1.0, 1.0],
+        [-1.0, 5.0, 4.0],
+        [1.0, 4.0, 5.0],
+    ], dtype="float32"))
+    assert np.allclose(ops.convert_to_numpy(adj), ops.convert_to_numpy(expected), atol=1e-5)
+
+    edge_index = ops.convert_to_tensor([[0, 1], [1, 2]], dtype="int64")
+    value = model.decode(z, edge_index)
+    expected_value = ops.sigmoid(ops.convert_to_tensor([-1.0, 4.0], dtype="float32"))
+    assert np.allclose(ops.convert_to_numpy(value), ops.convert_to_numpy(expected_value), atol=1e-5)
+
+    pos_edge_index = ops.convert_to_tensor(
+        [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], dtype="int64"
+    )
+    z2 = ops.convert_to_tensor(np.random.randn(11, 16).astype("float32"))
+    loss = model.recon_loss(z2, pos_edge_index)
+    assert float(ops.convert_to_numpy(loss)) > 0
+
+    auc, ap = model.test(z2, pos_edge_index, pos_edge_index)
+    assert 0 <= auc <= 1 and 0 <= ap <= 1
+
+
+def test_vgae():
+    model = VGAE(encoder=lambda x: (x, x))
+
+    x = ops.convert_to_tensor([[1.0, -1.0], [1.0, 2.0], [2.0, 1.0]], dtype="float32")
+    model.encode(x)
+    assert float(ops.convert_to_numpy(model.kl_loss())) > 0
+
+    model.eval()
+    model.encode(x)
+
+
+def test_arga():
+    model = ARGA(encoder=lambda x: x, discriminator=lambda x: ops.convert_to_tensor([0.5]))
+    model.reset_parameters()
+
+    x = ops.convert_to_tensor([[1.0, -1.0], [1.0, 2.0], [2.0, 1.0]], dtype="float32")
+    z = model.encode(x)
+
+    assert float(ops.convert_to_numpy(model.reg_loss(z))) > 0
+    assert float(ops.convert_to_numpy(model.discriminator_loss(z))) > 0
+
+
+def test_argva():
+    model = ARGVA(encoder=lambda x: (x, x), discriminator=lambda x: ops.convert_to_tensor([0.5]))
+
+    x = ops.convert_to_tensor([[1.0, -1.0], [1.0, 2.0], [2.0, 1.0]], dtype="float32")
+    model.encode(x)
+    model.reparametrize(model._mu, model._logstd)
+    assert float(ops.convert_to_numpy(model.kl_loss())) > 0
+
+
+def test_init():
+    import keras
+
+    encoder = keras.layers.Dense(32)
+    encoder.build((None, 16))
+    decoder = keras.layers.Dense(16)
+    decoder.build((None, 32))
+    discriminator = keras.layers.Dense(1)
+    discriminator.build((None, 32))
+
+    GAE(encoder, decoder)
+    VGAE(encoder, decoder)
+    ARGA(encoder, discriminator, decoder)
+    ARGVA(encoder, discriminator, decoder)
+
+
+def test_fit_evaluate_all_variants():
+    import keras
+    import numpy as np
+    from k3_node.data import Data
+    from k3_node.layers import GCNConv
+
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 20, size=(2, 60))
+    data = Data(x=rng.random((20, 8)).astype("float32"), edge_index=edge_index,
+                pos_edge_label_index=edge_index, neg_edge_label_index=rng.integers(0, 20, size=(2, 60)))
+
+    class Encoder(keras.Model):
+        def __init__(self, variational):
+            super().__init__()
+            self.variational = variational
+            self.conv_mu, self.conv_logstd = GCNConv(8, 4), GCNConv(8, 4)
+
+        def call(self, x, edge_index):
+            mu = self.conv_mu(x, edge_index)
+            return (mu, self.conv_logstd(x, edge_index)) if self.variational else mu
+
+    def discriminator():
+        return keras.Sequential([keras.layers.Dense(8, activation="relu"), keras.layers.Dense(1)])
+
+    for model in [GAE(Encoder(False)), VGAE(Encoder(True)), ARGA(Encoder(False), discriminator()),
+                  ARGVA(Encoder(True), discriminator())]:
+        model.compile(keras.optimizers.Adam(0.01), discriminator_optimizer=keras.optimizers.Adam(0.01))
+        model.embed(data)  # builds the encoder
+        before = [np.array(v) for v in model.encoder.trainable_variables]
+        history = model.fit(data, epochs=2, validation_data=data, verbose=0)
+        assert set(history) == {"loss", "val_auc", "val_ap"} and len(history["loss"]) == 2
+        assert any(not np.allclose(b, np.array(v)) for b, v in zip(before, model.encoder.trainable_variables))
+        assert 0.0 <= model.evaluate(data)["auc"] <= 1.0
+        assert tuple(model.embed(data).shape) == (20, 4)

k3_node/models/test_basic_gnn.py

@@ -0,0 +1,102 @@
+import pytest
+from keras import ops
+
+from k3_node.models import GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN
+
+
+@pytest.mark.parametrize('jk', [None, 'last', 'cat', 'max', 'lstm'])
+@pytest.mark.parametrize('out_dim', [None, 4])
+def test_gcn(jk, out_dim):
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_weight = ops.convert_to_tensor([1.0, 1.0, 0.5, 0.5])
+    out_channels = 16 if out_dim is None else out_dim
+
+    model = GCN(8, 16, num_layers=3, out_channels=out_dim, dropout=0.1,
+                act='relu', norm='batch_norm', jk=jk)
+    assert str(model) == f'GCN(8, {out_channels}, num_layers=3)'
+
+    out = model(x, edge_index, edge_weight=edge_weight)
+    assert ops.shape(out) == (4, out_channels)
+
+
+@pytest.mark.parametrize('jk', [None, 'cat', 'max'])
+def test_graph_sage(jk):
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    model = GraphSAGE(8, 16, num_layers=2, out_channels=8, jk=jk)
+    assert str(model) == 'GraphSAGE(8, 8, num_layers=2)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 8)
+
+
+@pytest.mark.parametrize('jk', [None, 'cat'])
+def test_gin(jk):
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    model = GIN(8, 16, num_layers=2, out_channels=4, jk=jk)
+    assert str(model) == 'GIN(8, 4, num_layers=2)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 4)
+
+
+@pytest.mark.parametrize('v2', [False, True])
+def test_gat(v2):
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    model = GAT(8, 16, num_layers=2, out_channels=8, heads=2, v2=v2)
+    assert str(model) == 'GAT(8, 8, num_layers=2)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 8)
+
+
+def test_pna():
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    deg = ops.convert_to_tensor([0, 1, 2, 1])
+
+    model = PNA(
+        8, 16, num_layers=2, out_channels=4,
+        aggregators=['mean', 'min', 'max', 'std'],
+        scalers=['identity', 'amplification', 'attenuation'],
+        deg=deg,
+    )
+    assert str(model) == 'PNA(8, 4, num_layers=2)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 4)
+
+
+def test_edge_cnn():
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    model = EdgeCNN(8, 16, num_layers=2, out_channels=4)
+    assert str(model) == 'EdgeCNN(8, 4, num_layers=2)'
+
+    out = model(x, edge_index)
+    assert ops.shape(out) == (4, 4)
+
+
+
+def test_trim_to_layer_keeps_seed_outputs():
+    import numpy as np
+    from keras import ops
+    from k3_node.data import Data
+    from k3_node.loader import NeighborLoader
+    from k3_node.models import GraphSAGE
+
+    rng = np.random.default_rng(0)
+    data = Data(x=rng.random((200, 8)).astype("float32"), edge_index=rng.integers(0, 200, (2, 1200)), num_nodes=200)
+    batch = next(iter(NeighborLoader(data, num_neighbors=[5, 4, 3], batch_size=16)))
+    model = GraphSAGE(8, hidden_channels=16, num_layers=3, out_channels=4)
+    full = ops.convert_to_numpy(model(batch.x, batch.edge_index))[:16]
+    trimmed = ops.convert_to_numpy(model(batch.x, batch.edge_index, num_sampled_nodes_per_hop=batch.num_sampled_nodes,
+                                         num_sampled_edges_per_hop=batch.num_sampled_edges))[:16]
+    np.testing.assert_allclose(full, trimmed, atol=1e-5)