k3-node 1.0.0__py3-none-any.whl

k3_node/layers/pool/max_pool.py

@@ -0,0 +1,159 @@
+from typing import Callable, Optional, Tuple
+from keras import ops
+
+from .consecutive import consecutive_cluster
+from .pool import as_mutable_graph, pool_batch, pool_edge, pool_pos
+from k3_node.ops.segment import segment_max
+from k3_node.ops.host import to_numpy
+
+
+def _max_pool_x(cluster, x, size: Optional[int] = None):
+    cluster = ops.cast(cluster, dtype="int32")
+    if size is None:
+        size = int(to_numpy(cluster).max()) + 1 if ops.shape(cluster)[0] > 0 else 0
+    return segment_max(x, cluster, num_segments=size)
+
+
+def max_pool_x(
+    cluster,
+    x,
+    batch,
+    batch_size: Optional[int] = None,
+    size: Optional[int] = None,
+) -> Tuple[any, Optional[any]]:
+    r"""Max-pools node features according to the clustering defined in `cluster`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import max_pool_x
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+        cluster = np.repeat(np.arange(5), 2)  # merge nodes pairwise into 5 clusters
+
+        x_pool, batch_pool = max_pool_x(cluster, x, batch)
+        print(tuple(x_pool.shape))  # (5, 8)
+        ```
+    """
+    if size is not None:
+        if batch_size is None:
+            batch_size = int(to_numpy(batch).max()) + 1
+        return _max_pool_x(cluster, x, batch_size * size), None
+
+    cluster, perm = consecutive_cluster(cluster)
+    x = _max_pool_x(cluster, x)
+    batch = pool_batch(perm, batch)
+    return x, batch
+
+
+def max_pool(
+    cluster,
+    data,
+    transform: Optional[Callable] = None,
+    edge_index: Optional[any] = None,
+    edge_attr: Optional[any] = None,
+    batch: Optional[any] = None,
+    pos: Optional[any] = None,
+):
+    r"""Pools and coarsens a graph given by `data` according to `cluster`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import max_pool
+
+        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
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+        cluster = np.repeat(np.arange(5), 2)  # merge nodes pairwise into 5 clusters
+
+        x_pool, edge_index_pool, batch_pool = max_pool(cluster, x, edge_index, batch=batch)
+        print(tuple(x_pool.shape))  # (5, 8)
+        ```
+    """
+    cluster, perm = consecutive_cluster(cluster)
+
+    if hasattr(data, "x"):
+        data = as_mutable_graph(data)
+        x = getattr(data, "x", None)
+        if x is not None:
+            data.x = _max_pool_x(cluster, x)
+
+        edge_index = getattr(data, "edge_index", None)
+        edge_attr = getattr(data, "edge_attr", None)
+        if edge_index is not None:
+            data.edge_index, data.edge_attr = pool_edge(cluster, edge_index, edge_attr)
+
+        batch = getattr(data, "batch", None)
+        if batch is not None:
+            data.batch = pool_batch(perm, batch)
+
+        pos = getattr(data, "pos", None)
+        if pos is not None:
+            data.pos = pool_pos(cluster, pos)
+
+        if transform is not None:
+            data = transform(data)
+
+        return data
+
+    # Raw tensor mode
+    pooled_x = _max_pool_x(cluster, data)
+    pooled_edge_index, pooled_edge_attr = (None, None)
+    if edge_index is not None:
+        pooled_edge_index, pooled_edge_attr = pool_edge(cluster, edge_index, edge_attr)
+    pooled_batch = pool_batch(perm, batch) if batch is not None else None
+
+    return pooled_x, pooled_edge_index, pooled_batch
+
+
+def max_pool_neighbor_x(
+    data,
+    edge_index=None,
+    flow: str = "source_to_target",
+):
+    r"""Max-pools neighboring node features.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import max_pool_neighbor_x
+
+        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
+
+        out = max_pool_neighbor_x(x, edge_index=edge_index)  # pool each node with its neighbors
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    if hasattr(data, "x"):
+        x = data.x
+        edge_index = data.edge_index
+        is_data_obj = True
+    else:
+        x = data
+        is_data_obj = False
+        if edge_index is None:
+            raise ValueError("edge_index must be provided if data is a tensor.")
+
+    num_nodes = getattr(data, "num_nodes", None) if is_data_obj else ops.shape(x)[0]
+    if num_nodes is None:
+        num_nodes = ops.shape(x)[0]
+
+    # Add self-loops
+    loop_idx = ops.arange(num_nodes, dtype=edge_index.dtype)
+    loop_edge = ops.stack([loop_idx, loop_idx], axis=0)
+    full_edge_index = ops.concatenate([edge_index, loop_edge], axis=1)
+
+    row = full_edge_index[0]
+    col = full_edge_index[1]
+    row, col = (row, col) if flow == "source_to_target" else (col, row)
+
+    col = ops.cast(col, dtype="int32")
+    x_src = ops.take(x, row, axis=0)
+    out_x = segment_max(x_src, col, num_segments=num_nodes)
+    if is_data_obj:
+        data.x = out_x
+        return data
+    return out_x

k3_node/layers/pool/mem_pool.py

@@ -0,0 +1,145 @@
+from typing import Optional, Tuple
+from k3_node.layers.aggr.base import to_dense_batch
+from keras import initializers, layers, ops
+
+EPS = 1e-15
+
+
+class MemPooling(layers.Layer):
+    r"""Memory based pooling layer from `"Memory-Based Graph Networks"
+    <https://arxiv.org/abs/2002.09518>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import MemPooling
+
+        x = np.random.rand(2, 5, 8).astype("float32")  # dense batch: 2 graphs with 5 nodes each
+        layer = MemPooling(in_channels=8, out_channels=16, heads=2, num_clusters=3)
+        x_pool, assignment = layer(x)
+        print(tuple(x_pool.shape))  # (2, 3, 16)
+        print(tuple(assignment.shape))  # (2, 5, 3): soft assignment of nodes to clusters
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: int,
+        num_clusters: int,
+        tau: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.num_clusters = num_clusters
+        self.tau = tau
+
+        self.k = self.add_weight(
+            shape=(heads, num_clusters, in_channels),
+            initializer=initializers.RandomUniform(minval=-1.0, maxval=1.0),
+            trainable=True,
+            name="k",
+        )
+        self.conv_weight = self.add_weight(
+            shape=(heads, 1),
+            initializer=initializers.GlorotUniform(),
+            trainable=True,
+            name="conv_weight",
+        )
+        self.lin = layers.Dense(out_channels, use_bias=False, name="lin")
+
+    def build(self, input_shape=None):
+        self.lin.build((None, self.num_clusters, self.in_channels))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        init = initializers.RandomUniform(minval=-1.0, maxval=1.0)
+        self.k.assign(init(self.k.shape, dtype=self.k.dtype))
+        glorot = initializers.GlorotUniform()
+        self.conv_weight.assign(glorot(self.conv_weight.shape, dtype=self.conv_weight.dtype))
+
+    @staticmethod
+    def kl_loss(S) -> any:
+        r"""The additional KL divergence-based loss."""
+        S_2 = ops.power(S, 2)
+        P = S_2 / ops.sum(S, axis=1, keepdims=True)
+        denom = ops.sum(P, axis=2, keepdims=True)
+        denom = ops.where(ops.sum(S, axis=2, keepdims=True) == 0.0, 1.0, denom)
+        P = P / denom
+
+        S_clamped = ops.maximum(S, EPS)
+        P_clamped = ops.maximum(P, EPS)
+        # KL(P || S) = sum(P * (log(P) - log(S)))
+        kl = P * (ops.log(P_clamped) - ops.log(S_clamped))
+        return ops.mean(ops.sum(kl, axis=(1, 2)))
+
+    def call(
+        self,
+        x,
+        batch: Optional[any] = None,
+        mask: Optional[any] = None,
+        max_num_nodes: Optional[int] = None,
+        batch_size: Optional[int] = None,
+    ) -> Tuple[any, any]:
+        r"""Forward pass."""
+        if len(ops.shape(x)) == 2:
+            # Node-level input: one dense [num_nodes, channels] block per graph (as in PyG)
+            if batch is None:
+                batch = ops.zeros((ops.shape(x)[0],), dtype="int32")
+            x, mask = to_dense_batch(x, ops.cast(batch, "int32"), dim_size=batch_size,
+                                     max_num_elements=max_num_nodes)
+        elif mask is None:
+            mask = ops.ones((ops.shape(x)[0], ops.shape(x)[1]), dtype=bool)
+
+        B = ops.shape(x)[0]
+        N = ops.shape(x)[1]
+        H = self.heads
+        K = self.num_clusters
+
+        # Compute pairwise squared Euclidean distance between k and x
+        # k: [H, K, C], x: [B, N, C]
+        # Reshape to [H * K, C] and [B * N, C]
+        k_flat = ops.reshape(self.k, (H * K, self.in_channels))
+        x_flat = ops.reshape(x, (B * N, self.in_channels))
+
+        k_sq = ops.sum(ops.power(k_flat, 2), axis=-1, keepdims=True)  # [HK, 1]
+        x_sq = ops.sum(ops.power(x_flat, 2), axis=-1, keepdims=True)  # [BN, 1]
+        dot = ops.matmul(k_flat, ops.transpose(x_flat))  # [HK, BN]
+        dist = ops.maximum(k_sq + ops.transpose(x_sq) - 2.0 * dot, 0.0)  # [HK, BN]
+
+        dist = ops.power(1.0 + dist / self.tau, -(self.tau + 1.0) / 2.0)
+        # Reshape to [H, K, B, N] then permute to [B, H, N, K]
+        dist = ops.reshape(dist, (H, K, B, N))
+        dist = ops.transpose(dist, (2, 0, 3, 1))
+
+        S = dist / ops.sum(dist, axis=-1, keepdims=True)  # [B, H, N, K]
+
+        # Conv over head dimension: reduce H -> 1
+        # conv_weight: [H, 1]
+        # Transpose S to [B, N, K, H] and multiply by [H, 1]
+        S_perm = ops.transpose(S, (0, 2, 3, 1))  # [B, N, K, H]
+        S_conv = ops.squeeze(ops.matmul(S_perm, self.conv_weight), axis=-1)  # [B, N, K]
+        # [B, N, K]. With one cluster the softmax is 1 with zero gradient (as in PyG); `0 * S_conv`
+        # keeps the keys and conv weights in the graph, so Keras doesn't warn about them.
+        S = ops.softmax(S_conv, axis=-1) if K > 1 else 1.0 + 0.0 * S_conv
+
+        mask_f = ops.cast(ops.reshape(mask, (B, N, 1)), S.dtype)
+        S = S * mask_f
+
+        # x_out: [B, K, out_channels]
+        # S.transpose(1, 2) is [B, K, N]
+        # x is [B, N, C]
+        pooled_x = ops.matmul(ops.swapaxes(S, 1, 2), x)  # [B, K, C]
+        x_out = self.lin(pooled_x)
+
+        return x_out, S
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self.in_channels}, "
+            f"{self.out_channels}, heads={self.heads}, "
+            f"num_clusters={self.num_clusters})"
+        )

k3_node/layers/pool/pan_pool.py

@@ -0,0 +1,144 @@
+from typing import Callable, Optional, Tuple, Union
+from keras import initializers, layers, ops
+
+from .connect.filter_edges import FilterEdges
+from .select.topk import SelectTopK
+from k3_node.ops.segment import segment_sum
+from k3_node.ops.creation import full
+
+
+class PANPooling(layers.Layer):
+    r"""The path integral based pooling operator from the
+    `"Path Integral Based Convolution and Pooling for Graph Neural Networks"
+    <https://arxiv.org/abs/2006.16811>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import PANPooling, PANConv
+
+        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
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        # PANPooling consumes the path weights produced by PANConv
+        x, weights = PANConv(in_channels=8, out_channels=8, filter_size=2)(x, edge_index)
+        layer = PANPooling(in_channels=8, ratio=0.5)
+        x_pool, edge_index_pool, edge_attr_pool, batch_pool, perm, score = layer(x, weights, batch=batch)
+        print(tuple(x_pool.shape))  # (6, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: float = 0.5,
+        min_score: Optional[float] = None,
+        multiplier: float = 1.0,
+        nonlinearity: Union[str, Callable] = "tanh",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.in_channels = in_channels
+        self.ratio = ratio
+        self.min_score = min_score
+        self.multiplier = multiplier
+        self.nonlinearity = nonlinearity
+
+        self.p = self.add_weight(
+            shape=(in_channels,),
+            initializer=initializers.Constant(1.0),
+            trainable=True,
+            name="p",
+        )
+        self.beta = self.add_weight(
+            shape=(2,),
+            initializer=initializers.Constant(0.5),
+            trainable=True,
+            name="beta",
+        )
+
+        self.select = SelectTopK(1, ratio, min_score, nonlinearity)
+        self.connect = FilterEdges()
+
+    def reset_parameters(self):
+        self.p.assign(ops.ones(self.p.shape, dtype=self.p.dtype))
+        self.beta.assign(full(self.beta.shape, 0.5, dtype=self.beta.dtype))
+        self.select.reset_parameters()
+
+    def build(self, input_shape=None):
+        if hasattr(self.select, "built") and not self.select.built:
+            self.select.build(None)
+        if hasattr(self.connect, "built") and not self.connect.built:
+            self.connect.build(None)
+        self.built = True
+
+    def call(
+        self,
+        x,
+        M,
+        batch: Optional[any] = None,
+    ) -> Tuple[any, any, any, Optional[any], any, any]:
+        r"""Forward pass.
+
+        Args:
+            x: Node feature matrix.
+            M: MET matrix, either a tuple/list (edge_index, edge_weight) or an object
+               with `.coo()` method (like PyG SparseTensor).
+            batch: Batch vector.
+        """
+        num_nodes = ops.shape(x)[0]
+        if batch is None:
+            batch = ops.zeros((num_nodes,), dtype="int32")
+
+        if hasattr(M, "coo"):
+            row, col, edge_weight = M.coo()
+        elif isinstance(M, (tuple, list)) and len(M) >= 2 and not isinstance(M[0], int):
+            edge_index, edge_weight = M[0], M[1]
+            row, col = edge_index[0], edge_index[1]
+        elif hasattr(M, "shape") and len(M.shape) == 2 and M.shape[0] == 2:
+            row, col = M[0], M[1]
+            edge_weight = ops.ones((ops.shape(col)[0],), dtype=x.dtype)
+        else:
+            nz = ops.where(M != 0)
+            row, col = nz[0], nz[1]
+            edge_weight = ops.take(ops.reshape(M, (-1,)), row * ops.shape(M)[1] + col, axis=0)
+
+        col = ops.cast(col, dtype="int32")
+        row = ops.cast(row, dtype="int32")
+        edge_weight = ops.cast(edge_weight, dtype=x.dtype)
+
+        score1 = ops.sum(x * self.p, axis=-1)
+        score2 = segment_sum(edge_weight, col, num_segments=num_nodes)
+        score = self.beta[0] * score1 + self.beta[1] * score2
+
+        select_out = self.select(score, batch)
+
+        perm = select_out.node_index
+        score_val = select_out.weight
+
+        x_pooled = ops.take(x, perm, axis=0) * ops.expand_dims(score_val, axis=-1)
+        if self.multiplier != 1.0:
+            x_pooled = x_pooled * self.multiplier
+
+        edge_index = ops.stack([col, row], axis=0)
+        connect_out = self.connect(select_out, edge_index, edge_weight, batch)
+
+        return (
+            x_pooled,
+            connect_out.edge_index,
+            connect_out.edge_attr,
+            connect_out.batch,
+            perm,
+            score_val,
+        )
+
+    def __repr__(self) -> str:
+        if self.min_score is None:
+            ratio = f"ratio={self.ratio}"
+        else:
+            ratio = f"min_score={self.min_score}"
+        return (
+            f"{self.__class__.__name__}({self.in_channels}, {ratio}, "
+            f"multiplier={self.multiplier})"
+        )

k3_node/layers/pool/point_cloud.py

@@ -0,0 +1,212 @@
+from typing import Optional
+from keras import ops
+import numpy as np
+
+from .knn import knn
+from k3_node.ops.host import to_numpy
+
+
+def fps(
+    x,
+    batch: Optional[any] = None,
+    ratio: float = 0.5,
+    random_start: bool = True,
+    batch_size: Optional[int] = None,
+):
+    r"""Farthest Point Sampling algorithm.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import fps
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+        batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1
+
+        index = fps(pos, batch, ratio=0.4)  # farthest point sampling: 40% of each graph's points
+        print(tuple(index.shape))  # (4,)
+        ```
+    """
+    x_np = to_numpy(x)
+    num_nodes = x_np.shape[0]
+
+    if batch is None:
+        batch_np = np.zeros(num_nodes, dtype=np.int64)
+    else:
+        batch_np = to_numpy(batch).astype(np.int64)
+
+    unique_batches = np.unique(batch_np)
+    selected_indices = []
+
+    for b in unique_batches:
+        idx_b = np.where(batch_np == b)[0]
+        n_b = len(idx_b)
+        if n_b == 0:
+            continue
+
+        if ratio >= 1:
+            num_samples = min(int(ratio), n_b)
+        else:
+            num_samples = max(1, int(np.ceil(ratio * n_b)))
+
+        pts_b = x_np[idx_b]
+        sampled = []
+
+        start_idx = np.random.randint(n_b) if random_start else 0
+        sampled.append(start_idx)
+        min_dist = np.sum((pts_b - pts_b[start_idx]) ** 2, axis=-1)
+
+        for _ in range(1, num_samples):
+            next_idx = int(np.argmax(min_dist))
+            sampled.append(next_idx)
+            dist_new = np.sum((pts_b - pts_b[next_idx]) ** 2, axis=-1)
+            min_dist = np.minimum(min_dist, dist_new)
+
+        selected_indices.extend(idx_b[sampled])
+
+    return ops.convert_to_tensor(np.array(selected_indices, dtype=np.int64), dtype="int64")
+
+
+def radius(
+    x,
+    y,
+    r: float,
+    batch_x: Optional[any] = None,
+    batch_y: Optional[any] = None,
+    max_num_neighbors: int = 32,
+    num_workers: int = 1,
+    batch_size: Optional[int] = None,
+):
+    r"""Finds for each element in `y` all points in `x` within distance `r`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import radius
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+        query = np.random.rand(4, 3).astype("float32")  # 4 query points
+
+        assign = radius(pos, query, r=0.5)  # all points within distance 0.5 of each query point
+        print(assign.shape[0])  # 2: rows: (query index, point index)
+        ```
+    """
+    x_np = to_numpy(x)
+    y_np = to_numpy(y)
+    if x_np.ndim == 1:
+        x_np = x_np[:, None]
+    if y_np.ndim == 1:
+        y_np = y_np[:, None]
+
+    N = x_np.shape[0]
+    M = y_np.shape[0]
+
+    if batch_x is None:
+        batch_x_np = np.zeros(N, dtype=np.int64)
+    else:
+        batch_x_np = to_numpy(batch_x).astype(np.int64)
+
+    if batch_y is None:
+        batch_y_np = np.zeros(M, dtype=np.int64)
+    else:
+        batch_y_np = to_numpy(batch_y).astype(np.int64)
+
+    rows = []
+    cols = []
+    r_sq = r * r
+
+    for i in range(M):
+        valid_b = batch_x_np == batch_y_np[i]
+        valid_idx = np.where(valid_b)[0]
+        if len(valid_idx) == 0:
+            continue
+
+        diff = x_np[valid_idx] - y_np[i]
+        dist_sq = np.sum(diff**2, axis=-1)
+        within_r = np.where(dist_sq <= r_sq)[0]
+        if len(within_r) > max_num_neighbors:
+            # Sort and take top max_num_neighbors
+            sort_order = np.argsort(dist_sq[within_r])[:max_num_neighbors]
+            within_r = within_r[sort_order]
+
+        for match_idx in valid_idx[within_r]:
+            rows.append(i)
+            cols.append(match_idx)
+
+    if len(rows) == 0:
+        return ops.zeros((2, 0), dtype="int64")
+
+    return ops.convert_to_tensor(np.stack([rows, cols], axis=0), dtype="int64")
+
+
+def radius_graph(
+    x,
+    r: float,
+    batch: Optional[any] = None,
+    loop: bool = False,
+    max_num_neighbors: int = 32,
+    flow: str = "source_to_target",
+    num_workers: int = 1,
+    batch_size: Optional[int] = None,
+):
+    r"""Computes graph edges to all points within a given distance `r`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import radius_graph
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+
+        edge_index = radius_graph(pos, r=0.5)  # connect points closer than 0.5
+        print(edge_index.shape[0])  # 2
+        ```
+    """
+    assert flow in ["source_to_target", "target_to_source"]
+    edge_index = radius(
+        x,
+        x,
+        r,
+        batch_x=batch,
+        batch_y=batch,
+        max_num_neighbors=max_num_neighbors if loop else max_num_neighbors + 1,
+    )
+    edge_index_np = to_numpy(edge_index)
+
+    if not loop and edge_index_np.shape[1] > 0:
+        mask = edge_index_np[0] != edge_index_np[1]
+        edge_index_np = edge_index_np[:, mask]
+
+    if flow == "source_to_target" and edge_index_np.shape[1] > 0:
+        edge_index_np = np.flip(edge_index_np, axis=0)
+
+    return ops.convert_to_tensor(edge_index_np, dtype=edge_index.dtype)
+
+
+def nearest(
+    x,
+    y,
+    batch_x: Optional[any] = None,
+    batch_y: Optional[any] = None,
+):
+    r"""Clusters each point in `x` to its nearest point in `y`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import nearest
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+        query = np.random.rand(4, 3).astype("float32")  # 4 query points
+
+        cluster = nearest(query, pos)  # index of the closest point in `pos` for every query point
+        print(tuple(cluster.shape))  # (4,)
+        ```
+    """
+    edge_index = knn(y, x, k=1, batch_x=batch_y, batch_y=batch_x)
+    return edge_index[1]
+