k3-node 1.0.0__py3-none-any.whl

k3_node/layers/aggr/gmt.py

@@ -0,0 +1,89 @@
+from typing import Optional
+from keras import ops
+
+from .base import Aggregation
+from .utils import PoolingByMultiheadAttention, SetAttentionBlock
+
+
+class GraphMultisetTransformer(Aggregation):
+    r"""The Graph Multiset Transformer pooling operator from the
+    `"Accurate Learning of Graph Representations
+    with Graph Multiset Pooling" <https://arxiv.org/abs/2102.11533>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GraphMultisetTransformer
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = GraphMultisetTransformer(channels=8, k=2)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        k: int,
+        num_encoder_blocks: int = 1,
+        heads: int = 1,
+        layer_norm: bool = False,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.k = k
+        self.heads = heads
+        self.layer_norm = layer_norm
+        self.dropout = dropout
+
+        self.pma1 = PoolingByMultiheadAttention(channels, k, heads, layer_norm, dropout)
+        self.encoders = [
+            SetAttentionBlock(channels, heads, layer_norm, dropout)
+            for _ in range(num_encoder_blocks)
+        ]
+        self.pma2 = PoolingByMultiheadAttention(channels, 1, heads, layer_norm, dropout)
+
+    def reset_parameters(self):
+        self.pma1.reset_parameters()
+        for encoder in self.encoders:
+            encoder.reset_parameters()
+        self.pma2.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+        training: bool = False,
+        **kwargs,
+    ):
+        x_dense, mask = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=max_num_elements,
+        )
+
+        x_dense = self.pma1(x_dense, mask=mask, training=training)
+
+        for encoder in self.encoders:
+            x_dense = encoder(x_dense, training=training)
+
+        x_dense = self.pma2(x_dense, training=training)
+
+        # Output shape: [B, channels]
+        return ops.squeeze(x_dense, axis=1)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.channels}, k={self.k}, "
+            f"heads={self.heads}, layer_norm={self.layer_norm}, "
+            f"dropout={self.dropout})"
+        )
+

k3_node/layers/aggr/gru.py

@@ -0,0 +1,58 @@
+from typing import Optional
+from keras import layers
+
+from .base import Aggregation
+
+
+class GRUAggregation(Aggregation):
+    r"""Performs GRU aggregation in which the elements to aggregate are
+    interpreted as a sequence, as described in the `"Graph Neural Networks
+    with Adaptive Readouts" <https://arxiv.org/abs/2211.04952>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GRUAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = GRUAggregation(in_channels=8, out_channels=16)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.gru = layers.GRU(out_channels, return_sequences=False)
+
+    def build(self, input_shape=None):
+        self.gru.build((None, None, self.in_channels))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        self.gru.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+        training: bool = False,
+        **kwargs,
+    ):
+        x_dense, _ = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=max_num_elements,
+        )
+        return self.gru(x_dense, training=training)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
+

k3_node/layers/aggr/lcm.py

@@ -0,0 +1,143 @@
+from math import ceil, log2
+from typing import Optional
+from keras import layers, ops
+
+from .base import Aggregation
+
+
+class LCMAggregation(Aggregation):
+    r"""The Learnable Commutative Monoid aggregation from the
+    `"Learnable Commutative Monoids for Graph Neural Networks"
+    <https://arxiv.org/abs/2212.08541>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import LCMAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = LCMAggregation(in_channels=8, out_channels=16)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        project: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if in_channels != out_channels and not project:
+            raise ValueError(
+                f"Inputs of '{self.__class__.__name__}' must be projected if `in_channels != out_channels`"
+            )
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.project = project
+
+        self.lin = layers.Dense(out_channels) if project else None
+        self.gru_cell = layers.GRUCell(out_channels)
+
+    def reset_parameters(self):
+        if self.lin is not None:
+            self.lin.reset_parameters()
+        self.gru_cell.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+        **kwargs,
+    ):
+        if self.lin is not None:
+            x = ops.relu(self.lin(x))
+
+        x_dense, _ = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=max_num_elements,
+        )
+
+        # Transpose to [num_neighbors, num_nodes, num_features]
+        x_dense = ops.transpose(x_dense, (1, 0, 2))
+        num_neighbors = ops.shape(x_dense)[0]
+        num_nodes = ops.shape(x_dense)[1]
+        num_features = ops.shape(x_dense)[2]
+
+        if not isinstance(num_neighbors, int):  # a tensor while tracing (e.g. TensorFlow's fit)
+            return self._reduce_traced(x_dense)
+        if num_neighbors == 0:
+            return ops.zeros((num_nodes, self.out_channels), dtype=x.dtype)
+
+        depth = ceil(log2(max(num_neighbors, 1)))
+        for _ in range(depth):
+            curr_len = ops.shape(x_dense)[0]
+            if curr_len <= 1:
+                break
+            half_size = ceil(curr_len / 2)
+
+            if curr_len % 2 == 1:
+                x_pair = x_dense[:-1]
+                remainder = x_dense[-1:]
+            else:
+                x_pair = x_dense
+                remainder = None
+
+            # x_pair: [2 * half, num_nodes, num_features]
+            pair_count = ops.shape(x_pair)[0] // 2
+            x_pair = ops.reshape(x_pair, (pair_count, 2, num_nodes, num_features))
+            left = x_pair[:, 0]  # [pair_count, num_nodes, num_features]
+            right = x_pair[:, 1]  # [pair_count, num_nodes, num_features]
+
+            left_flat = ops.reshape(left, (-1, num_features))
+            right_flat = ops.reshape(right, (-1, num_features))
+
+            # GRUCell: inputs=left, state=[right]
+            out1, _ = self.gru_cell(left_flat, [right_flat])
+            out2, _ = self.gru_cell(right_flat, [left_flat])
+            out = 0.5 * (out1 + out2)
+            out = ops.reshape(out, (pair_count, num_nodes, num_features))
+
+            if remainder is not None:
+                out = ops.concatenate([out, remainder], axis=0)
+
+            x_dense = out
+
+        return ops.squeeze(x_dense, axis=0)
+
+    def _combine(self, left, right):
+        num_features = ops.shape(left)[-1]
+        left_flat, right_flat = ops.reshape(left, (-1, num_features)), ops.reshape(right, (-1, num_features))
+        out1, _ = self.gru_cell(left_flat, [right_flat])
+        out2, _ = self.gru_cell(right_flat, [left_flat])
+        return ops.reshape(0.5 * (out1 + out2), ops.shape(left))
+
+    def _reduce_traced(self, x_dense):
+        """The pairwise reduction of ``call`` for a neighbor count only known at run time: position
+        ``i`` absorbs position ``i + stride`` when ``i`` is a multiple of ``2 * stride``, with the
+        stride doubling every level. This pairs the same elements as ``call``, in a fixed shape."""
+        length = ops.shape(x_dense)[0]
+        if not self.gru_cell.built:  # no weights may be created inside the loop
+            self.gru_cell.build((None, x_dense.shape[-1]))
+        position = ops.arange(length, dtype="int32")
+
+        def body(stride, h):
+            combined = self._combine(h, ops.roll(h, -stride, axis=0))
+            active = ops.logical_and(position % (2 * stride) == 0, position + stride < length)
+            return stride * 2, ops.where(active[:, None, None], combined, h)
+
+        _, h = ops.while_loop(lambda stride, h: stride < length, body, (ops.convert_to_tensor(1, "int32"), x_dense))
+        return h[0]
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels}, project={self.project})"
+

k3_node/layers/aggr/lstm.py

@@ -0,0 +1,58 @@
+from typing import Optional
+from keras import layers
+
+from .base import Aggregation
+
+
+class LSTMAggregation(Aggregation):
+    r"""Performs LSTM-style aggregation in which the elements to aggregate are
+    interpreted as a sequence, as described in the `"Inductive Representation
+    Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import LSTMAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = LSTMAggregation(in_channels=8, out_channels=16)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.lstm = layers.LSTM(out_channels, return_sequences=False)
+
+    def build(self, input_shape=None):
+        self.lstm.build((None, None, self.in_channels))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        self.lstm.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+        training: bool = False,
+        **kwargs,
+    ):
+        x_dense, _ = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=max_num_elements,
+        )
+        return self.lstm(x_dense, training=training)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
+

k3_node/layers/aggr/mlp.py

@@ -0,0 +1,75 @@
+from typing import Optional
+from keras import layers, ops
+
+from .base import Aggregation
+
+
+class MLPAggregation(Aggregation):
+    r"""Performs MLP aggregation in which the elements to aggregate are
+    flattened into a single vectorial representation, and are then processed by
+    a Multi-Layer Perceptron (MLP).
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MLPAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MLPAggregation(in_channels=8, out_channels=16, max_num_elements=5)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        max_num_elements: int,
+        mlp: Optional[any] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.max_num_elements = max_num_elements
+
+        if mlp is None:
+            self.mlp = layers.Dense(out_channels)
+        else:
+            self.mlp = mlp
+
+    def build(self, input_shape=None):
+        if hasattr(self.mlp, "build"):
+            self.mlp.build((None, self.in_channels * self.max_num_elements))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        if hasattr(self.mlp, "reset_parameters"):
+            self.mlp.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        x_dense, _ = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=self.max_num_elements,
+        )
+        B = ops.shape(x_dense)[0]
+        flattened = ops.reshape(x_dense, (B, self.max_num_elements * self.in_channels))
+        return self.mlp(flattened)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.in_channels}, {self.out_channels}, "
+            f"max_num_elements={self.max_num_elements})"
+        )
+

k3_node/layers/aggr/multi.py

@@ -0,0 +1,154 @@
+import copy
+from typing import Any, Dict, List, Optional, Union
+from keras import layers, ops
+
+from .base import Aggregation
+
+
+class MultiAggregation(Aggregation):
+    r"""Performs aggregations with one or more aggregators and combines
+    aggregated results, as described in the `"Principal Neighbourhood
+    Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718>`_ and
+    `"Adaptive Filters and Aggregator Fusion for Efficient Graph Convolutions"
+    <https://arxiv.org/abs/2104.01481>`_ papers.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MultiAggregation
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        index = np.repeat([0, 1], 5)  # aggregate nodes 0-4 into set 0 and nodes 5-9 into set 1
+
+        aggr = MultiAggregation(aggrs=["sum", "mean", "max"])
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 24)
+        ```
+    """
+
+    def __init__(
+        self,
+        aggrs: List[Union[Aggregation, str]],
+        aggrs_kwargs: Optional[List[Dict[str, Any]]] = None,
+        mode: Optional[str] = "cat",
+        mode_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if not isinstance(aggrs, (list, tuple)):
+            raise ValueError(f"'aggrs' of '{self.__class__.__name__}' should be a list or tuple.")
+
+        if len(aggrs) == 0:
+            raise ValueError(f"'aggrs' of '{self.__class__.__name__}' should not be empty.")
+
+        if aggrs_kwargs is None:
+            aggrs_kwargs = [{}] * len(aggrs)
+        elif len(aggrs) != len(aggrs_kwargs):
+            raise ValueError(
+                f"'aggrs_kwargs' with invalid length passed to '{self.__class__.__name__}' "
+                f"(got '{len(aggrs_kwargs)}', expected '{len(aggrs)}')."
+            )
+
+        from .resolver import aggregation_resolver
+        self.aggrs = [
+            aggregation_resolver(aggr, **aggr_kw)
+            for aggr, aggr_kw in zip(aggrs, aggrs_kwargs)
+        ]
+
+        self.mode = mode
+        mode_kwargs = copy.copy(mode_kwargs) or {}
+        self.in_channels = mode_kwargs.pop("in_channels", None)
+        self.out_channels = mode_kwargs.pop("out_channels", None)
+
+        if mode in ["proj", "attn"]:
+            if len(aggrs) == 1:
+                raise ValueError("Multiple aggregations are required for 'proj' or 'attn' combine mode.")
+            if self.in_channels is None or self.out_channels is None:
+                raise ValueError(f"Combine mode '{mode}' must have `in_channels` and `out_channels` specified.")
+
+            if isinstance(self.in_channels, int):
+                self.in_channels = [self.in_channels] * len(aggrs)
+
+            if mode == "proj":
+                self.lin = layers.Dense(self.out_channels, **mode_kwargs)
+            elif mode == "attn":
+                from ..dense import HeteroDictLinear
+                channels = {str(k): v for k, v in enumerate(self.in_channels)}
+                self.lin_heads = HeteroDictLinear(channels, self.out_channels)
+                num_heads = mode_kwargs.pop("num_heads", 1)
+                self.multihead_attn = layers.MultiHeadAttention(
+                    num_heads=num_heads,
+                    key_dim=max(self.out_channels // num_heads, 1),
+                    **mode_kwargs,
+                )
+
+    def reset_parameters(self):
+        for aggr in self.aggrs:
+            if hasattr(aggr, "reset_parameters"):
+                aggr.reset_parameters()
+        if hasattr(self, "lin") and hasattr(self.lin, "reset_parameters"):
+            self.lin.reset_parameters()
+        if hasattr(self, "lin_heads") and hasattr(self.lin_heads, "reset_parameters"):
+            self.lin_heads.reset_parameters()
+
+    def get_out_channels(self, in_channels: int) -> int:
+        if self.out_channels is not None:
+            return self.out_channels
+        if self.mode == "cat":
+            return in_channels * len(self.aggrs)
+        return in_channels
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        outs = [aggr(x, index=index, ptr=ptr, dim_size=dim_size, dim=dim, **kwargs) for aggr in self.aggrs]
+        return self.combine(outs)
+
+    def combine(self, inputs: List[any]):
+        if len(inputs) == 1:
+            return inputs[0]
+
+        if self.mode == "cat":
+            return ops.concatenate(inputs, axis=-1)
+
+        if hasattr(self, "lin"):
+            return self.lin(ops.concatenate(inputs, axis=-1))
+
+        if hasattr(self, "multihead_attn"):
+            x_dict = {str(k): v for k, v in enumerate(inputs)}
+            x_dict = self.lin_heads(x_dict)
+            xs = [x_dict[str(key)] for key in range(len(inputs))]
+            # xs: [num_aggrs, B, D] -> transpose to [B, num_aggrs, D]
+            x_stack = ops.transpose(ops.stack(xs, axis=0), (1, 0, 2))
+            attn_out = self.multihead_attn(x_stack, x_stack, x_stack)
+            return ops.mean(attn_out, axis=1)
+
+        stacked = ops.stack(inputs, axis=0)  # [num_aggrs, B, D]
+        if self.mode == "sum":
+            return ops.sum(stacked, axis=0)
+        elif self.mode == "mean":
+            return ops.mean(stacked, axis=0)
+        elif self.mode == "max":
+            return ops.max(stacked, axis=0)
+        elif self.mode == "min":
+            return ops.min(stacked, axis=0)
+        elif self.mode == "logsumexp":
+            return ops.logsumexp(stacked, axis=0)
+        elif self.mode == "std":
+            return ops.std(stacked, axis=0)
+        elif self.mode == "var":
+            return ops.var(stacked, axis=0)
+
+        raise ValueError(f"Combine mode '{self.mode}' is not supported.")
+
+    def __repr__(self) -> str:
+        aggrs = ",\n".join([f"  {aggr}" for aggr in self.aggrs])
+        return f"{self.__class__.__name__}([\n{aggrs}\n], mode={self.mode})"
+