k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/ssg_conv.py

@@ -0,0 +1,101 @@
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm, is_tracing
+
+
+class SSGConv(MessagePassing):
+    r"""The simple spectral graph convolutional operator from the
+    `"Simple Spectral Graph Convolution" <https://arxiv.org/abs/2109.07191>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        alpha: Teleport probability :math:`\alpha`.
+        K: Number of hops :math:`K`. (default: ``1``)
+        cached: If set to :obj:`True`, the layer will cache normalization coefficients.
+            (default: ``False``)
+        add_self_loops: If set to :obj:`False`, will not add self-loops.
+            (default: ``True``)
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SSGConv
+
+        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
+
+        layer = SSGConv(in_channels=8, out_channels=16, alpha=0.1, K=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        alpha: float,
+        K: int = 1,
+        cached: bool = False,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.alpha = alpha
+        self.K = K
+        self.cached = cached
+        self.add_self_loops = add_self_loops
+        self.use_bias = bias
+
+        self.lin = layers.Dense(out_channels, use_bias=bias)
+        self._cached_edge_index = None
+        self._cached_norm = None
+
+    def build(self, input_shape):
+        feat_shape = input_shape[0] if isinstance(input_shape, (tuple, list)) and isinstance(input_shape[0], (tuple, list)) else input_shape
+        self.lin.build(feat_shape)
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        if self.cached and self._cached_edge_index is not None:
+            edge_index = self._cached_edge_index
+            edge_weight = self._cached_norm
+        else:
+            num_nodes = x.shape[self.node_dim] if hasattr(x, "shape") and x.shape[self.node_dim] is not None else ops.shape(x)[self.node_dim]
+            edge_index, edge_weight = gcn_norm(
+                edge_index,
+                edge_weight,
+                num_nodes=num_nodes,
+                add_self_loops=self.add_self_loops,
+                flow=self.flow,
+                dtype=x.dtype,
+            )
+            if self.cached and not is_tracing(edge_index):
+                self._cached_edge_index = edge_index
+                self._cached_norm = edge_weight
+
+        out = self.alpha * x
+        h = x
+        for _ in range(self.K):
+            h = self.propagate(edge_index, x=h, edge_weight=edge_weight)
+            out = out + ((1.0 - self.alpha) / self.K) * h
+
+        return self.lin(out)
+
+    def message(self, x_j, edge_weight=None):
+        if edge_weight is None:
+            return x_j
+        return ops.expand_dims(edge_weight, -1) * x_j
+

k3_node/layers/conv/supergat_conv.py

@@ -0,0 +1,195 @@
+import math
+import keras
+from keras import ops
+from keras.layers import Dense, Dropout
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import (
+    add_self_loops,
+    extend_mask_for_self_loops,
+    mask_edge_logits,
+    remove_self_loops_masked,
+    softmax,
+)
+from k3_node.ops.creation import full
+
+
+class SuperGATConv(MessagePassing):
+    r"""The self-supervised graph attentional operator from the
+    `"How to Find Your Friendly Neighborhood: Graph Attention Design with Self-Supervision"
+    <https://openreview.net/forum?id=Wi5KUNlqWty>`_ paper.
+
+    Args:
+        attention_type (str): ``"MX"`` (mixed GO/DP) or ``"SD"`` (scaled dot-product).
+        neg_sample_ratio (float): Negative (random) pairs per positive edge in the attention loss.
+        edge_sample_ratio (float): Fraction of edges used as positives in the attention loss.
+        attention_loss_weight (float): If positive, the self-supervised attention loss is added to
+            the model loss (``model.losses``) with this weight while training, so ``fit`` optimizes
+            it automatically. PyG's example uses ``4.0``. (default: ``0.0``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SuperGATConv
+
+        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
+
+        layer = SuperGATConv(in_channels=8, out_channels=16, heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 32)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        dropout: float = 0.0,
+        add_self_loops: bool = True,
+        bias: bool = True,
+        attention_type: str = "MX",
+        neg_sample_ratio: float = 0.5,
+        edge_sample_ratio: float = 1.0,
+        is_undirected: bool = False,
+        attention_loss_weight: float = 0.0,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(node_dim=0, **kwargs)
+
+        assert attention_type in ["MX", "SD"]
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout_rate = dropout
+        self.add_self_loops = add_self_loops
+        self.attention_type = attention_type
+        self.neg_sample_ratio = neg_sample_ratio
+        self.edge_sample_ratio = edge_sample_ratio
+        self.is_undirected = is_undirected
+        self.use_bias = bias
+
+        self.attention_loss_weight = attention_loss_weight
+        self.lin = Dense(heads * out_channels, use_bias=False)
+        self.dropout = Dropout(dropout)
+        self.seed_generator = keras.random.SeedGenerator()
+        self._last_attention_loss = None
+
+        if self.attention_type == "MX":
+            self.att_l = self.add_weight(
+                shape=(1, heads, out_channels),
+                initializer="glorot_uniform",
+                name="att_l",
+            )
+            self.att_r = self.add_weight(
+                shape=(1, heads, out_channels),
+                initializer="glorot_uniform",
+                name="att_r",
+            )
+        else:
+            self.att_l = None
+            self.att_r = None
+
+        if bias:
+            out_dim = heads * out_channels if concat else out_channels
+            self.bias = self.add_weight(
+                shape=(out_dim,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+
+    def build(self, input_shape=None):
+        self.lin.build((None, self.in_channels))
+        self.built = True
+
+    def call(self, inputs, edge_index=None, training=None, **kwargs):
+        if edge_index is None:
+            if isinstance(inputs, (list, tuple)) and len(inputs) == 2:
+                x, edge_index = inputs
+            else:
+                raise ValueError("Expected (x, edge_index) or x and edge_index")
+        else:
+            x = inputs
+
+        if not self.built:
+            self.build()
+
+        num_nodes = ops.shape(x)[0]
+        keep_mask = None
+        if self.add_self_loops:
+            edge_index, _, keep_mask = remove_self_loops_masked(edge_index)
+            edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)
+            keep_mask = extend_mask_for_self_loops(keep_mask, num_nodes)
+
+        x = self.lin(x)
+        x = ops.reshape(x, (-1, self.heads, self.out_channels))
+
+        out = self.propagate(edge_index, x=x, keep_mask=keep_mask, training=training, size=(num_nodes, num_nodes))
+
+        if training:
+            loss = self._attention_loss(x, edge_index, num_nodes)
+            if self.attention_loss_weight:
+                self.add_loss(self.attention_loss_weight * loss)
+            from k3_node.layers.conv.utils import is_tracing
+            self._last_attention_loss = None if is_tracing(loss) else loss
+
+        if self.concat:
+            out = ops.reshape(out, (-1, self.heads * self.out_channels))
+        else:
+            out = ops.mean(out, axis=1)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        return out
+
+    def message(self, x_i, x_j, index=None, size_i=None, keep_mask=None, training=None):
+        if self.attention_type == "MX":
+            logits = ops.sum(x_i * x_j, axis=-1)
+            alpha = ops.sum(x_j * self.att_l, axis=-1) + ops.sum(x_i * self.att_r, axis=-1)
+            alpha = alpha * ops.sigmoid(logits)
+        else:  # SD
+            alpha = ops.sum(x_i * x_j, axis=-1) / math.sqrt(self.out_channels)
+
+        alpha = ops.leaky_relu(alpha, negative_slope=self.negative_slope)
+        alpha = mask_edge_logits(alpha, keep_mask)
+        alpha = softmax(alpha, index, num_nodes=size_i, dim=0)
+        alpha = self.dropout(alpha, training=training)
+        return x_j * ops.expand_dims(alpha, -1)
+
+    def _attention_logits(self, x_i, x_j):
+        logits = ops.sum(x_i * x_j, axis=-1)
+        if self.attention_type == "SD":
+            logits = logits / math.sqrt(self.out_channels)
+        return logits
+
+    def _attention_loss(self, x, edge_index, num_nodes):
+        # Self-supervised attention loss: attention logits should separate real edges (label 1)
+        # from random node pairs (label 0). Edges are kept with probability `edge_sample_ratio` and
+        # random pairs with probability `neg_sample_ratio * edge_sample_ratio` (the expected counts
+        # used by PyG); weighting instead of slicing keeps all shapes static for XLA / jax.jit.
+        edge_index = ops.cast(edge_index, "int32")
+        num_edges = ops.shape(edge_index)[1]
+        neg = keras.random.randint(ops.shape(edge_index), 0, num_nodes, seed=self.seed_generator, dtype="int32")
+        pairs = ops.concatenate([edge_index, neg], axis=1)
+        logits = ops.mean(self._attention_logits(ops.take(x, pairs[1], axis=0), ops.take(x, pairs[0], axis=0)), axis=-1)
+        labels = ops.concatenate([ops.ones((num_edges,)), ops.zeros((num_edges,))])
+        keep = keras.random.uniform(ops.shape(labels), seed=self.seed_generator) < ops.concatenate([
+            full((num_edges,), self.edge_sample_ratio),
+            full((num_edges,), self.neg_sample_ratio * self.edge_sample_ratio),
+        ])
+        weights = ops.cast(keep, "float32")
+        losses = ops.binary_crossentropy(labels, logits, from_logits=True)
+        return ops.sum(losses * weights) / ops.maximum(ops.sum(weights), 1.0)
+
+    def get_attention_loss(self):
+        r"""The self-supervised attention loss of the last training call (as in PyG)."""
+        return self._last_attention_loss

k3_node/layers/conv/tag_conv.py

@@ -0,0 +1,98 @@
+from keras import layers, ops
+
+from k3_node.layers.conv.message_passing import MessagePassing
+from k3_node.layers.conv.utils import gcn_norm
+
+
+class TAGConv(MessagePassing):
+    r"""The topology adaptive graph convolutional operator from the
+    `"Topology Adaptive Graph Convolutional Networks"
+    <https://arxiv.org/abs/1710.10370>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample.
+        out_channels: Size of each output sample.
+        K: Number of hops :math:`K`. (default: ``3``)
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+        normalize: Whether to apply symmetric normalization. (default: ``True``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import TAGConv
+
+        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
+
+        layer = TAGConv(in_channels=8, out_channels=16, K=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+
+    weighted_sum_message = True
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        K: int = 3,
+        bias: bool = True,
+        normalize: bool = True,
+        **kwargs,
+    ):
+        super().__init__(aggr="add", **kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.K = K
+        self.normalize = normalize
+        self.use_bias = bias
+
+        self.lins = [layers.Dense(out_channels, use_bias=False) for _ in range(K + 1)]
+
+    def build(self, input_shape):
+        feat_shape = input_shape[0] if isinstance(input_shape, (tuple, list)) and isinstance(input_shape[0], (tuple, list)) else input_shape
+        for lin in self.lins:
+            lin.build(feat_shape)
+
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(self.out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        if self.normalize:
+            num_nodes = x.shape[self.node_dim] if hasattr(x, "shape") and x.shape[self.node_dim] is not None else ops.shape(x)[self.node_dim]
+            edge_index, edge_weight = gcn_norm(
+                edge_index,
+                edge_weight,
+                num_nodes=num_nodes,
+                add_self_loops=False,
+                flow=self.flow,
+                dtype=x.dtype,
+            )
+
+        out = self.lins[0](x)
+        h = x
+        for k in range(1, self.K + 1):
+            h = self.propagate(edge_index, x=h, edge_weight=edge_weight)
+            out = out + self.lins[k](h)
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def message(self, x_j, edge_weight=None):
+        if edge_weight is None:
+            return x_j
+        return ops.expand_dims(edge_weight, -1) * x_j
+

k3_node/layers/conv/test_backend_consistency.py

@@ -0,0 +1,164 @@
+"""Cross-backend and compiled-mode consistency tests for convolution layers.
+
+Each layer is run with identical seeded inputs and weights, and its eager output is
+compared against golden outputs recorded with the torch backend (the backend that
+``tests_reference`` validates against PyG). On TensorFlow and JAX the output under
+``jit_compile=True`` must also match the eager output.
+
+The input graph deliberately contains self-loops: layers that remove and re-add
+self-loops must handle them without dynamic shapes when compiled.
+
+To regenerate the golden file after an intentional numerical change::
+
+    KERAS_BACKEND=torch python k3_node/layers/conv/test_backend_consistency.py
+"""
+import os
+import os.path as osp
+
+import numpy as np
+import pytest
+import keras
+from keras import layers
+
+import k3_node.layers as L
+
+GOLDEN_PATH = osp.join(osp.dirname(__file__), "testdata", "backend_consistency_golden.npz")
+
+N, E, C, O = 12, 30, 8, 5
+
+
+def _graph():
+    rng = np.random.default_rng(0)
+    edge_index = rng.integers(0, N, size=(2, E))
+    loops = np.array([[0, 3, 7], [0, 3, 7]])  # guarantee pre-existing self-loops
+    return {
+        "x": rng.standard_normal((N, C)).astype("float32"),
+        "pos": rng.standard_normal((N, 3)).astype("float32"),
+        "normal": rng.standard_normal((N, 3)).astype("float32"),
+        "edge_index": np.concatenate([edge_index, loops], axis=1).astype("int32"),
+    }
+
+
+def _mlp(units):
+    return keras.Sequential([layers.Dense(units, activation="relu"), layers.Dense(units)])
+
+
+def _xe(layer, g):
+    return layer(g["x"], g["edge_index"])
+
+
+# name -> (layer factory, call function)
+CASES = {
+    "GCNConv": (lambda: L.GCNConv(C, O), _xe),
+    "GATConv": (lambda: L.GATConv(C, O, heads=2), _xe),
+    "GATv2Conv": (lambda: L.GATv2Conv(C, O, heads=2), _xe),
+    "SuperGATConv": (lambda: L.SuperGATConv(C, O, heads=2), _xe),
+    "ClusterGCNConv": (lambda: L.ClusterGCNConv(C, O), _xe),
+    "AGNNConv": (lambda: L.AGNNConv(), _xe),
+    "FeaStConv": (lambda: L.FeaStConv(C, O, heads=2), _xe),
+    "PointNetConv": (lambda: L.PointNetConv(local_nn=_mlp(O)), lambda l, g: l(g["x"], g["pos"], g["edge_index"])),
+    "PointTransformerConv": (
+        lambda: L.PointTransformerConv(C, O),
+        lambda l, g: l(g["x"], g["pos"], g["edge_index"]),
+    ),
+    "PPFConv": (
+        lambda: L.PPFConv(local_nn=_mlp(O)),
+        lambda l, g: l(g["x"], g["pos"], g["normal"], g["edge_index"]),
+    ),
+    "SAGEConv": (lambda: L.SAGEConv(C, O), _xe),
+    "GraphConv": (lambda: L.GraphConv(C, O), _xe),
+    "TransformerConv": (lambda: L.TransformerConv(C, O, heads=2), _xe),
+    "ChebConv": (lambda: L.ChebConv(C, O, K=3), _xe),
+    "TAGConv": (lambda: L.TAGConv(C, O, K=2), _xe),
+    "SGConv": (lambda: L.SGConv(C, O, K=2), _xe),
+    "LEConv": (lambda: L.LEConv(C, O), _xe),
+    "ResGatedGraphConv": (lambda: L.ResGatedGraphConv(C, O), _xe),
+    "GENConv": (lambda: L.GENConv(C, O), _xe),
+    "GINConv": (lambda: L.GINConv(_mlp(O)), _xe),
+    "EdgeConv": (lambda: L.EdgeConv(keras.Sequential([layers.Dense(O)])), _xe),
+    "MFConv": (lambda: L.MFConv(C, O), _xe),
+    "FiLMConv": (lambda: L.FiLMConv(C, O), _xe),
+    "GeneralConv": (lambda: L.GeneralConv(C, O), _xe),
+    "PNAConv": (
+        lambda: L.PNAConv(
+            C,
+            O,
+            aggregators=["mean", "max", "min", "std"],
+            scalers=["identity", "amplification"],
+            deg=np.array([0, 2, 4, 3, 2, 1], "int32"),
+        ),
+        _xe,
+    ),
+}
+
+
+class ConsistencyWrapper(keras.Model):
+    def __init__(self, layer, call_fn):
+        super().__init__()
+        self.layer = layer
+        self.call_fn = call_fn
+
+    def call(self, inputs):
+        out = self.call_fn(self.layer, inputs)
+        return out[0] if isinstance(out, (tuple, list)) else out
+
+
+def _build(name):
+    factory, call_fn = CASES[name]
+    g = _graph()
+    model = ConsistencyWrapper(factory(), call_fn)
+    model(g)
+    rng = np.random.default_rng(123)
+    weights = []
+    for v in model.weights:
+        w = rng.standard_normal(v.shape) * 0.3
+        if "moving_variance" in v.path:
+            w = np.abs(w) + 0.5
+        weights.append(w.astype(v.dtype))
+    model.set_weights(weights)
+    shapes = np.array([str(tuple(v.shape)) for v in model.weights])
+    return model, g, shapes
+
+
+def _eager(model, g):
+    return keras.ops.convert_to_numpy(model(g))
+
+
+@pytest.fixture(scope="module")
+def golden():
+    if not osp.exists(GOLDEN_PATH):
+        pytest.fail(f"Golden file missing; regenerate it (see module docstring): {GOLDEN_PATH}")
+    return np.load(GOLDEN_PATH)
+
+
+@pytest.mark.parametrize("name", sorted(CASES))
+def test_eager_matches_torch_golden(name, golden):
+    model, g, shapes = _build(name)
+    assert list(shapes) == list(golden[f"{name}/shapes"]), (
+        f"{name}: weight layout differs from the torch backend, so weights cannot be shared"
+    )
+    np.testing.assert_allclose(_eager(model, g), golden[f"{name}/out"], rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.skipif(
+    keras.backend.backend() not in ("tensorflow", "jax"),
+    reason="jit_compile=True means XLA only on the TensorFlow and JAX backends",
+)
+@pytest.mark.parametrize("name", sorted(CASES))
+def test_jit_matches_eager(name):
+    model, g, _ = _build(name)
+    expected = _eager(model, g)
+    model.compile(jit_compile=True)
+    np.testing.assert_allclose(np.asarray(model.predict_on_batch(g)), expected, rtol=1e-4, atol=1e-5)
+
+
+if __name__ == "__main__":
+    assert keras.backend.backend() == "torch", "Golden outputs must be generated with KERAS_BACKEND=torch"
+    arrays = {}
+    for case in sorted(CASES):
+        m, graph, layer_shapes = _build(case)
+        arrays[f"{case}/out"] = _eager(m, graph)
+        arrays[f"{case}/shapes"] = layer_shapes
+    os.makedirs(osp.dirname(GOLDEN_PATH), exist_ok=True)
+    np.savez(GOLDEN_PATH, **arrays)
+    print(f"Wrote {len(CASES)} golden outputs to {GOLDEN_PATH}")