k3-node 1.0.0__py3-none-any.whl

k3_node/layers/aggr/patch_transformer.py

@@ -0,0 +1,137 @@
+import math
+from typing import List, Optional, Union
+from keras import layers, ops
+import numpy as np
+
+from .base import Aggregation
+from .utils import MultiheadAttentionBlock
+
+
+class PatchTransformerAggregation(Aggregation):
+    r"""Performs patch transformer aggregation in which the elements to
+    aggregate are processed by multi-head attention blocks across patches.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PatchTransformerAggregation
+
+        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 = PatchTransformerAggregation(in_channels=8, out_channels=16, patch_size=2, hidden_channels=8)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        patch_size: int,
+        hidden_channels: int,
+        num_transformer_blocks: int = 1,
+        heads: int = 1,
+        dropout: float = 0.0,
+        aggr: Union[str, List[str]] = "mean",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.patch_size = patch_size
+        self.hidden_channels = hidden_channels
+        self.aggrs = [aggr] if isinstance(aggr, str) else list(aggr)
+
+        self.lin = layers.Dense(hidden_channels)
+        self.pad_projector = layers.Dense(hidden_channels)
+        self.blocks = [
+            MultiheadAttentionBlock(
+                channels=hidden_channels,
+                heads=heads,
+                layer_norm=True,
+                dropout=dropout,
+            )
+            for _ in range(num_transformer_blocks)
+        ]
+        self.fc = layers.Dense(out_channels)
+
+    def reset_parameters(self):
+        self.lin.reset_parameters()
+        self.pad_projector.reset_parameters()
+        for block in self.blocks:
+            block.reset_parameters()
+        self.fc.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,
+    ):
+        if max_num_elements is None:
+            from k3_node.layers.conv.utils import is_tracing
+            if is_tracing(x) or is_tracing(index):
+                if hasattr(x, "shape") and x.shape[0] is not None:
+                    max_num_elements = int(x.shape[0])
+                else:
+                    max_num_elements = 16
+            elif ptr is not None:
+                ptr_np = ops.convert_to_numpy(ptr)
+                count = ptr_np[1:] - ptr_np[:-1]
+                max_num_elements = int(np.max(count)) if len(count) > 0 else 1
+            else:
+                idx_np = ops.convert_to_numpy(index).astype(np.int64)
+                counts = np.bincount(idx_np)
+                max_num_elements = int(np.max(counts)) if len(counts) > 0 else 1
+
+        # Ensure max_num_elements is a multiple of patch_size
+        num_patches = max(math.ceil(max_num_elements / self.patch_size), 1)
+        target_elements = num_patches * self.patch_size
+
+        x_dense, _ = self.to_dense_batch(
+            x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+            max_num_elements=target_elements,
+        )
+
+        B = ops.shape(x_dense)[0]
+        x_proj = self.lin(x_dense)  # [B, target_elements, hidden_channels]
+
+        # Reshape to patches: [B, num_patches, patch_size * hidden_channels]
+        x_patches = ops.reshape(x_proj, (B, num_patches, self.patch_size * self.hidden_channels))
+        x_patches = self.pad_projector(x_patches)  # [B, num_patches, hidden_channels]
+
+        # Process through transformer blocks
+        for block in self.blocks:
+            x_patches = block(x_patches, x_patches, training=training)
+
+        outs = []
+        for aggr_mode in self.aggrs:
+            if aggr_mode == "mean":
+                outs.append(ops.mean(x_patches, axis=1))
+            elif aggr_mode == "sum":
+                outs.append(ops.sum(x_patches, axis=1))
+            elif aggr_mode == "max":
+                outs.append(ops.max(x_patches, axis=1))
+            elif aggr_mode == "min":
+                outs.append(ops.min(x_patches, axis=1))
+            elif aggr_mode == "var":
+                outs.append(ops.var(x_patches, axis=1))
+            elif aggr_mode == "std":
+                outs.append(ops.std(x_patches, axis=1))
+
+        combined = ops.concatenate(outs, axis=-1) if len(outs) > 1 else outs[0]
+        return self.fc(combined)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.in_channels}, "
+            f"{self.out_channels}, patch_size={self.patch_size})"
+        )
+

k3_node/layers/aggr/quantile.py

@@ -0,0 +1,125 @@
+from typing import List, Optional, Union
+from keras import ops
+import numpy as np
+
+from .base import Aggregation
+
+
+class QuantileAggregation(Aggregation):
+    r"""An aggregation operator that returns the feature-wise :math:`q`-th
+    quantile of a set :math:`\mathcal{X}`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import QuantileAggregation
+
+        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 = QuantileAggregation(q=0.75)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+    interpolations = {"linear", "lower", "higher", "nearest", "midpoint"}
+
+    def __init__(
+        self,
+        q: Union[float, List[float]],
+        interpolation: str = "linear",
+        fill_value: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        qs = [q] if not isinstance(q, (list, tuple)) else list(q)
+        if len(qs) == 0:
+            raise ValueError("Provide at least one quantile value for `q`.")
+        if not all(0.0 <= quantile <= 1.0 for quantile in qs):
+            raise ValueError("`q` must be in the range [0, 1].")
+        if interpolation not in self.interpolations:
+            raise ValueError(f"Invalid interpolation method got ('{interpolation}')")
+
+        self.q = qs
+        self.interpolation = interpolation
+        self.fill_value = fill_value
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        self.assert_index_present(index)
+        from k3_node.ops.segment import segment_sum
+
+        # Sort every set's values in a dense [sets, max_size, features] tensor; the padding sorts
+        # last (+inf) and is then zeroed. Static shapes and differentiable, like PyG's version.
+        dense, _ = self.to_dense_batch(x, index=index, ptr=ptr, dim_size=dim_size, dim=dim,
+                                       fill_value=float("inf"))
+        dense = ops.sort(dense, axis=1)
+        dense = ops.where(ops.isinf(dense), ops.zeros_like(dense), dense)
+
+        index_i = ops.cast(index, "int32")
+        count = segment_sum(ops.ones_like(index_i), index_i, num_segments=ops.shape(dense)[0])
+        count_f = ops.cast(count, dense.dtype)
+        last = ops.maximum(count - 1, 0)
+
+        def gather(position):  # the value at `position` of every set, per feature
+            position = ops.minimum(ops.maximum(ops.cast(position, "int32"), 0), last)
+            position = ops.broadcast_to(ops.reshape(position, (-1, 1, 1)),
+                                        (ops.shape(dense)[0], 1, ops.shape(dense)[2]))
+            return ops.take_along_axis(dense, position, axis=1)[:, 0]
+
+        outs = []
+        for q_val in self.q:
+            q_point = q_val * (count_f - 1.0)
+            if self.interpolation == "lower":
+                quantile = gather(ops.floor(q_point))
+            elif self.interpolation == "higher":
+                quantile = gather(ops.ceil(q_point))
+            elif self.interpolation == "nearest":
+                quantile = gather(ops.round(q_point))
+            else:
+                low, high = gather(ops.floor(q_point)), gather(ops.ceil(q_point))
+                if self.interpolation == "linear":
+                    frac = ops.expand_dims(q_point - ops.floor(q_point), -1)
+                    quantile = low + (high - low) * frac
+                else:  # midpoint
+                    quantile = 0.5 * low + 0.5 * high
+            empty = ops.expand_dims(count == 0, -1)
+            outs.append(ops.where(empty, ops.cast(self.fill_value, quantile.dtype), quantile))
+        return outs[0] if len(outs) == 1 else ops.concatenate(outs, axis=-1)
+
+    def __repr__(self) -> str:
+        q_str = self.q[0] if len(self.q) == 1 else self.q
+        return f"{self.__class__.__name__}(q={q_str})"
+
+
+class MedianAggregation(QuantileAggregation):
+    r"""An aggregation operator that returns the feature-wise median of a set.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MedianAggregation
+
+        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 = MedianAggregation()
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(self, fill_value: float = 0.0, **kwargs):
+        super().__init__(0.5, "lower", fill_value=fill_value, **kwargs)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+

k3_node/layers/aggr/resolver.py

@@ -0,0 +1,68 @@
+from typing import Union
+
+from .base import Aggregation
+
+
+def aggregation_resolver(
+    query: Union[str, Aggregation],
+    *args,
+    **kwargs,
+) -> Aggregation:
+    r"""Resolves an aggregation string or instance to an `Aggregation` object.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import aggregation_resolver
+
+        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 = aggregation_resolver("mean")  # build an Aggregation from its name
+        print(type(aggr).__name__, tuple(aggr(x, index=index, dim_size=2).shape))  # MeanAggregation (2, 8)
+        ```
+    """
+    if isinstance(query, Aggregation):
+        return query
+
+    if not isinstance(query, str):
+        raise ValueError(f"Expected string or Aggregation instance, got {type(query)}")
+
+    query_norm = query.lower().strip()
+
+    from .basic import (
+        MaxAggregation,
+        MeanAggregation,
+        MinAggregation,
+        MulAggregation,
+        PowerMeanAggregation,
+        SoftmaxAggregation,
+        StdAggregation,
+        SumAggregation,
+        VarAggregation,
+    )
+    from .quantile import MedianAggregation, QuantileAggregation
+    from .variance_preserving import VariancePreservingAggregation
+
+    AGGR_DICT = {
+        "sum": SumAggregation,
+        "add": SumAggregation,
+        "mean": MeanAggregation,
+        "max": MaxAggregation,
+        "min": MinAggregation,
+        "mul": MulAggregation,
+        "var": VarAggregation,
+        "std": StdAggregation,
+        "softmax": SoftmaxAggregation,
+        "powermean": PowerMeanAggregation,
+        "median": MedianAggregation,
+        "quantile": QuantileAggregation,
+        "variance_preserving": VariancePreservingAggregation,
+        "vpa": VariancePreservingAggregation,
+    }
+
+    if query_norm in AGGR_DICT:
+        return AGGR_DICT[query_norm](*args, **kwargs)
+
+    raise ValueError(f"Could not resolve aggregation '{query}'")
+

k3_node/layers/aggr/scaler.py

@@ -0,0 +1,133 @@
+from typing import Any, Dict, List, Optional, Union
+from keras import initializers, ops
+import numpy as np
+
+from .base import Aggregation
+from k3_node.ops.segment import segment_sum
+from k3_node.ops.creation import full
+
+
+class DegreeScalerAggregation(Aggregation):
+    r"""Combines one or more aggregators and transforms its output with one or
+    more scalers as introduced in the `"Principal Neighbourhood Aggregation for
+    Graph Nets" <https://arxiv.org/abs/2004.05718>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import DegreeScalerAggregation
+
+        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
+
+        deg = np.array([0, 3, 5, 2])  # in-degree histogram of the training graphs
+        aggr = DegreeScalerAggregation(aggr=["mean", "max"], scaler=["identity", "amplification"], deg=deg)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 32): 2 aggregators x 2 scalers x 8 features
+        ```
+    """
+
+    def __init__(
+        self,
+        aggr: Union[str, List[str], Aggregation],
+        scaler: Union[str, List[str]],
+        deg,
+        train_norm: bool = False,
+        aggr_kwargs: Optional[List[Dict[str, Any]]] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        from .resolver import aggregation_resolver
+        from .multi import MultiAggregation
+
+        if isinstance(aggr, (str, Aggregation)):
+            self.aggr = aggregation_resolver(aggr, **(aggr_kwargs or {}))
+        elif isinstance(aggr, (tuple, list)):
+            self.aggr = MultiAggregation(aggr, aggr_kwargs)
+        else:
+            raise ValueError(
+                f"Only strings, list, tuples and instances of "
+                f"`Aggregation` are valid aggregation schemes (got '{type(aggr)}')"
+            )
+
+        self.scaler = [scaler] if isinstance(scaler, str) else list(scaler)
+
+        deg_np = ops.convert_to_numpy(deg).astype(np.float32)
+        N = float(np.sum(deg_np))
+        bin_degree = np.arange(len(deg_np), dtype=np.float32)
+
+        self.init_avg_deg_lin = float(np.sum(bin_degree * deg_np)) / max(N, 1.0)
+        self.init_avg_deg_log = float(np.sum(np.log(bin_degree + 1.0) * deg_np)) / max(N, 1.0)
+        self.train_norm = train_norm
+
+        if train_norm:
+            self.avg_deg_lin = self.add_weight(
+                shape=(1,),
+                initializer=initializers.Constant(self.init_avg_deg_lin),
+                trainable=True,
+                name="avg_deg_lin",
+            )
+            self.avg_deg_log = self.add_weight(
+                shape=(1,),
+                initializer=initializers.Constant(self.init_avg_deg_log),
+                trainable=True,
+                name="avg_deg_log",
+            )
+        else:
+            self.avg_deg_lin = self.init_avg_deg_lin
+            self.avg_deg_log = self.init_avg_deg_log
+
+    def reset_parameters(self):
+        if hasattr(self.aggr, "reset_parameters"):
+            self.aggr.reset_parameters()
+        if self.train_norm:
+            self.avg_deg_lin.assign(full((1,), self.init_avg_deg_lin, dtype=self.avg_deg_lin.dtype))
+            self.avg_deg_log.assign(full((1,), self.init_avg_deg_log, dtype=self.avg_deg_log.dtype))
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        self.assert_index_present(index)
+
+        out = self.aggr(x, index=index, ptr=ptr, dim_size=dim_size, dim=dim)
+
+        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
+
+        # Compute degree per index
+        ones = ops.ones((ops.shape(index)[0],), dtype=out.dtype)
+        deg = segment_sum(ones, index, num_segments=dim_size)
+        deg = ops.reshape(deg, (dim_size,) + (1,) * (len(ops.shape(out)) - 1))
+
+        avg_deg_log = self.avg_deg_log
+        avg_deg_lin = self.avg_deg_lin
+
+        outs = []
+        for scaler in self.scaler:
+            if scaler == "identity":
+                out_scaler = out
+            elif scaler == "amplification":
+                out_scaler = out * (ops.log(deg + 1.0) / avg_deg_log)
+            elif scaler == "attenuation":
+                out_scaler = out * (avg_deg_log / ops.log(ops.maximum(deg, 1.0) + 1.0))
+            elif scaler == "linear":
+                out_scaler = out * (deg / avg_deg_lin)
+            elif scaler == "inverse_linear":
+                out_scaler = out * (avg_deg_lin / ops.maximum(deg, 1.0))
+            else:
+                raise ValueError(f"Unknown scaler '{scaler}'")
+            outs.append(out_scaler)
+
+        return ops.concatenate(outs, axis=-1) if len(outs) > 1 else outs[0]
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(aggr={self.aggr}, scaler={self.scaler})"
+

k3_node/layers/aggr/set2set.py

@@ -0,0 +1,87 @@
+from typing import Optional
+from keras import layers, ops
+
+from .base import Aggregation
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class Set2Set(Aggregation):
+    r"""The Set2Set aggregation operator based on iterative content-based
+    attention, as described in the `"Order Matters: Sequence to sequence for
+    Sets" <https://arxiv.org/abs/1511.06391>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import Set2Set
+
+        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 = Set2Set(in_channels=8, processing_steps=2)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(self, in_channels: int, processing_steps: int, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = 2 * in_channels
+        self.processing_steps = processing_steps
+        self.lstm_cell = layers.LSTMCell(in_channels)
+
+    def build(self, input_shape=None):
+        self.lstm_cell.build((None, self.out_channels))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        self.lstm_cell.reset_parameters()
+
+    def call(
+        self,
+        x,
+        index: Optional[any] = None,
+        ptr: Optional[any] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        **kwargs,
+    ):
+        if ptr is not None and index is None:
+            from .base import ptr2index
+            index = ptr2index(ptr)
+
+        self.assert_index_present(index)
+        self.assert_two_dimensional_input(x, dim)
+
+        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
+
+        # Initial hidden states: [dim_size, in_channels]
+        h = [
+            ops.zeros((dim_size, self.in_channels), dtype=x.dtype),
+            ops.zeros((dim_size, self.in_channels), dtype=x.dtype),
+        ]
+        q_star = ops.zeros((dim_size, self.out_channels), dtype=x.dtype)
+
+        for _ in range(self.processing_steps):
+            q, h = self.lstm_cell(q_star, h)
+            q_taken = ops.take(q, index, axis=0)
+            e = ops.sum(x * q_taken, axis=-1, keepdims=True)
+
+            max_e = segment_max(e, index, num_segments=dim_size)
+            max_e_exp = ops.take(max_e, index, axis=0)
+            exp_e = ops.exp(e - max_e_exp)
+            sum_exp = segment_sum(exp_e, index, num_segments=dim_size)
+            sum_exp_exp = ops.take(sum_exp, index, axis=0)
+            a = exp_e / ops.maximum(sum_exp_exp, 1e-12)
+
+            r = segment_sum(a * x, index, num_segments=dim_size)
+            q_star = ops.concatenate([q, r], axis=-1)
+
+        return q_star
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels}, {self.out_channels})"
+

k3_node/layers/aggr/set_transformer.py

@@ -0,0 +1,107 @@
+from typing import Optional
+from keras import layers, ops
+
+from .base import Aggregation
+from .utils import PoolingByMultiheadAttention, SetAttentionBlock
+
+
+class SetTransformerAggregation(Aggregation):
+    r"""Performs "Set Transformer" aggregation in which the elements to
+    aggregate are processed by multi-head attention blocks, 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 SetTransformerAggregation
+
+        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 = SetTransformerAggregation(channels=8, num_seed_points=2)
+        out = aggr(x, index=index, dim_size=2)
+        print(tuple(out.shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        num_seed_points: int = 1,
+        num_encoder_blocks: int = 1,
+        num_decoder_blocks: int = 1,
+        heads: int = 1,
+        concat: bool = True,
+        layer_norm: bool = False,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.num_seed_points = num_seed_points
+        self.heads = heads
+        self.concat = concat
+        self.layer_norm = layer_norm
+        self.dropout = dropout
+
+        self.encoders = [
+            SetAttentionBlock(channels, heads, layer_norm, dropout)
+            for _ in range(num_encoder_blocks)
+        ]
+        self.pma = PoolingByMultiheadAttention(
+            channels, num_seed_points, heads, layer_norm, dropout
+        )
+        self.decoders = [
+            SetAttentionBlock(channels, heads, layer_norm, dropout)
+            for _ in range(num_decoder_blocks)
+        ]
+
+    def reset_parameters(self):
+        for encoder in self.encoders:
+            encoder.reset_parameters()
+        self.pma.reset_parameters()
+        for decoder in self.decoders:
+            decoder.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,
+        )
+
+        for encoder in self.encoders:
+            x_dense = encoder(x_dense, mask=mask, training=training)
+
+        x_dense = self.pma(x_dense, mask=mask, training=training)
+
+        for decoder in self.decoders:
+            x_dense = decoder(x_dense, training=training)
+
+        # Handle NaNs if any
+        x_dense = ops.where(ops.isnan(x_dense), 0.0, x_dense)
+
+        if self.concat:
+            B = ops.shape(x_dense)[0]
+            return ops.reshape(x_dense, (B, self.num_seed_points * self.channels))
+        else:
+            return ops.mean(x_dense, axis=1)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.channels}, "
+            f"num_seed_points={self.num_seed_points}, "
+            f"heads={self.heads}, layer_norm={self.layer_norm}, "
+            f"dropout={self.dropout})"
+        )
+