k3-node 1.0.0__py3-none-any.whl

k3_node/layers/attention/polynormer.py

@@ -0,0 +1,160 @@
+import keras
+from keras import ops
+
+
+class PolynormerAttention(keras.layers.Layer):
+    """Layer ``PolynormerAttention``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PolynormerAttention
+
+        x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]
+
+        mask = np.ones((1, 10), dtype=bool)
+        attn = PolynormerAttention(channels=8, heads=2)
+        print(tuple(attn(x, mask).shape))  # (1, 10, 128)
+        ```
+    """
+    def __init__(
+        self,
+        channels,
+        heads,
+        head_channels=64,
+        beta=0.9,
+        qkv_bias=False,
+        qk_shared=True,
+        dropout=0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+
+        self.channels = channels
+        self.heads = heads
+        self.head_channels = head_channels
+        self.beta = beta
+        self.qk_shared = qk_shared
+
+        inner_channels = heads * head_channels
+
+        self.h_lins = keras.layers.Dense(inner_channels)
+
+        if not qk_shared:
+            self.q = keras.layers.Dense(
+                inner_channels,
+                use_bias=qkv_bias,
+            )
+
+        self.k = keras.layers.Dense(
+            inner_channels,
+            use_bias=qkv_bias,
+        )
+
+        self.v = keras.layers.Dense(
+            inner_channels,
+            use_bias=qkv_bias,
+        )
+
+        self.lns = keras.layers.LayerNormalization(epsilon=1e-5)
+        self.lin_out = keras.layers.Dense(inner_channels)
+        self.dropout = keras.layers.Dropout(dropout)
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def call(self, x, mask=None, training=None):
+        shape = ops.shape(x)
+
+        B = shape[0]
+        N = shape[1]
+
+        h = self.h_lins(x)
+
+        k = self.k(x)
+        k = ops.sigmoid(k)
+        k = ops.reshape(
+            k,
+            (B, N, self.head_channels, self.heads),
+        )
+
+        if self.qk_shared:
+            q = k
+        else:
+            q = ops.sigmoid(self.q(x))
+            q = ops.reshape(
+                q,
+                (B, N, self.head_channels, self.heads),
+            )
+
+        v = self.v(x)
+        v = ops.reshape(
+            v,
+            (B, N, self.head_channels, self.heads),
+        )
+
+        if mask is not None:
+            mask = ops.expand_dims(mask, axis=-1)
+            mask = ops.expand_dims(mask, axis=-1)
+            v = v * ops.cast(mask, v.dtype)
+
+        # numerator
+        kv = ops.einsum(
+            "bndh,bnmh->bdmh",
+            k,
+            v,
+        )
+
+        num = ops.einsum(
+            "bndh,bdmh->bnmh",
+            q,
+            kv,
+        )
+
+        # denominator
+        k_sum = ops.einsum(
+            "bndh->bdh",
+            k,
+        )
+
+        den = ops.einsum(
+            "bndh,bdh->bnh",
+            q,
+            k_sum,
+        )
+
+        den = ops.expand_dims(den, axis=2)
+
+        x = num / (den + 1e-6)
+
+        x = ops.reshape(
+            x,
+            (B, N, self.head_channels * self.heads),
+        )
+
+        x = self.lns(x)
+        x = x * (h + self.beta)
+
+        x = self.lin_out(x)
+        x = ops.relu(x)
+
+        x = self.dropout(
+            x,
+            training=training,
+        )
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "channels": self.channels,
+                "heads": self.heads,
+                "head_channels": self.head_channels,
+                "beta": self.beta,
+                "qk_shared": self.qk_shared,
+            }
+        )
+        return config

k3_node/layers/attention/qformer.py

@@ -0,0 +1,143 @@
+from keras import layers, ops
+
+
+class QFormerEncoderLayer(layers.Layer):
+    """Layer ``QFormerEncoderLayer``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import QFormerEncoderLayer
+
+        x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]
+
+        layer = QFormerEncoderLayer(input_dim=8, hidden_dim=16, num_heads=2)
+        print(tuple(layer(x).shape))  # (1, 10, 8)
+        ```
+    """
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        num_heads,
+        dropout=0.0,
+        activation="relu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.self_attn = layers.MultiHeadAttention(
+            num_heads=num_heads,
+            key_dim=input_dim // num_heads,
+            dropout=dropout,
+            use_bias=True,
+        )
+
+        self.linear1 = layers.Dense(hidden_dim)
+        self.linear2 = layers.Dense(input_dim)
+
+        self.norm1 = layers.LayerNormalization(epsilon=1e-5)
+        self.norm2 = layers.LayerNormalization(epsilon=1e-5)
+
+        self.dropout = layers.Dropout(dropout)
+        self.dropout1 = layers.Dropout(dropout)
+        self.dropout2 = layers.Dropout(dropout)
+
+        self.activation = activation
+
+    def call(self, x, training=None):
+        # PyTorch TransformerEncoderLayer, norm_first=False
+
+        attn = self.self_attn(
+            x,
+            x,
+            training=training,
+        )
+
+        x = self.norm1(
+            x + self.dropout1(attn, training=training)
+        )
+
+        ff = self.linear1(x)
+
+        if self.activation == "relu":
+            ff = ops.relu(ff)
+        elif self.activation == "gelu":
+            ff = ops.gelu(ff)
+        else:
+            raise ValueError(
+                f"Unsupported activation: {self.activation}"
+            )
+
+        ff = self.dropout(ff, training=training)
+        ff = self.linear2(ff)
+
+        x = self.norm2(
+            x + self.dropout2(ff, training=training)
+        )
+
+        return x
+
+
+class QFormer(layers.Layer):
+    """Layer ``QFormer``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import QFormer
+
+        x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]
+
+        model = QFormer(input_dim=8, hidden_dim=16, output_dim=4, num_heads=2, num_layers=1)
+        print(tuple(model(x).shape))  # (1, 10, 4)
+        ```
+    """
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        output_dim,
+        num_heads,
+        num_layers,
+        dropout=0.0,
+        activation="relu",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+
+        self.layer_norm = layers.LayerNormalization(
+            epsilon=1e-5
+        )
+
+        self.encoder = [
+            QFormerEncoderLayer(
+                input_dim=input_dim,
+                hidden_dim=hidden_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                activation=activation,
+            )
+            for _ in range(num_layers)
+        ]
+
+        self.project = layers.Dense(output_dim)
+
+    def call(self, x, training=None):
+        x = self.layer_norm(x)
+
+        for layer in self.encoder:
+            x = layer(x, training=training)
+
+        return self.project(x)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            "num_layers": self.num_layers,
+            "num_heads": self.num_heads,
+        })
+        return config

k3_node/layers/attention/sgformer.py

@@ -0,0 +1,106 @@
+from typing import Optional
+import keras
+from keras import ops
+
+
+class SGFormerAttention(keras.layers.Layer):
+    r"""The simple global attention mechanism from the
+    `"SGFormer: Simplifying and Empowering Transformers for
+    Large-Graph Representations"
+    <https://arxiv.org/abs/2306.10759>`_ paper.
+
+    Args:
+        channels (int): Size of each input sample.
+        heads (int, optional): Number of parallel attention heads.
+            (default: :obj:`1`)
+        head_channels (int, optional): Size of each attention head.
+            (default: :obj:`64`)
+        qkv_bias (bool, optional): If specified, add bias to query, key
+            and value in the self attention. (default: :obj:`False`)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SGFormerAttention
+
+        x = np.random.rand(1, 10, 8).astype("float32")  # [batch, num_nodes, channels]
+
+        attn = SGFormerAttention(channels=8, heads=2)
+        print(tuple(attn(x).shape))  # (1, 10, 64)
+        ```
+    """
+    def __init__(
+        self,
+        channels: int,
+        heads: int = 1,
+        head_channels: Optional[int] = 64,
+        qkv_bias: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+        assert channels % heads == 0
+        if head_channels is None:
+            head_channels = channels // heads
+
+        self.channels = channels
+        self.heads = heads
+        self.head_channels = head_channels
+        self.qkv_bias = qkv_bias
+
+        inner_channels = head_channels * heads
+        self.q = keras.layers.Dense(inner_channels, use_bias=qkv_bias)
+        self.k = keras.layers.Dense(inner_channels, use_bias=qkv_bias)
+        self.v = keras.layers.Dense(inner_channels, use_bias=qkv_bias)
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def call(self, x, mask: Optional[any] = None):
+        shape = ops.shape(x)
+        B = shape[0]
+        N = shape[1]
+
+        qs = self.q(x)
+        ks = self.k(x)
+        vs = self.v(x)
+
+        qs = ops.reshape(qs, (B, N, self.heads, self.head_channels))
+        ks = ops.reshape(ks, (B, N, self.heads, self.head_channels))
+        vs = ops.reshape(vs, (B, N, self.heads, self.head_channels))
+
+        if mask is not None:
+            m = ops.expand_dims(ops.expand_dims(mask, -1), -1)
+            m = ops.cast(m, "bool")
+            vs = ops.where(m, vs, ops.zeros_like(vs))
+
+        epsilon = 1e-6
+        qs = ops.where(ops.equal(qs, 0.0), epsilon, qs)
+        ks = ops.where(ops.equal(ks, 0.0), epsilon, ks)
+
+        qs_norm = ops.norm(qs, axis=-1, keepdims=True)
+        ks_norm = ops.norm(ks, axis=-1, keepdims=True)
+        qs = qs / ops.maximum(qs_norm, 1e-12)
+        ks = ks / ops.maximum(ks_norm, 1e-12)
+
+        # numerator
+        kvs = ops.einsum("blhm,blhd->bhmd", ks, vs)
+        attention_num = ops.einsum("bnhm,bhmd->bnhd", qs, kvs)
+        n_float = ops.cast(N, vs.dtype)
+        attention_num = attention_num + n_float * vs
+
+        # denominator
+        all_ones = ops.ones((B, N), dtype=ks.dtype)
+        ks_sum = ops.einsum("blhm,bl->bhm", ks, all_ones)
+        attention_normalizer = ops.einsum("bnhm,bhm->bnh", qs, ks_sum)
+        attention_normalizer = ops.expand_dims(attention_normalizer, -1)
+        attention_normalizer = attention_normalizer + n_float
+        attn_output = attention_num / attention_normalizer
+
+        return ops.mean(attn_output, axis=2)
+
+    def __repr__(self) -> str:
+        return (f"{self.__class__.__name__}("
+                f"heads={self.heads}, "
+                f"head_channels={self.head_channels})")
+

k3_node/layers/attention/test_attention.py

@@ -0,0 +1,68 @@
+import pytest
+from keras import ops, random
+
+from k3_node.layers import PerformerAttention, PolynormerAttention, QFormer
+
+
+@pytest.mark.parametrize("in_channels", [32, 64])
+@pytest.mark.parametrize("out_channels", [32, 64])
+@pytest.mark.parametrize("heads", [4, 8])
+@pytest.mark.parametrize("num_nodes", [4, 8])
+def test_performer_attention(num_nodes, in_channels, out_channels, heads):
+    x = random.normal((1, num_nodes, in_channels))
+    mask = ops.ones([1, num_nodes])
+    attn = PerformerAttention(channels=out_channels, heads=heads)
+    out = attn(x, mask)
+    assert out.shape == (1, num_nodes, out_channels)
+
+
+
+@pytest.mark.parametrize("in_channels", [32, 64])
+@pytest.mark.parametrize("heads", [4, 8])
+@pytest.mark.parametrize("num_nodes", [4, 8])
+def test_polynormer_attention(num_nodes, in_channels, heads):
+
+    x = random.normal((1, num_nodes, in_channels))
+
+    mask = ops.ones([1, num_nodes])
+
+    attn = PolynormerAttention(
+        channels=in_channels,
+        heads=heads,
+    )
+
+    out = attn(x, mask)
+
+    assert out.shape == (1, num_nodes, heads * 64)
+
+@pytest.mark.parametrize("input_dim", [16, 32])
+@pytest.mark.parametrize("hidden_dim", [16, 32])
+@pytest.mark.parametrize("output_dim", [16, 32])
+@pytest.mark.parametrize("num_heads", [2, 4])
+@pytest.mark.parametrize("num_layers", [1, 2])
+@pytest.mark.parametrize("num_nodes", [4, 8])
+def test_qformer(
+    input_dim,
+    hidden_dim,
+    output_dim,
+    num_heads,
+    num_layers,
+    num_nodes,
+):
+    x = random.normal((1, num_nodes, input_dim))
+
+    attn = QFormer(
+        input_dim=input_dim,
+        hidden_dim=hidden_dim,
+        output_dim=output_dim,
+        num_heads=num_heads,
+        num_layers=num_layers,
+    )
+
+    out = attn(x)
+
+    assert out.shape == (
+        1,
+        num_nodes,
+        output_dim,
+    )

k3_node/layers/attention/test_pair_attention.py

@@ -0,0 +1,91 @@
+import numpy as np
+import pytest
+import keras
+from keras import ops
+
+from k3_node.layers.attention import (
+    SelfMultiheadAttentionWithPair,
+    TransformerEncoderLayerWithPair,
+    TriangleMultiplication,
+    OuterProduct,
+)
+
+
+def test_self_multihead_attention_with_pair():
+    bsz = 2
+    seq_len = 6
+    embed_dim = 32
+    num_heads = 4
+
+    x = ops.convert_to_tensor(np.random.randn(bsz, seq_len, embed_dim).astype("float32"))
+    attn_bias = ops.convert_to_tensor(np.random.randn(bsz, num_heads, seq_len, seq_len).astype("float32"))
+    padding_mask = ops.convert_to_tensor(np.zeros((bsz, seq_len), dtype="bool"))
+
+    layer = SelfMultiheadAttentionWithPair(embed_dim=embed_dim, num_heads=num_heads)
+    out = layer(x, key_padding_mask=padding_mask, attn_bias=attn_bias)
+    assert ops.shape(out) == (bsz, seq_len, embed_dim)
+
+    # Test return_attn=True
+    out, weights, probs = layer(x, key_padding_mask=padding_mask, attn_bias=attn_bias, return_attn=True)
+    assert ops.shape(out) == (bsz, seq_len, embed_dim)
+    assert ops.shape(weights) == (bsz, num_heads, seq_len, seq_len)
+    assert ops.shape(probs) == (bsz, num_heads, seq_len, seq_len)
+
+
+def test_transformer_encoder_layer_with_pair():
+    bsz = 2
+    seq_len = 6
+    embed_dim = 32
+    num_heads = 4
+
+    x = ops.convert_to_tensor(np.random.randn(bsz, seq_len, embed_dim).astype("float32"))
+    attn_bias = ops.convert_to_tensor(np.random.randn(bsz, num_heads, seq_len, seq_len).astype("float32"))
+
+    layer = TransformerEncoderLayerWithPair(
+        embed_dim=embed_dim,
+        ffn_embed_dim=64,
+        attention_heads=num_heads,
+        post_ln=False,
+    )
+    out = layer(x, attn_bias=attn_bias)
+    assert ops.shape(out) == (bsz, seq_len, embed_dim)
+
+    # Post-LN
+    layer_post = TransformerEncoderLayerWithPair(
+        embed_dim=embed_dim,
+        ffn_embed_dim=64,
+        attention_heads=num_heads,
+        post_ln=True,
+    )
+    out_post = layer_post(x, attn_bias=attn_bias)
+    assert ops.shape(out_post) == (bsz, seq_len, embed_dim)
+
+
+def test_triangle_multiplication():
+    bsz = 2
+    seq_len = 5
+    pair_dim = 16
+    hidden_dim = 8
+
+    pair = ops.convert_to_tensor(np.random.randn(bsz, seq_len, seq_len, pair_dim).astype("float32"))
+
+    tri_out = TriangleMultiplication(pair_dim=pair_dim, hidden_dim=hidden_dim, mode="outgoing")
+    out = tri_out(pair)
+    assert ops.shape(out) == (bsz, seq_len, seq_len, pair_dim)
+
+    tri_in = TriangleMultiplication(pair_dim=pair_dim, hidden_dim=hidden_dim, mode="incoming")
+    out = tri_in(pair)
+    assert ops.shape(out) == (bsz, seq_len, seq_len, pair_dim)
+
+
+def test_outer_product():
+    bsz = 2
+    seq_len = 6
+    embed_dim = 32
+    pair_dim = 16
+
+    x = ops.convert_to_tensor(np.random.randn(bsz, seq_len, embed_dim).astype("float32"))
+    op = OuterProduct(embed_dim=embed_dim, pair_dim=pair_dim, hidden_dim=8)
+    out = op(x)
+    assert ops.shape(out) == (bsz, seq_len, seq_len, pair_dim)
+