k3-node 1.0.0__py3-none-any.whl

k3_node/models/test_meta.py

@@ -0,0 +1,159 @@
+import pytest
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.models import MetaLayer
+from k3_node.layers.conv.utils import scatter
+
+
+# Global counter to verify model call counts
+_count = 0
+
+
+def test_meta_layer_repr():
+    """MetaLayer should have the correct string representation."""
+    assert str(MetaLayer()) == (
+        'MetaLayer(\n'
+        '  edge_model=None,\n'
+        '  node_model=None,\n'
+        '  global_model=None\n'
+        ')'
+    )
+
+
+def test_meta_layer_none_models():
+    """All None models should return unchanged tensors."""
+    x = ops.ones((20, 10))
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 0]])
+
+    model = MetaLayer()
+    x_out, edge_attr_out, u_out = model(x, edge_index)
+
+    assert ops.shape(x_out) == (20, 10)
+    assert edge_attr_out is None
+    assert u_out is None
+
+
+def test_meta_layer_call_counting():
+    """Verify that the correct submodels are called for each combination."""
+    global _count
+    _count = 0
+
+    def dummy_model(*args):
+        global _count
+        _count += 1
+        return None
+
+    x = ops.ones((20, 10))
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]])
+
+    for edge_model in (dummy_model, None):
+        for node_model in (dummy_model, None):
+            for global_model in (dummy_model, None):
+                model = MetaLayer(edge_model, node_model, global_model)
+                out = model(x, edge_index)
+                assert isinstance(out, tuple) and len(out) == 3
+
+    assert _count == 12  # 3 models × 4 combinations with at least 1 non-None
+
+
+def test_meta_layer_with_edge_model():
+    """Edge model should receive correct src/dst/edge_attr/u/batch args."""
+    received = {}
+
+    def edge_model(src, dst, edge_attr, u, batch):
+        received['src_shape'] = ops.shape(src)
+        received['dst_shape'] = ops.shape(dst)
+        received['edge_attr_shape'] = ops.shape(edge_attr) if edge_attr is not None else None
+        return edge_attr  # pass-through
+
+    x = ops.ones((5, 4))
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 3]])
+    edge_attr = ops.zeros((3, 7))
+
+    model = MetaLayer(edge_model=edge_model)
+    model(x, edge_index, edge_attr=edge_attr)
+
+    assert received['src_shape'] == (3, 4)
+    assert received['dst_shape'] == (3, 4)
+    assert received['edge_attr_shape'] == (3, 7)
+
+
+def test_meta_layer_full_example():
+    """Full graph network test matching PyG's test_meta_layer_example."""
+
+    class EdgeModel(keras.layers.Layer):
+        def __init__(self):
+            super().__init__()
+            self.mlp = keras.Sequential([
+                keras.layers.Dense(5),
+                keras.layers.ReLU(),
+                keras.layers.Dense(5),
+            ])
+
+        def call(self, src, dst, edge_attr, u, batch):
+            assert edge_attr is not None
+            assert u is not None
+            assert batch is not None
+            out = ops.concatenate([src, dst, edge_attr, ops.take(u, batch, axis=0)], axis=1)
+            return self.mlp(out)
+
+    class NodeModel(keras.layers.Layer):
+        def __init__(self):
+            super().__init__()
+            self.mlp1 = keras.Sequential([
+                keras.layers.Dense(10),
+                keras.layers.ReLU(),
+                keras.layers.Dense(10),
+            ])
+            self.mlp2 = keras.Sequential([
+                keras.layers.Dense(10),
+                keras.layers.ReLU(),
+                keras.layers.Dense(10),
+            ])
+
+        def call(self, x, edge_index, edge_attr, u, batch):
+            assert edge_attr is not None
+            assert u is not None
+            assert batch is not None
+            row = edge_index[0]
+            col = edge_index[1]
+            out = ops.concatenate([ops.take(x, row, axis=0), edge_attr], axis=1)
+            out = self.mlp1(out)
+            out = scatter(out, col, dim_size=ops.shape(x)[0], reduce="mean")
+            out = ops.concatenate([x, out, ops.take(u, batch, axis=0)], axis=1)
+            return self.mlp2(out)
+
+    class GlobalModel(keras.layers.Layer):
+        def __init__(self):
+            super().__init__()
+            self.mlp = keras.Sequential([
+                keras.layers.Dense(20),
+                keras.layers.ReLU(),
+                keras.layers.Dense(20),
+            ])
+
+        def call(self, x, edge_index, edge_attr, u, batch):
+            assert u is not None
+            assert batch is not None
+            batch_size = ops.shape(u)[0]
+            x_mean = scatter(x, batch, dim_size=batch_size, reduce="mean")
+            out = ops.concatenate([u, x_mean], axis=1)
+            return self.mlp(out)
+
+    op = MetaLayer(EdgeModel(), NodeModel(), GlobalModel())
+
+    x = ops.ones((20, 10))
+    edge_attr = ops.ones((40, 5))
+    u = ops.ones((2, 20))
+    batch = ops.convert_to_tensor([0] * 10 + [1] * 10)
+    row_idx = list(range(20)) + list(range(20))
+    col_idx = list(range(1, 20)) + [0] + list(range(1, 20)) + [0]
+    edge_index = ops.convert_to_tensor([row_idx, col_idx])
+
+    x_out, edge_attr_out, u_out = op(x, edge_index, edge_attr, u, batch)
+    assert ops.shape(x_out) == (20, 10)
+    assert ops.shape(edge_attr_out) == (40, 5)
+    assert ops.shape(u_out) == (2, 20)
+

k3_node/models/test_metapath2vec.py

@@ -0,0 +1,45 @@
+from keras import ops
+from k3_node.models import MetaPath2Vec
+
+
+def test_metapath2vec():
+    edge_index_dict = {
+        ("author", "writes", "paper"): ops.convert_to_tensor([[0, 1, 1], [0, 0, 1]], dtype="int64"),
+        ("paper", "written_by", "author"): ops.convert_to_tensor([[0, 0, 1], [0, 1, 1]], dtype="int64"),
+    }
+    metapath = [
+        ("author", "writes", "paper"),
+        ("paper", "written_by", "author"),
+    ]
+
+    model = MetaPath2Vec(
+        edge_index_dict=edge_index_dict,
+        embedding_dim=16,
+        metapath=metapath,
+        walk_length=2,
+        context_size=2,
+        walks_per_node=2,
+    )
+
+    # Test forward
+    out_author = model("author")
+    assert out_author.shape == (2, 16)
+
+    out_paper = model("paper")
+    assert out_paper.shape == (2, 16)
+
+    batch = ops.convert_to_tensor([0], dtype="int64")
+    out_batch = model("author", batch)
+    assert out_batch.shape == (1, 16)
+
+    # Test sampling
+    pos_rw = model._pos_sample(batch)
+    neg_rw = model._neg_sample(batch)
+    assert pos_rw.shape[1] == 2
+    assert neg_rw.shape[1] == 2
+
+    # Test loss
+    loss = model.loss(pos_rw, neg_rw)
+    assert loss.shape == ()
+    assert float(ops.convert_to_numpy(loss)) > 0
+

k3_node/models/test_mlp.py

@@ -0,0 +1,62 @@
+import pytest
+from keras import ops
+
+from k3_node.models import MLP
+
+
+@pytest.mark.parametrize("norm", ["batch_norm", None])
+@pytest.mark.parametrize("act_first", [False, True])
+@pytest.mark.parametrize("plain_last", [False, True])
+def test_mlp(norm, act_first, plain_last):
+    x = ops.ones((4, 16))
+
+    mlp = MLP([16, 32, 32, 64], norm=norm, act_first=act_first, plain_last=plain_last)
+    assert str(mlp) == "MLP(16, 32, 32, 64)"
+    out = mlp(x)
+    assert ops.shape(out) == (4, 64)
+
+    mlp2 = MLP(
+        16,
+        hidden_channels=32,
+        out_channels=64,
+        num_layers=3,
+        norm=norm,
+        act_first=act_first,
+        plain_last=plain_last,
+    )
+    assert ops.shape(mlp2(x)) == (4, 64)
+
+
+@pytest.mark.parametrize("norm", ["BatchNorm", "GraphNorm", "InstanceNorm", "LayerNorm"])
+def test_batch(norm):
+    x = ops.ones((3, 8))
+    batch = ops.convert_to_tensor([0, 0, 1], dtype="int64")
+
+    model = MLP(8, hidden_channels=16, out_channels=32, num_layers=2, norm=norm)
+    assert model.supports_norm_batch == (norm != "BatchNorm")
+
+    out = model(x, batch=batch)
+    assert ops.shape(out) == (3, 32)
+
+
+def test_mlp_return_emb():
+    x = ops.ones((4, 16))
+
+    mlp = MLP([16, 32, 1])
+
+    out, emb = mlp(x, return_emb=True)
+    assert ops.shape(out) == (4, 1)
+    assert ops.shape(emb) == (4, 32)
+
+    out, emb = mlp(x, return_emb=False)
+    assert ops.shape(out) == (4, 1)
+    assert emb is None
+
+    out = mlp(x)
+    assert ops.shape(out) == (4, 1)
+
+
+@pytest.mark.parametrize("plain_last", [False, True])
+def test_fine_grained_mlp(plain_last):
+    mlp = MLP([16, 32, 32, 64], dropout=[0.1, 0.2, 0.3], bias=[False, True, False], plain_last=plain_last)
+    assert ops.shape(mlp(ops.ones((4, 16)))) == (4, 64)

k3_node/models/test_mole_bert.py

@@ -0,0 +1,164 @@
+import os
+import numpy as np
+import pytest
+import keras.ops as ops
+try:
+    import torch
+except ImportError:
+    torch = None
+
+from k3_node.models.mole_bert import (
+    MoleBERT,
+    MoleBERTGNN,
+    MoleBERTGINConv,
+    load_mole_bert_weights,
+    download_mole_bert_checkpoint,
+)
+
+
+def _make_dummy_graph(num_nodes=5, num_edges=8, emb_dim=300):
+    np.random.seed(42)
+    x = np.stack(
+        [
+            np.random.randint(0, 119, size=num_nodes),
+            np.random.randint(0, 3, size=num_nodes),
+        ],
+        axis=1,
+    )
+    src = np.random.randint(0, num_nodes, size=num_edges)
+    dst = np.random.randint(0, num_nodes, size=num_edges)
+    edge_index = np.stack([src, dst], axis=0)
+    edge_attr = np.stack(
+        [
+            np.random.randint(0, 5, size=num_edges),
+            np.random.randint(0, 3, size=num_edges),
+        ],
+        axis=1,
+    )
+    batch = np.array([0, 0, 0, 1, 1], dtype=np.int64)
+
+    return (
+        ops.convert_to_tensor(x, dtype="int64"),
+        ops.convert_to_tensor(edge_index, dtype="int64"),
+        ops.convert_to_tensor(edge_attr, dtype="int64"),
+        ops.convert_to_tensor(batch, dtype="int64"),
+    )
+
+
+def test_mole_bert_gin_conv():
+    emb_dim = 64
+    conv = MoleBERTGINConv(emb_dim=emb_dim)
+    conv.build(None)
+
+    num_nodes = 4
+    x = ops.convert_to_tensor(np.random.randn(num_nodes, emb_dim).astype(np.float32))
+    edge_index = ops.convert_to_tensor(
+        np.array([[0, 1, 1, 2], [1, 0, 2, 1]]), dtype="int64"
+    )
+    edge_attr = ops.convert_to_tensor(
+        np.array([[0, 0], [0, 0], [1, 0], [1, 0]]), dtype="int64"
+    )
+
+    out = conv(x, edge_index, edge_attr)
+    assert ops.shape(out) == (num_nodes, emb_dim)
+
+
+def test_mole_bert_gnn_jk_modes():
+    x, edge_index, edge_attr, _ = _make_dummy_graph(num_nodes=5, num_edges=8, emb_dim=32)
+    emb_dim = 32
+
+    # 1. JK = 'last'
+    gnn_last = MoleBERTGNN(num_layer=3, emb_dim=emb_dim, JK="last")
+    gnn_last.build(None)
+    out_last = gnn_last(x, edge_index, edge_attr)
+    assert ops.shape(out_last) == (5, emb_dim)
+
+    # 2. JK = 'concat'
+    gnn_concat = MoleBERTGNN(num_layer=3, emb_dim=emb_dim, JK="concat")
+    gnn_concat.build(None)
+    out_concat = gnn_concat(x, edge_index, edge_attr)
+    assert ops.shape(out_concat) == (5, 4 * emb_dim)
+
+    # 3. JK = 'sum'
+    gnn_sum = MoleBERTGNN(num_layer=3, emb_dim=emb_dim, JK="sum")
+    gnn_sum.build(None)
+    out_sum = gnn_sum(x, edge_index, edge_attr)
+    assert ops.shape(out_sum) == (5, emb_dim)
+
+    # 4. JK = 'max'
+    gnn_max = MoleBERTGNN(num_layer=3, emb_dim=emb_dim, JK="max")
+    gnn_max.build(None)
+    out_max = gnn_max(x, edge_index, edge_attr)
+    assert ops.shape(out_max) == (5, emb_dim)
+
+
+def test_mole_bert_pooling_and_prediction():
+    x, edge_index, edge_attr, batch = _make_dummy_graph(num_nodes=5, num_edges=8, emb_dim=32)
+    emb_dim = 32
+
+    # Multi-task graph classification (e.g. ClinTox: 2 tasks, Tox21: 12 tasks)
+    for pooling in ["mean", "sum", "max"]:
+        model = MoleBERT(
+            num_layer=3,
+            emb_dim=emb_dim,
+            num_tasks=2,
+            JK="last",
+            graph_pooling=pooling,
+        )
+        model.build(None)
+
+        logits, node_rep = model((x, edge_index, edge_attr, batch))
+        assert ops.shape(logits) == (2, 2)
+        assert ops.shape(node_rep) == (5, emb_dim)
+
+
+def test_mole_bert_synthetic_checkpoint_load(tmp_path):
+    if torch is None:
+        pytest.skip("PyTorch is required for checkpoint loading test")
+    emb_dim = 16
+    num_layer = 3
+    model = MoleBERT(num_layer=num_layer, emb_dim=emb_dim, num_tasks=1)
+    model.build(None)
+
+    state_dict = {
+        "x_embedding1.weight": torch.randn(120, emb_dim),
+        "x_embedding2.weight": torch.randn(3, emb_dim),
+    }
+    for l in range(num_layer):
+        state_dict[f"gnns.{l}.edge_embedding1.weight"] = torch.randn(6, emb_dim)
+        state_dict[f"gnns.{l}.edge_embedding2.weight"] = torch.randn(3, emb_dim)
+        state_dict[f"gnns.{l}.mlp.0.weight"] = torch.randn(2 * emb_dim, emb_dim)
+        state_dict[f"gnns.{l}.mlp.0.bias"] = torch.zeros(2 * emb_dim)
+        state_dict[f"gnns.{l}.mlp.2.weight"] = torch.randn(emb_dim, 2 * emb_dim)
+        state_dict[f"gnns.{l}.mlp.2.bias"] = torch.zeros(emb_dim)
+        state_dict[f"batch_norms.{l}.weight"] = torch.ones(emb_dim)
+        state_dict[f"batch_norms.{l}.bias"] = torch.zeros(emb_dim)
+        state_dict[f"batch_norms.{l}.running_mean"] = torch.zeros(emb_dim)
+        state_dict[f"batch_norms.{l}.running_var"] = torch.ones(emb_dim)
+
+    ckpt_file = str(tmp_path / "synthetic_mole_bert.pth")
+    torch.save(state_dict, ckpt_file)
+
+    load_mole_bert_weights(model, ckpt_file)
+
+    x, edge_index, edge_attr, batch = _make_dummy_graph(num_nodes=5, num_edges=8, emb_dim=emb_dim)
+    logits, node_rep = model((x, edge_index, edge_attr, batch))
+    assert ops.shape(logits) == (2, 1)
+    assert ops.shape(node_rep) == (5, emb_dim)
+
+
+def test_mole_bert_official_checkpoint_load():
+    torch = pytest.importorskip("torch")
+    ckpt_path = "Mole-BERT/model_gin/Mole-BERT.pth"
+    if not os.path.exists(ckpt_path):
+        ckpt_path = download_mole_bert_checkpoint()
+
+    model = MoleBERT(num_layer=5, emb_dim=300, num_tasks=1)
+    model.build(None)
+    load_mole_bert_weights(model, ckpt_path)
+
+    x, edge_index, edge_attr, batch = _make_dummy_graph(num_nodes=5, num_edges=8, emb_dim=300)
+    logits, node_rep = model((x, edge_index, edge_attr, batch))
+    assert ops.shape(logits) == (2, 1)
+    assert ops.shape(node_rep) == (5, 300)
+

k3_node/models/test_neural_fingerprint.py

@@ -0,0 +1,13 @@
+from keras import ops, random
+from k3_node.models import NeuralFingerprint
+
+
+def test_neural_fingerprint():
+    model = NeuralFingerprint(in_channels=16, hidden_channels=32, out_channels=8, num_layers=3)
+    x = random.normal((6, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3, 4], [1, 2, 0, 4, 5]], dtype="int64")
+    batch = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int64")
+
+    out = model(x, edge_index, batch)
+    assert out.shape == (2, 8)
+

k3_node/models/test_node2vec.py

@@ -0,0 +1,57 @@
+from keras import ops
+from k3_node.models import Node2Vec
+
+
+def test_node2vec():
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 0, 2],
+        [1, 2, 3, 0, 2, 0],
+    ], dtype="int64")
+
+    model = Node2Vec(
+        edge_index,
+        embedding_dim=16,
+        walk_length=4,
+        context_size=3,
+        walks_per_node=2,
+        num_negative_samples=1,
+    )
+
+    # Test forward
+    out = model()
+    assert out.shape == (4, 16)
+
+    batch = ops.convert_to_tensor([0, 1], dtype="int64")
+    out_batch = model(batch)
+    assert out_batch.shape == (2, 16)
+
+    # Test pos and neg sampling
+    pos_rw = model.pos_sample(batch)
+    assert pos_rw.shape[1] == 3
+
+    neg_rw = model.neg_sample(batch)
+    assert neg_rw.shape[1] == 3
+
+    # Test loss
+    loss = model.loss(pos_rw, neg_rw)
+    assert loss.shape == ()
+    assert float(ops.convert_to_numpy(loss)) > 0
+
+
+
+def test_node2vec_fit_and_biased_walks():
+    import numpy as np
+    from k3_node.models import Node2Vec
+
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, 20, size=(2, 80))
+    edge_index = np.concatenate([edge_index, edge_index[::-1]], axis=1)
+    import keras
+    model = Node2Vec(edge_index, embedding_dim=8, walk_length=5, context_size=3, walks_per_node=2, num_nodes=20)
+    model.compile(keras.optimizers.Adam(0.01))
+    history = model.fit(epochs=2, batch_size=8, verbose=0)
+    assert len(history["loss"]) == 2 and np.isfinite(history["loss"]).all()
+    # With a tiny p the walk (almost) always steps back to where it came from
+    model.p, model.q = 1e-6, 1.0
+    walk = model._random_walk(int(edge_index[0, 0]))
+    assert all(walk[i] == walk[i - 2] for i in range(2, len(walk)))

k3_node/models/test_pmlp.py

@@ -0,0 +1,81 @@
+import pytest
+from keras import ops
+
+from k3_node.models import PMLP
+
+
+def test_pmlp():
+    x = ops.ones((4, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    pmlp = PMLP(in_channels=16, hidden_channels=32, out_channels=2, num_layers=4)
+    assert str(pmlp) == 'PMLP(16, 2, num_layers=4)'
+
+    pmlp.training = True
+    out = pmlp(x)
+    assert ops.shape(out) == (4, 2)
+
+    pmlp.training = False
+    out = pmlp(x, edge_index)
+    assert ops.shape(out) == (4, 2)
+
+
+def test_pmlp_raises_without_edge_index():
+    """Should raise ValueError when edge_index is missing during inference."""
+    x = ops.ones((4, 16))
+    pmlp = PMLP(in_channels=16, hidden_channels=32, out_channels=2, num_layers=4)
+
+    with pytest.raises(ValueError, match="'edge_index' needs to be present"):
+        pmlp.training = False
+        pmlp(x)
+
+
+def test_pmlp_call_training_override():
+    """The training kwarg at call-time should override self.training."""
+    x = ops.ones((4, 16))
+    edge_index = ops.convert_to_tensor([[0, 1], [1, 0]])
+    pmlp = PMLP(in_channels=16, hidden_channels=32, out_channels=4, num_layers=2)
+
+    # Instance says training=True but call says training=False (requires edge_index)
+    pmlp.training = True
+    out = pmlp(x, edge_index, training=False)
+    assert ops.shape(out) == (4, 4)
+
+    # Instance says training=False but call says training=True (no edge_index needed)
+    pmlp.training = False
+    out = pmlp(x, training=True)
+    assert ops.shape(out) == (4, 4)
+
+
+@pytest.mark.parametrize("norm", [True, False])
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize("dropout", [0.0, 0.3])
+def test_pmlp_configs(norm, bias, dropout):
+    x = ops.ones((8, 10))
+    pmlp = PMLP(
+        in_channels=10,
+        hidden_channels=20,
+        out_channels=5,
+        num_layers=3,
+        norm=norm,
+        bias=bias,
+        dropout=dropout,
+    )
+    pmlp.training = True
+    out = pmlp(x)
+    assert ops.shape(out) == (8, 5)
+
+
+def test_pmlp_single_layer():
+    """Edge case: num_layers=1 should work (no hidden layers, direct in→out)."""
+    x = ops.ones((4, 8))
+    edge_index = ops.convert_to_tensor([[0, 1], [1, 0]])
+    pmlp = PMLP(in_channels=8, hidden_channels=16, out_channels=3, num_layers=1)
+    pmlp.training = True
+    out = pmlp(x)
+    assert ops.shape(out) == (4, 3)
+
+    pmlp.training = False
+    out = pmlp(x, edge_index)
+    assert ops.shape(out) == (4, 3)
+

k3_node/models/test_polynormer.py

@@ -0,0 +1,104 @@
+import pytest
+from keras import ops
+
+from k3_node.models import Polynormer
+
+
+@pytest.mark.parametrize('local_attn', [True, False])
+@pytest.mark.parametrize('qk_shared', [True, False])
+@pytest.mark.parametrize('pre_ln', [True, False])
+@pytest.mark.parametrize('post_bn', [True, False])
+def test_polynormer_local(local_attn, qk_shared, pre_ln, post_bn):
+    """Local-mode Polynormer (default _global=False) should output [N, out_channels]."""
+    x = ops.ones((10, 16))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+        [1, 2, 3, 4, 0, 6, 7, 8, 9, 5],
+    ])
+    batch = ops.convert_to_tensor([0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
+
+    model = Polynormer(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=40,
+        local_layers=2,
+        global_layers=1,
+        qk_shared=qk_shared,
+        pre_ln=pre_ln,
+        post_bn=post_bn,
+        local_attn=local_attn,
+        heads=2,
+    )
+    out = model(x, edge_index, batch)
+    assert ops.shape(out) == (10, 40)
+
+
+@pytest.mark.parametrize('local_attn', [True, False])
+@pytest.mark.parametrize('qk_shared', [True, False])
+def test_polynormer_global(local_attn, qk_shared):
+    """Global-mode Polynormer (_global=True) should output [N, out_channels]."""
+    x = ops.ones((10, 16))
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+        [1, 2, 3, 4, 0, 6, 7, 8, 9, 5],
+    ])
+    batch = ops.convert_to_tensor([0, 0, 0, 0, 1, 1, 1, 1, 1, 1])
+
+    model = Polynormer(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=40,
+        local_layers=2,
+        global_layers=1,
+        qk_shared=qk_shared,
+        local_attn=local_attn,
+    )
+    model._global = True
+    out = model(x, edge_index, batch)
+    assert ops.shape(out) == (10, 40)
+
+
+def test_polynormer_log_softmax_sums():
+    """Output of Polynormer is log_softmax, so exp(out).sum(axis=-1) ≈ 1."""
+    import numpy as np
+    x = ops.ones((6, 8))
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 0]])
+    batch = ops.convert_to_tensor([0, 0, 0, 0, 0, 0])
+
+    model = Polynormer(
+        in_channels=8,
+        hidden_channels=16,
+        out_channels=5,
+        local_layers=1,
+        global_layers=1,
+    )
+    out = model(x, edge_index, batch)
+    assert ops.shape(out) == (6, 5)
+
+    out_np = ops.convert_to_numpy(out)
+    sums = np.exp(out_np).sum(axis=-1)
+    assert np.allclose(sums, 1.0, atol=1e-5)
+
+
+def test_polynormer_global_toggle():
+    """Switching between local and global modes on the same model instance."""
+    x = ops.ones((8, 16))
+    edge_index = ops.convert_to_tensor([[0, 1, 2, 3], [1, 2, 3, 0]])
+    batch = ops.convert_to_tensor([0, 0, 0, 0, 1, 1, 1, 1])
+
+    model = Polynormer(
+        in_channels=16,
+        hidden_channels=16,
+        out_channels=10,
+        local_layers=1,
+        global_layers=1,
+    )
+
+    model._global = False
+    out_local = model(x, edge_index, batch)
+    assert ops.shape(out_local) == (8, 10)
+
+    model._global = True
+    out_global = model(x, edge_index, batch)
+    assert ops.shape(out_global) == (8, 10)
+