k3-node 1.0.0__py3-none-any.whl

k3_node/layers/conv/gat_conv.py

@@ -0,0 +1,244 @@
+from typing import Optional, Union, Tuple
+from keras import layers, ops
+
+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.segment import segment_sum
+
+
+class GATConv(MessagePassing):
+    r"""The graph attentional operator from the `"Graph Attention Networks"
+    <https://arxiv.org/abs/1710.10903>`_ paper.
+
+    Args:
+        in_channels: Size of each input sample, or a tuple for bipartite graphs.
+        out_channels: Size of each output sample.
+        heads: Number of multi-head-attentions. (default: ``1``)
+        concat: If set to :obj:`False`, the multi-head-attentions are averaged
+            instead of concatenated. (default: ``True``)
+        negative_slope: LeakyReLU angle of the negative slope. (default: ``0.2``)
+        dropout: Dropout probability of the normalized attention coefficients.
+            (default: ``0.0``)
+        add_self_loops: If set to :obj:`False`, will not add self-loops to
+            the input graph. (default: ``True``)
+        edge_dim: Edge feature dimensionality (in case there are any).
+            (default: :obj:`None`)
+        fill_value: The way to generate edge features of self-loops
+            (default: ``"mean"``)
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+        share_weights: If set to :obj:`True`, the same matrix will be applied
+            to the source and target node features. (default: ``False``)
+        residual: If set to :obj:`True`, will compute residual connections.
+            (default: ``False``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GATConv
+
+        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 = GATConv(in_channels=8, out_channels=16, heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        dropout: float = 0.0,
+        add_self_loops: bool = True,
+        edge_dim: Optional[int] = None,
+        fill_value: Union[float, str] = "mean",
+        bias: bool = True,
+        share_weights: bool = False,
+        residual: bool = False,
+        **kwargs,
+    ):
+        super().__init__(node_dim=0, **kwargs)
+        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.edge_dim = edge_dim
+        self.fill_value = fill_value
+        self.use_bias = bias
+        self.share_weights = share_weights
+        self.residual = residual
+
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        if isinstance(in_channels, int):
+            self.lin = layers.Dense(heads * out_channels, use_bias=False)
+            self.lin_src = self.lin
+            self.lin_dst = self.lin
+        else:
+            self.lin = None
+            self.lin_src = layers.Dense(heads * out_channels, use_bias=False)
+            if share_weights:
+                self.lin_dst = self.lin_src
+            else:
+                self.lin_dst = layers.Dense(heads * out_channels, use_bias=False)
+
+        if edge_dim is not None:
+            self.lin_edge = layers.Dense(heads * out_channels, use_bias=False)
+        else:
+            self.lin_edge = None
+
+        if residual:
+            self.res = layers.Dense(total_out_channels, use_bias=False)
+        else:
+            self.res = None
+
+        self.dropout = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+            in_channels_src = input_shape[0][-1]
+            in_channels_dst = input_shape[1][-1] if len(input_shape) > 1 and input_shape[1] is not None else in_channels_src
+        else:
+            in_channels_src = input_shape[-1]
+            in_channels_dst = input_shape[-1]
+
+        self.lin_src.build((None, in_channels_src))
+        if self.lin_dst is not self.lin_src:
+            self.lin_dst.build((None, in_channels_dst))
+        if self.lin_edge is not None:
+            self.lin_edge.build((None, self.edge_dim))
+        if self.res is not None:
+            self.res.build((None, in_channels_dst))
+
+        self.att_src = self.add_weight(
+            shape=(1, self.heads, self.out_channels),
+            initializer="glorot_uniform",
+            name="att_src",
+        )
+        self.att_dst = self.add_weight(
+            shape=(1, self.heads, self.out_channels),
+            initializer="glorot_uniform",
+            name="att_dst",
+        )
+        if self.edge_dim is not None:
+            self.att_edge = self.add_weight(
+                shape=(1, self.heads, self.out_channels),
+                initializer="glorot_uniform",
+                name="att_edge",
+            )
+        else:
+            self.att_edge = None
+
+        total_out_channels = self.out_channels * (self.heads if self.concat else 1)
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(total_out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=None, size=None, return_attention_weights=None, training=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        H, C = self.heads, self.out_channels
+        if isinstance(x, (tuple, list)):
+            x_src, x_dst = x[0], x[1]
+        else:
+            x_src, x_dst = x, x
+
+        x_src_proj = ops.reshape(self.lin_src(x_src), (-1, H, C))
+        x_dst_proj = ops.reshape(self.lin_dst(x_dst), (-1, H, C)) if x_dst is not None else None
+
+        alpha_src = ops.sum(x_src_proj * self.att_src, axis=-1)
+        alpha_dst = ops.sum(x_dst_proj * self.att_dst, axis=-1) if x_dst_proj is not None else None
+
+        if self.add_self_loops:
+            if not isinstance(x, (tuple, list)):
+                num_nodes = ops.shape(x)[0]
+            else:
+                num_nodes = ops.shape(x_src)[0]
+                if x_dst is not None:
+                    num_nodes = ops.minimum(num_nodes, ops.shape(x_dst)[0])
+            edge_index, edge_attr, keep_mask = remove_self_loops_masked(edge_index, edge_attr)
+            edge_index, edge_attr = add_self_loops(
+                edge_index, edge_attr, fill_value=self.fill_value, num_nodes=num_nodes
+            )
+            keep_mask = extend_mask_for_self_loops(keep_mask, num_nodes)
+        else:
+            keep_mask = None
+
+        row, col = edge_index[0], edge_index[1]
+        row, col = ops.cast(row, "int32"), ops.cast(col, "int32")
+
+        alpha_j = ops.take(alpha_src, row, axis=0)
+        alpha_i = ops.take(alpha_dst, col, axis=0) if alpha_dst is not None else 0.0
+        alpha = alpha_j + alpha_i
+
+        if edge_attr is not None and self.lin_edge is not None and self.att_edge is not None:
+            edge_attr_proj = ops.reshape(self.lin_edge(edge_attr), (-1, H, C))
+            alpha = alpha + ops.sum(edge_attr_proj * self.att_edge, axis=-1)
+
+        alpha = ops.leaky_relu(alpha, negative_slope=self.negative_slope)
+        alpha = mask_edge_logits(alpha, keep_mask)
+        num_nodes_dst = ops.shape(x_dst)[0] if x_dst is not None else ops.shape(x_src)[0]
+        alpha = softmax(alpha, col, num_nodes=num_nodes_dst, dim=0)
+
+        if self.dropout is not None:
+            alpha = self.dropout(alpha, training=training)
+
+        # Message & aggregate
+        x_src_j = ops.take(x_src_proj, row, axis=0)
+        out = ops.expand_dims(alpha, -1) * x_src_j
+        out = segment_sum(out, col, num_segments=num_nodes_dst)
+
+        if self.concat:
+            out = ops.reshape(out, (-1, H * C))
+        else:
+            out = ops.mean(out, axis=1)
+
+        if self.res is not None and x_dst is not None:
+            out = out + self.res(x_dst)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        if return_attention_weights:
+            return out, (edge_index, alpha)
+        return out
+
+
+class FusedGATConv(GATConv):
+    r"""The fused graph attentional operator.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FusedGATConv
+
+        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 = FusedGATConv(in_channels=8, out_channels=16)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    pass

k3_node/layers/conv/gated_graph_conv.py

@@ -0,0 +1,136 @@
+# ported from spektral
+
+from keras import ops
+from keras.layers import GRUCell
+
+from k3_node.layers.conv.message_passing import MessagePassing
+
+
+class GatedGraphConv(MessagePassing):
+    """
+    `k3_node.layers.GatedGraphConv` 
+
+    Implementation of Gated Graph Convolution (GGC) layer
+
+    Args:
+        channels: The number of output channels.
+        n_layers: The number of GGC layers to stack.
+        activation: Activation function to use.
+        use_bias: Whether to add a bias to the linear transformation.
+        kernel_initializer: Initializer for the `kernel` weights matrix.
+        bias_initializer: Initializer for the bias vector.
+        kernel_regularizer: Regularizer for the `kernel` weights matrix.
+        bias_regularizer: Regularizer for the bias vector.
+        activity_regularizer: Regularizer for the output.
+        kernel_constraint: Constraint for the `kernel` weights matrix.
+        bias_constraint: Constraint for the bias vector.
+        **kwargs: Additional arguments to pass to the `MessagePassing` superclass.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GatedGraphConv
+
+        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 = GatedGraphConv(out_channels=16, num_layers=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        channels=None,
+        n_layers=None,
+        out_channels=None,
+        num_layers=None,
+        activation=None,
+        use_bias=True,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        kernel_regularizer=None,
+        bias_regularizer=None,
+        activity_regularizer=None,
+        kernel_constraint=None,
+        bias_constraint=None,
+        **kwargs,
+    ):
+        channels = out_channels if out_channels is not None else channels
+        n_layers = num_layers if num_layers is not None else n_layers
+        super().__init__(
+            activation=activation,
+            use_bias=use_bias,
+            kernel_initializer=kernel_initializer,
+            bias_initializer=bias_initializer,
+            kernel_regularizer=kernel_regularizer,
+            bias_regularizer=bias_regularizer,
+            activity_regularizer=activity_regularizer,
+            kernel_constraint=kernel_constraint,
+            bias_constraint=bias_constraint,
+            **kwargs,
+        )
+        self.channels = channels
+        self.out_channels = channels
+        self.n_layers = n_layers
+        self.num_layers = n_layers
+
+    def build(self, input_shape=None):
+        self.kernel = self.add_weight(
+            name="kernel",
+            shape=(self.n_layers, self.channels, self.channels),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+        )
+        self.rnn = GRUCell(
+            self.channels,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+            kernel_regularizer=self.kernel_regularizer,
+            bias_regularizer=self.bias_regularizer,
+            activity_regularizer=self.activity_regularizer,
+            kernel_constraint=self.kernel_constraint,
+            bias_constraint=self.bias_constraint,
+            use_bias=self.use_bias,
+            dtype=self.dtype,
+        )
+        self.rnn.build((self.channels,))
+        super().build(input_shape)
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_weight=None, **kwargs):
+        is_legacy = False
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, a, _ = self.get_inputs(x)
+            edge_index = a
+            is_legacy = True
+
+        F = ops.shape(x)[-1]
+        if F < self.channels:
+            to_pad = self.channels - F
+            ndims = len(ops.shape(x)) - 1
+            output = ops.pad(x, [[0, 0]] * ndims + [[0, to_pad]])
+        elif F > self.channels:
+            output = x[..., :self.channels]
+        else:
+            output = x
+
+        for i in range(self.n_layers):
+            m = ops.matmul(output, self.kernel[i])
+            if is_legacy:
+                m = self.propagate(m, edge_index)
+            else:
+                m = self.propagate(edge_index, x=m, edge_weight=edge_weight)
+            output = self.rnn(m, [output])[0]
+
+        if hasattr(self, "activation") and self.activation is not None and callable(self.activation):
+            output = self.activation(output)
+        return output
+
+    @property
+    def config(self):
+        return {
+            "channels": self.channels,
+            "n_layers": self.n_layers,
+        }

k3_node/layers/conv/gatv2_conv.py

@@ -0,0 +1,205 @@
+from typing import Optional, Union, Tuple
+from keras import layers, ops
+
+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.segment import segment_sum
+
+
+class GATv2Conv(MessagePassing):
+    r"""The GATv2 operator from the `"How Attentive are Graph Attention Networks?"
+    <https://arxiv.org/abs/2105.14491>`_ paper, which fixes the static
+    attention problem of standard :class:`~k3_node.layers.conv.GATConv`.
+
+    Args:
+        in_channels: Size of each input sample, or a tuple for bipartite graphs.
+        out_channels: Size of each output sample.
+        heads: Number of multi-head-attentions. (default: ``1``)
+        concat: If set to :obj:`False`, the multi-head-attentions are averaged
+            instead of concatenated. (default: ``True``)
+        negative_slope: LeakyReLU angle of the negative slope. (default: ``0.2``)
+        dropout: Dropout probability of the normalized attention coefficients.
+            (default: ``0.0``)
+        add_self_loops: If set to :obj:`False`, will not add self-loops to
+            the input graph. (default: ``True``)
+        edge_dim: Edge feature dimensionality (in case there are any).
+            (default: :obj:`None`)
+        fill_value: The way to generate edge features of self-loops
+            (default: ``"mean"``)
+        bias: If set to :obj:`False`, the layer will not learn an additive bias.
+            (default: ``True``)
+        share_weights: If set to :obj:`True`, the same matrix will be applied
+            to the source and target node features. (default: ``False``)
+        residual: If set to :obj:`True`, will compute residual connections.
+            (default: ``False``)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import GATv2Conv
+
+        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 = GATv2Conv(in_channels=8, out_channels=16, heads=2)
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 32)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: Union[int, Tuple[int, int]],
+        out_channels: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        dropout: float = 0.0,
+        add_self_loops: bool = True,
+        edge_dim: Optional[int] = None,
+        fill_value: Union[float, str] = "mean",
+        bias: bool = True,
+        share_weights: bool = False,
+        residual: bool = False,
+        **kwargs,
+    ):
+        super().__init__(node_dim=0, **kwargs)
+        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.edge_dim = edge_dim
+        self.fill_value = fill_value
+        self.use_bias = bias
+        self.share_weights = share_weights
+        self.residual = residual
+
+        total_out_channels = out_channels * (heads if concat else 1)
+
+        self.lin_l = layers.Dense(heads * out_channels, use_bias=bias)
+        if share_weights:
+            self.lin_r = self.lin_l
+        else:
+            self.lin_r = layers.Dense(heads * out_channels, use_bias=bias)
+
+        if edge_dim is not None:
+            self.lin_edge = layers.Dense(heads * out_channels, use_bias=False)
+        else:
+            self.lin_edge = None
+
+        if residual:
+            self.res = layers.Dense(total_out_channels, use_bias=False)
+        else:
+            self.res = None
+
+        self.dropout = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def build(self, input_shape):
+        if isinstance(input_shape, (tuple, list)) and len(input_shape) > 0 and isinstance(input_shape[0], (tuple, list)):
+            in_channels_src = input_shape[0][-1]
+            in_channels_dst = input_shape[1][-1] if len(input_shape) > 1 and input_shape[1] is not None else in_channels_src
+        else:
+            in_channels_src = input_shape[-1]
+            in_channels_dst = input_shape[-1]
+
+        self.lin_l.build((None, in_channels_src))
+        if self.lin_r is not self.lin_l:
+            self.lin_r.build((None, in_channels_dst))
+        if self.lin_edge is not None:
+            self.lin_edge.build((None, self.edge_dim))
+        if self.res is not None:
+            self.res.build((None, in_channels_dst))
+
+        self.att = self.add_weight(
+            shape=(1, self.heads, self.out_channels),
+            initializer="glorot_uniform",
+            name="att",
+        )
+
+        total_out_channels = self.out_channels * (self.heads if self.concat else 1)
+        if self.use_bias:
+            self.bias = self.add_weight(
+                shape=(total_out_channels,),
+                initializer="zeros",
+                name="bias",
+            )
+        else:
+            self.bias = None
+        self.built = True
+
+    def call(self, x, edge_index=None, edge_attr=None, size=None, return_attention_weights=None, training=None, **kwargs):
+        if edge_index is None and isinstance(x, (tuple, list)):
+            x, edge_index = x[0], x[1]
+
+        H, C = self.heads, self.out_channels
+        if isinstance(x, (tuple, list)):
+            x_src, x_dst = x[0], x[1]
+        else:
+            x_src, x_dst = x, x
+
+        x_l = ops.reshape(self.lin_l(x_src), (-1, H, C))
+        x_r = ops.reshape(self.lin_r(x_dst), (-1, H, C)) if x_dst is not None else x_l
+
+        if self.add_self_loops:
+            if not isinstance(x, (tuple, list)):
+                num_nodes = ops.shape(x)[0]
+            else:
+                num_nodes = ops.shape(x_l)[0]
+                if x_r is not None:
+                    num_nodes = ops.minimum(num_nodes, ops.shape(x_r)[0])
+            edge_index, edge_attr, keep_mask = remove_self_loops_masked(edge_index, edge_attr)
+            edge_index, edge_attr = add_self_loops(
+                edge_index, edge_attr, fill_value=self.fill_value, num_nodes=num_nodes
+            )
+            keep_mask = extend_mask_for_self_loops(keep_mask, num_nodes)
+        else:
+            keep_mask = None
+
+        row, col = edge_index[0], edge_index[1]
+        row, col = ops.cast(row, "int32"), ops.cast(col, "int32")
+
+        x_l_j = ops.take(x_l, row, axis=0)
+        x_r_i = ops.take(x_r, col, axis=0)
+        alpha = x_l_j + x_r_i
+
+        if edge_attr is not None and self.lin_edge is not None:
+            edge_attr_proj = ops.reshape(self.lin_edge(edge_attr), (-1, H, C))
+            alpha = alpha + edge_attr_proj
+
+        alpha = ops.leaky_relu(alpha, negative_slope=self.negative_slope)
+        alpha = ops.sum(alpha * self.att, axis=-1)
+        alpha = mask_edge_logits(alpha, keep_mask)
+
+        num_nodes_dst = ops.shape(x_r)[0]
+        alpha = softmax(alpha, col, num_nodes=num_nodes_dst, dim=0)
+
+        if self.dropout is not None:
+            alpha = self.dropout(alpha, training=training)
+
+        out = ops.expand_dims(alpha, -1) * x_l_j
+        out = segment_sum(out, col, num_segments=num_nodes_dst)
+
+        if self.concat:
+            out = ops.reshape(out, (-1, H * C))
+        else:
+            out = ops.mean(out, axis=1)
+
+        if self.res is not None and x_dst is not None:
+            out = out + self.res(x_dst)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        if return_attention_weights:
+            return out, (edge_index, alpha)
+        return out
+