k3-node 1.0.0__py3-none-any.whl

k3_node/layers/aggr/sort.py

@@ -0,0 +1,68 @@
+from typing import Optional
+from keras import ops
+import numpy as np
+
+from .base import Aggregation
+from k3_node.ops.creation import repeat
+
+
+class SortAggregation(Aggregation):
+    r"""The pooling operator from the `"An End-to-End Deep Learning
+    Architecture for Graph Classification"
+    <https://www.cse.wustl.edu/~muhan/papers/AAAI_2018_DGCNN.pdf>`_ paper,
+    where node features are sorted in descending order based on their last
+    feature channel. The first :math:`k` nodes form the output of the layer.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import SortAggregation
+
+        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 = SortAggregation(k=2)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(self, k: int, **kwargs):
+        super().__init__(**kwargs)
+        self.k = k
+
+    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,
+    ):
+        fill_value = -1e9
+        batch_x, mask = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            fill_value=fill_value, max_num_elements=max_num_elements,
+        )
+
+        B = ops.shape(batch_x)[0]
+        N = ops.shape(batch_x)[1]
+        D = ops.shape(batch_x)[2]
+
+        scores = batch_x[:, :, -1]  # [B, N]
+        k = self.k
+        k_val = min(k, int(scores.shape[1])) if hasattr(scores, "shape") and isinstance(scores.shape[1], int) else k
+        _, perm = ops.top_k(scores, k=k_val, sorted=True)  # [B, k]
+        perm_expanded = repeat(ops.expand_dims(perm, -1), D, axis=-1)
+        out_x = ops.take_along_axis(batch_x, perm_expanded, axis=1)
+        out_x = ops.where(ops.equal(out_x, fill_value), ops.zeros_like(out_x), out_x)
+        static_d = x.shape[-1]
+        out_x = ops.reshape(out_x, (-1, k_val * static_d) if isinstance(static_d, int) else (B, -1))
+
+        return out_x
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(k={self.k})"
+

k3_node/layers/aggr/test_aggr.py

@@ -0,0 +1,337 @@
+import numpy as np
+import pytest
+from keras import layers, ops
+
+from k3_node.layers.aggr import (
+    Aggregation,
+    AttentionalAggregation,
+    DeepSetsAggregation,
+    DegreeScalerAggregation,
+    EquilibriumAggregation,
+    FusedAggregation,
+    GRUAggregation,
+    GraphMultisetTransformer,
+    LCMAggregation,
+    LSTMAggregation,
+    MLPAggregation,
+    MaxAggregation,
+    MeanAggregation,
+    MedianAggregation,
+    MinAggregation,
+    MulAggregation,
+    MultiAggregation,
+    PatchTransformerAggregation,
+    PowerMeanAggregation,
+    QuantileAggregation,
+    Set2Set,
+    SetTransformerAggregation,
+    SoftmaxAggregation,
+    SortAggregation,
+    StdAggregation,
+    SumAggregation,
+    VarAggregation,
+    VariancePreservingAggregation,
+    aggregation_resolver,
+    ptr2index,
+    to_dense_batch,
+)
+
+
+def test_ptr2index_and_to_dense_batch():
+    ptr = ops.convert_to_tensor([0, 2, 5, 6], dtype="int32")
+    index = ptr2index(ptr)
+    assert np.array_equal(ops.convert_to_numpy(index), [0, 0, 1, 1, 1, 2])
+
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+
+    dense_x, mask = to_dense_batch(x, index)
+    assert ops.shape(dense_x) == (3, 3, 2)
+    assert ops.shape(mask) == (3, 3)
+    assert np.array_equal(ops.convert_to_numpy(mask), [
+        [True, True, False],
+        [True, True, True],
+        [True, False, False],
+    ])
+
+
+@pytest.mark.parametrize(
+    "AggrClass,expected",
+    [
+        (SumAggregation, [[4.0, 6.0], [21.0, 24.0], [11.0, 12.0]]),
+        (MeanAggregation, [[2.0, 3.0], [7.0, 8.0], [11.0, 12.0]]),
+        (MaxAggregation, [[3.0, 4.0], [9.0, 10.0], [11.0, 12.0]]),
+        (MinAggregation, [[1.0, 2.0], [5.0, 6.0], [11.0, 12.0]]),
+        (MulAggregation, [[3.0, 8.0], [315.0, 480.0], [11.0, 12.0]]),
+    ],
+)
+def test_basic_reductions(AggrClass, expected):
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    aggr = AggrClass()
+    out = aggr(x, index)
+    assert np.allclose(ops.convert_to_numpy(out), expected)
+
+    # Test with ptr
+    ptr = ops.convert_to_tensor([0, 2, 5, 6], dtype="int32")
+    out_ptr = aggr(x, ptr=ptr)
+    assert np.allclose(ops.convert_to_numpy(out_ptr), expected)
+
+
+def test_var_and_std_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    var_aggr = VarAggregation()
+    var_out = var_aggr(x, index)
+    assert ops.shape(var_out) == (3, 2)
+    assert np.all(ops.convert_to_numpy(var_out) >= 0.0)
+
+    std_aggr = StdAggregation()
+    std_out = std_aggr(x, index)
+    assert ops.shape(std_out) == (3, 2)
+    assert np.all(ops.convert_to_numpy(std_out) >= 0.0)
+
+
+def test_softmax_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    # Fixed temperature
+    aggr = SoftmaxAggregation(t=1.0)
+    out = aggr(x, index)
+    assert ops.shape(out) == (3, 2)
+
+    # Learnable multi-channel temperature
+    learnable_aggr = SoftmaxAggregation(t=1.0, learn=True, channels=2)
+    out_learn = learnable_aggr(x, index)
+    assert ops.shape(out_learn) == (3, 2)
+
+
+def test_powermean_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    # Fixed p
+    aggr = PowerMeanAggregation(p=2.0)
+    out = aggr(x, index)
+    assert ops.shape(out) == (3, 2)
+
+    # Learnable p
+    learnable_aggr = PowerMeanAggregation(p=1.0, learn=True, channels=2)
+    out_learn = learnable_aggr(x, index)
+    assert ops.shape(out_learn) == (3, 2)
+
+
+def test_quantile_and_median_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    # Median
+    med_aggr = MedianAggregation()
+    med_out = med_aggr(x, index)
+    assert ops.shape(med_out) == (3, 2)
+
+    # Quantile with multiple q
+    q_aggr = QuantileAggregation(q=[0.25, 0.75], interpolation="linear")
+    q_out = q_aggr(x, index)
+    assert ops.shape(q_out) == (3, 4)
+
+    # Interpolations
+    for interp in ["lower", "higher", "nearest", "midpoint"]:
+        q_interp = QuantileAggregation(q=0.5, interpolation=interp)
+        res = q_interp(x, index)
+        assert ops.shape(res) == (3, 2)
+
+
+def test_attentional_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 12.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    gate_nn = layers.Dense(1)
+    nn = layers.Dense(8)
+    aggr = AttentionalAggregation(gate_nn=gate_nn, nn=nn)
+
+    out = aggr(x, index)
+    assert ops.shape(out) == (3, 8)
+
+
+def test_set2set_aggregation():
+    in_channels = 4
+    x = ops.ones((6, in_channels), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    s2s = Set2Set(in_channels=in_channels, processing_steps=2)
+    out = s2s(x, index)
+    assert ops.shape(out) == (3, 2 * in_channels)
+
+
+def test_degree_scaler_aggregation():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+    deg = ops.convert_to_tensor([0.0, 1.0, 2.0, 3.0], dtype="float32")
+
+    scalers = ["identity", "amplification", "attenuation", "linear", "inverse_linear"]
+    aggr = DegreeScalerAggregation(aggr="mean", scaler=scalers, deg=deg)
+
+    out = aggr(x, index)
+    assert ops.shape(out) == (3, 4 * len(scalers))
+
+
+def test_sort_aggregation():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 1.0],
+        [5.0, 9.0],
+        [7.0, 8.0],
+        [9.0, 10.0],
+        [11.0, 5.0],
+    ], dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    aggr = SortAggregation(k=2)
+    out = aggr(x, index)
+    assert ops.shape(out) == (3, 2 * 2)
+
+
+def test_multi_aggregation():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    # Cat mode
+    aggr_cat = MultiAggregation(["mean", "max"], mode="cat")
+    assert ops.shape(aggr_cat(x, index)) == (3, 8)
+
+    # Proj mode
+    aggr_proj = MultiAggregation(["mean", "max"], mode="proj", mode_kwargs={"in_channels": 4, "out_channels": 16})
+    assert ops.shape(aggr_proj(x, index)) == (3, 16)
+
+    # Elementwise modes
+    for m in ["sum", "mean", "max", "min", "var", "std"]:
+        aggr_m = MultiAggregation(["sum", "mean"], mode=m)
+        assert ops.shape(aggr_m(x, index)) == (3, 4)
+
+
+def test_deep_sets_and_mlp_aggregation():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    # DeepSets
+    ds = DeepSetsAggregation(local_nn=layers.Dense(8), global_nn=layers.Dense(16))
+    assert ops.shape(ds(x, index)) == (3, 16)
+
+    # MLP
+    mlp_aggr = MLPAggregation(in_channels=4, out_channels=16, max_num_elements=3)
+    assert ops.shape(mlp_aggr(x, index)) == (3, 16)
+
+
+def test_lstm_and_gru_aggregation():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    lstm_aggr = LSTMAggregation(in_channels=4, out_channels=8)
+    assert ops.shape(lstm_aggr(x, index)) == (3, 8)
+
+    gru_aggr = GRUAggregation(in_channels=4, out_channels=8)
+    assert ops.shape(gru_aggr(x, index)) == (3, 8)
+
+
+def test_set_transformer_and_gmt():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    st = SetTransformerAggregation(channels=4, num_seed_points=2, heads=1, concat=True)
+    out_st = st(x, index)
+    assert ops.shape(out_st) == (3, 8)
+
+    gmt = GraphMultisetTransformer(channels=4, k=2, heads=1)
+    out_gmt = gmt(x, index)
+    assert ops.shape(out_gmt) == (3, 4)
+
+
+def test_variance_preserving_and_patch_transformer():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    vpa = VariancePreservingAggregation()
+    assert ops.shape(vpa(x, index)) == (3, 4)
+
+    pt = PatchTransformerAggregation(in_channels=4, out_channels=8, patch_size=2, hidden_channels=4)
+    assert ops.shape(pt(x, index)) == (3, 8)
+
+
+def test_lcm_and_equilibrium_aggregation():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    lcm = LCMAggregation(in_channels=4, out_channels=4, project=False)
+    assert ops.shape(lcm(x, index)) == (3, 4)
+
+    eq = EquilibriumAggregation(in_channels=4, out_channels=4, num_layers=[8])
+    assert ops.shape(eq(x, index)) == (3, 4)
+
+
+def test_fused_and_resolver():
+    x = ops.ones((6, 4), dtype="float32")
+    index = ops.convert_to_tensor([0, 0, 1, 1, 1, 2], dtype="int32")
+
+    fused = FusedAggregation(["sum", "mean", "max"])
+    outs = fused(x, index)
+    assert len(outs) == 3
+    for o in outs:
+        assert ops.shape(o) == (3, 4)
+
+    # Resolver
+    res_sum = aggregation_resolver("sum")
+    assert isinstance(res_sum, SumAggregation)
+    res_mean = aggregation_resolver("mean")
+    assert isinstance(res_mean, MeanAggregation)

k3_node/layers/aggr/utils.py

@@ -0,0 +1,210 @@
+from typing import Optional
+from keras import initializers, layers, ops
+
+
+class MultiheadAttentionBlock(layers.Layer):
+    r"""The Multihead Attention Block (MAB) from the `"Set Transformer: A
+    Framework for Attention-based Permutation-Invariant Neural Networks"
+    <https://arxiv.org/abs/1810.00825>`_ paper.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        heads: int = 1,
+        layer_norm: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+        self.channels = channels
+        self.heads = heads
+        self.dropout = dropout
+
+        key_dim = max(channels // heads, 1)
+        self.attn = layers.MultiHeadAttention(
+            num_heads=heads,
+            key_dim=key_dim,
+            dropout=dropout,
+        )
+        self.lin = layers.Dense(channels)
+        self.layer_norm1 = layers.LayerNormalization() if layer_norm else None
+        self.layer_norm2 = layers.LayerNormalization() if layer_norm else None
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def reset_parameters(self):
+        pass
+
+    def call(
+        self,
+        x,
+        y,
+        x_mask: Optional[any] = None,
+        y_mask: Optional[any] = None,
+        training: bool = False,
+    ):
+        attn_mask = None
+        if y_mask is not None:
+            # y_mask shape: [B, S], expand to [B, 1, S]
+            attn_mask = ops.expand_dims(ops.cast(y_mask, "bool"), axis=1)
+
+        out = self.attn(
+            query=x,
+            value=y,
+            key=y,
+            attention_mask=attn_mask,
+            training=training,
+        )
+
+        if x_mask is not None:
+            out = out * ops.cast(ops.expand_dims(x_mask, axis=-1), out.dtype)
+
+        out = out + x
+        if self.layer_norm1 is not None:
+            out = self.layer_norm1(out)
+
+        out = out + ops.relu(self.lin(out))
+        if self.layer_norm2 is not None:
+            out = self.layer_norm2(out)
+
+        return out
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.channels}, "
+            f"heads={self.heads}, layer_norm={self.layer_norm1 is not None}, "
+            f"dropout={self.dropout})"
+        )
+
+
+class SetAttentionBlock(layers.Layer):
+    r"""The Set Attention Block (SAB) from the `"Set Transformer"` paper."""
+
+    def __init__(
+        self,
+        channels: int,
+        heads: int = 1,
+        layer_norm: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+        self.mab = MultiheadAttentionBlock(channels, heads, layer_norm, dropout)
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def reset_parameters(self):
+        self.mab.reset_parameters()
+
+    def call(self, x, mask: Optional[any] = None, training: bool = False):
+        return self.mab(x, x, x_mask=mask, y_mask=mask, training=training)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.mab.channels}, "
+            f"heads={self.mab.heads}, layer_norm={self.mab.layer_norm1 is not None}, "
+            f"dropout={self.mab.dropout})"
+        )
+
+
+class InducedSetAttentionBlock(layers.Layer):
+    r"""The Induced Set Attention Block (ISAB) from the `"Set Transformer"` paper."""
+
+    def __init__(
+        self,
+        channels: int,
+        num_induced_points: int,
+        heads: int = 1,
+        layer_norm: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+        self.channels = channels
+        self.num_induced_points = num_induced_points
+
+        self.ind = self.add_weight(
+            shape=(1, num_induced_points, channels),
+            initializer=initializers.GlorotUniform(),
+            trainable=True,
+            name="ind",
+        )
+        self.mab1 = MultiheadAttentionBlock(channels, heads, layer_norm, dropout)
+        self.mab2 = MultiheadAttentionBlock(channels, heads, layer_norm, dropout)
+
+    def compute_mask(self, inputs, mask=None):
+        return mask
+
+    def reset_parameters(self):
+        init = initializers.GlorotUniform()
+        self.ind.assign(init(self.ind.shape, dtype=self.ind.dtype))
+        self.mab1.reset_parameters()
+        self.mab2.reset_parameters()
+
+    def call(self, x, mask: Optional[any] = None, training: bool = False):
+        B = ops.shape(x)[0]
+        ind_expanded = ops.broadcast_to(self.ind, (B, self.num_induced_points, self.channels))
+        h = self.mab1(ind_expanded, x, y_mask=mask, training=training)
+        return self.mab2(x, h, x_mask=mask, training=training)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.channels}, "
+            f"num_induced_points={self.num_induced_points}, "
+            f"heads={self.mab1.heads})"
+        )
+
+
+class PoolingByMultiheadAttention(layers.Layer):
+    r"""The Pooling by Multihead Attention (PMA) layer from the `"Set Transformer"` paper."""
+
+    def __init__(
+        self,
+        channels: int,
+        num_seed_points: int = 1,
+        heads: int = 1,
+        layer_norm: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.supports_masking = True
+        self.channels = channels
+        self.num_seed_points = num_seed_points
+        self.lin = layers.Dense(channels)
+        self.seed = self.add_weight(
+            shape=(1, num_seed_points, channels),
+            initializer=initializers.GlorotUniform(),
+            trainable=True,
+            name="seed",
+        )
+        self.mab = MultiheadAttentionBlock(channels, heads, layer_norm, dropout)
+
+    def compute_mask(self, inputs, mask=None):
+        return None
+
+    def reset_parameters(self):
+        self.lin.reset_parameters()
+        init = initializers.GlorotUniform()
+        self.seed.assign(init(self.seed.shape, dtype=self.seed.dtype))
+        self.mab.reset_parameters()
+
+    def call(self, x, mask: Optional[any] = None, training: bool = False):
+        B = ops.shape(x)[0]
+        x_proj = ops.relu(self.lin(x))
+        seed_expanded = ops.broadcast_to(self.seed, (B, self.num_seed_points, self.channels))
+        return self.mab(seed_expanded, x_proj, y_mask=mask, training=training)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.channels}, "
+            f"num_seed_points={self.num_seed_points}, "
+            f"heads={self.mab.heads})"
+        )
+

k3_node/layers/aggr/variance_preserving.py

@@ -0,0 +1,54 @@
+from typing import Optional
+from keras import ops
+
+from .base import Aggregation
+from k3_node.ops.segment import segment_sum
+
+
+class VariancePreservingAggregation(Aggregation):
+    r"""Performs the Variance Preserving Aggregation (VPA) from the `"GNN-VPA:
+    A Variance-Preserving Aggregation Strategy for Graph Neural Networks"
+    <https://arxiv.org/abs/2403.04747>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import VariancePreservingAggregation
+
+        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 = VariancePreservingAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        out = self.reduce(x, index, ptr, dim_size, dim, reduce="sum")
+
+        if ptr is not None and index is None:
+            from .base import ptr2index
+            index = ptr2index(ptr)
+
+        index = ops.cast(index, dtype="int32")
+        if dim_size is None:  # a tensor while tracing; don't test its truth value
+            dim_size = int(ops.max(index)) + 1 if ops.shape(index)[0] > 0 else 0
+
+        ones = ops.ones((ops.shape(index)[0], 1), dtype=out.dtype)
+        count = segment_sum(ones, index, num_segments=dim_size)
+        count = ops.maximum(ops.sqrt(count), 1.0)
+
+        return out / count
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/layers/attention/__init__.py

@@ -0,0 +1,5 @@
+from .performer import *
+from .polynormer import *
+from .qformer import *
+from .sgformer import *
+from .pair_attention import *