k3-node 1.0.0__py3-none-any.whl

k3_node/layers/pool/edge_pool.py

@@ -0,0 +1,189 @@
+from typing import Callable, List, NamedTuple, Optional, Tuple
+from keras import layers, ops
+import numpy as np
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class UnpoolInfo(NamedTuple):
+    edge_index: any
+    cluster: any
+    batch: any
+    new_edge_score: any
+
+
+class EdgePooling(layers.Layer):
+    r"""The edge pooling operator from the `"Towards Graph Pooling by Edge
+    Contraction" <https://graphreason.github.io/papers/17.pdf>`__ and
+    `"Edge Contraction Pooling for Graph Neural Networks"
+    <https://arxiv.org/abs/1905.10990>`__ papers.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import EdgePooling
+
+        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
+
+        layer = EdgePooling(in_channels=8)
+        x_pool, edge_index_pool, batch_pool, unpool_info = layer(x, edge_index, batch)
+        # The number of clusters depends on the learned edge scores
+        print(x_pool.shape[0] <= 10, x_pool.shape[1])  # True 8: fewer nodes, same features
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        edge_score_method: Optional[Callable] = None,
+        dropout: float = 0.0,
+        add_to_edge_score: float = 0.5,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        if edge_score_method is None:
+            edge_score_method = self.compute_edge_score_softmax
+        self.compute_edge_score = edge_score_method
+        self.add_to_edge_score = add_to_edge_score
+        self.dropout_rate = dropout
+        self.drop = layers.Dropout(dropout) if dropout > 0.0 else None
+
+        self.lin = layers.Dense(1, use_bias=True, name="lin")
+
+    def build(self, input_shape=None):
+        self.lin.build((None, 2 * self.in_channels))
+        super().build(input_shape)
+
+    def reset_parameters(self):
+        self.lin.reset_parameters()
+
+    @staticmethod
+    def compute_edge_score_softmax(raw_edge_score, edge_index, num_nodes: int):
+        col = ops.cast(edge_index[1], dtype="int32")
+        max_score = segment_max(raw_edge_score, col, num_segments=num_nodes)
+        max_exp = ops.take(max_score, col, axis=0)
+        exp_score = ops.exp(raw_edge_score - max_exp)
+        sum_exp = segment_sum(exp_score, col, num_segments=num_nodes)
+        sum_exp_taken = ops.take(sum_exp, col, axis=0)
+        return exp_score / (sum_exp_taken + 1e-12)
+
+    @staticmethod
+    def compute_edge_score_tanh(raw_edge_score, edge_index=None, num_nodes=None):
+        return ops.tanh(raw_edge_score)
+
+    @staticmethod
+    def compute_edge_score_sigmoid(raw_edge_score, edge_index=None, num_nodes=None):
+        return ops.sigmoid(raw_edge_score)
+
+    def call(
+        self,
+        x,
+        edge_index,
+        batch,
+        training: bool = False,
+    ) -> Tuple[any, any, any, UnpoolInfo]:
+        r"""Forward pass."""
+        row = ops.cast(edge_index[0], dtype="int32")
+        col = ops.cast(edge_index[1], dtype="int32")
+
+        e = ops.concatenate([ops.take(x, row, axis=0), ops.take(x, col, axis=0)], axis=-1)
+        e = ops.reshape(self.lin(e), (-1,))
+        if self.drop is not None:
+            e = self.drop(e, training=training)
+
+        num_nodes = ops.shape(x)[0]
+        e = self.compute_edge_score(e, edge_index, num_nodes)
+        e = e + self.add_to_edge_score
+
+        return self._merge_edges(x, edge_index, batch, e)
+
+    def _merge_edges(
+        self,
+        x,
+        edge_index,
+        batch,
+        edge_score,
+    ) -> Tuple[any, any, any, UnpoolInfo]:
+        from k3_node.layers.conv.utils import eager_only_placeholder, host_callback
+        if eager_only_placeholder("EdgePooling", x, edge_index, edge_score):
+            num_nodes = ops.shape(x)[0]
+            unpool_info = UnpoolInfo(edge_index, ops.arange(num_nodes, dtype="int32"), batch, ops.ones((num_nodes,), dtype=x.dtype))
+            return x, edge_index, batch, unpool_info
+
+        num_nodes = int(ops.shape(x)[0])
+        num_edges = int(ops.shape(edge_index)[1])
+
+        def contract(edge_index_np, edge_score_np, batch_np):
+            # Greedy contraction on the host; outputs are padded to fixed sizes (N clusters, E edges).
+            cluster_np = np.zeros(num_nodes, dtype=np.int64)
+            cluster_edge = np.full(num_nodes, -1, dtype=np.int64)  # edge merged into each cluster
+            mask = np.ones(num_nodes, dtype=bool)
+            i = 0
+            for edge_idx in np.argsort(-edge_score_np, kind="stable"):
+                source, target = edge_index_np[0, edge_idx], edge_index_np[1, edge_idx]
+                if not mask[source] or not mask[target]:
+                    continue
+                cluster_edge[i] = edge_idx
+                cluster_np[source] = i
+                mask[source] = False
+                if source != target:
+                    cluster_np[target] = i
+                    mask[target] = False
+                i += 1
+            remaining = np.where(mask)[0]
+            cluster_np[remaining] = np.arange(i, i + len(remaining))
+            num_clusters = i + len(remaining)
+
+            unique_edges = np.unique(cluster_np[edge_index_np.astype(np.int64)], axis=1)
+            edges_pad = np.zeros((2, num_edges), dtype=np.int64)
+            edges_pad[:, : unique_edges.shape[1]] = unique_edges
+
+            batch_pad = np.zeros(num_nodes, dtype=np.int64)
+            batch_pad[cluster_np] = batch_np
+            return cluster_np, cluster_edge, num_clusters, edges_pad, unique_edges.shape[1], batch_pad
+
+        cluster, cluster_edge, num_clusters, edges_pad, num_new_edges, batch_pad = host_callback(
+            contract,
+            [((num_nodes,), "int32"), ((num_nodes,), "int32"), ((), "int32"),
+             ((2, num_edges), "int32"), ((), "int32"), ((num_nodes,), "int32")],
+            edge_index, edge_score, batch,
+        )
+        num_clusters, num_new_edges = int(num_clusters), int(num_new_edges)
+
+        new_x = segment_sum(x, cluster, num_segments=num_clusters)
+
+        # Score of the edge merged into each cluster (1 for unmatched nodes), gathered from the score
+        # tensor so gradients reach the scoring layer.
+        cluster_edge = cluster_edge[:num_clusters]
+        if num_edges > 0:
+            gathered = ops.take(ops.cast(edge_score, x.dtype), ops.maximum(cluster_edge, 0), axis=0)
+            new_edge_score = ops.where(cluster_edge >= 0, gathered, ops.ones_like(gathered))
+        else:
+            new_edge_score = ops.ones((num_clusters,), dtype=x.dtype)
+        new_x = new_x * ops.reshape(new_edge_score, (-1, 1))
+
+        new_edge_index = ops.cast(edges_pad[:, :num_new_edges], edge_index.dtype)
+        new_batch = ops.cast(batch_pad[:num_clusters], batch.dtype)
+
+        unpool_info = UnpoolInfo(
+            edge_index=edge_index,
+            cluster=cluster,
+            batch=batch,
+            new_edge_score=new_edge_score,
+        )
+
+        return new_x, new_edge_index, new_batch, unpool_info
+
+    def unpool(
+        self,
+        x,
+        unpool_info: UnpoolInfo,
+    ) -> Tuple[any, any, any]:
+        r"""Unpools a previous edge pooling step."""
+        new_x = x / ops.reshape(unpool_info.new_edge_score, (-1, 1))
+        new_x = ops.take(new_x, ops.cast(unpool_info.cluster, dtype="int32"), axis=0)
+        return new_x, unpool_info.edge_index, unpool_info.batch
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels})"

k3_node/layers/pool/glob.py

@@ -0,0 +1,139 @@
+from typing import Optional
+from keras import ops
+
+
+from k3_node.layers.conv.utils import is_tracing
+from k3_node.ops.segment import segment_max, segment_sum
+from k3_node.ops.host import to_numpy
+
+
+def _infer_size(batch, size=None):
+    if size is not None:
+        return size
+    import keras
+    from keras.src.backend.common.symbolic_scope import in_symbolic_scope
+
+    if in_symbolic_scope() or (hasattr(batch, "is_meta") and batch.is_meta):
+        # Keras shape inference on placeholder values: any size gives the right output rank.
+        return 1 if keras.config.backend() == "jax" else None
+    if is_tracing(batch):
+        if keras.config.backend() == "jax":
+            # The number of graphs depends on the values in `batch`, which are unknown inside
+            # jax.jit; guessing would silently merge every graph into one.
+            raise ValueError(
+                "Cannot infer the number of graphs from `batch` inside a compiled JAX function. "
+                "Pass `size=` to the pooling function (or `batch_size=` to the model)."
+            )
+        # TensorFlow graph mode supports a dynamic number of segments.
+        return ops.cast(ops.max(batch), "int32") + 1
+    # Use the maximum (as PyG does), not the last entry: pooling layers such as EdgePooling
+    # return batch vectors that are not sorted by graph.
+    try:
+        if ops.shape(batch)[0] == 0:
+            return 0
+        return int(to_numpy(batch).max()) + 1
+    except Exception:
+        pass
+    try:
+        return ops.cast(ops.max(batch), "int32") + 1
+    except Exception:
+        return None
+
+
+def global_add_pool(x, batch: Optional[any] = None, size: Optional[int] = None):
+    r"""Returns batch-wise graph-level-outputs by adding node features
+    across the node dimension.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import global_add_pool
+
+        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
+
+        out = global_add_pool(x, batch)  # one row per graph
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+    if batch is None:
+        dim = -1 if len(ops.shape(x)) == 1 else 0
+        keepdims = len(ops.shape(x)) <= 2
+        out = ops.sum(x, axis=dim, keepdims=keepdims)
+        if len(ops.shape(x)) == 1:
+            out = ops.reshape(out, (1, 1))
+        return out
+
+    batch = ops.cast(batch, dtype="int32")
+    if batch.shape[0] is not None and batch.shape[0] == 0:
+        num_seg = size if size is not None else 0
+        return ops.zeros((num_seg,) + tuple(ops.shape(x)[1:]), dtype=x.dtype)
+    size = _infer_size(batch, size)
+    return segment_sum(x, batch, num_segments=size)
+
+
+def global_mean_pool(x, batch: Optional[any] = None, size: Optional[int] = None):
+    r"""Returns batch-wise graph-level-outputs by averaging node features
+    across the node dimension.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import global_mean_pool
+
+        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
+
+        out = global_mean_pool(x, batch)  # one row per graph
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+    if batch is None:
+        dim = -1 if len(ops.shape(x)) == 1 else 0
+        keepdims = len(ops.shape(x)) <= 2
+        out = ops.mean(x, axis=dim, keepdims=keepdims)
+        if len(ops.shape(x)) == 1:
+            out = ops.reshape(out, (1, 1))
+        return out
+
+    batch = ops.cast(batch, dtype="int32")
+    if batch.shape[0] is not None and batch.shape[0] == 0:
+        num_seg = size if size is not None else 0
+        return ops.zeros((num_seg,) + tuple(ops.shape(x)[1:]), dtype=x.dtype)
+    size = _infer_size(batch, size)
+    sum_val = segment_sum(x, batch, num_segments=size)
+    ones = ops.ones_like(x)
+    count = segment_sum(ones, batch, num_segments=size)
+    return sum_val / ops.maximum(count, 1.0)
+
+
+def global_max_pool(x, batch: Optional[any] = None, size: Optional[int] = None):
+    r"""Returns batch-wise graph-level-outputs by taking the channel-wise
+    maximum across the node dimension.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import global_max_pool
+
+        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
+
+        out = global_max_pool(x, batch)  # one row per graph
+        print(tuple(out.shape))  # (2, 8)
+        ```
+    """
+    if batch is None:
+        dim = -1 if len(ops.shape(x)) == 1 else 0
+        keepdims = len(ops.shape(x)) <= 2
+        out = ops.max(x, axis=dim, keepdims=keepdims)
+        if len(ops.shape(x)) == 1:
+            out = ops.reshape(out, (1, 1))
+        return out
+
+    batch = ops.cast(batch, dtype="int32")
+    if batch.shape[0] is not None and batch.shape[0] == 0:
+        num_seg = size if size is not None else 0
+        return ops.zeros((num_seg,) + tuple(ops.shape(x)[1:]), dtype=x.dtype)
+    size = _infer_size(batch, size)
+    return segment_max(x, batch, num_segments=size)

k3_node/layers/pool/graclus.py

@@ -0,0 +1,66 @@
+from typing import Optional
+from keras import ops
+import numpy as np
+from k3_node.ops.host import to_numpy
+
+
+def graclus(
+    edge_index,
+    weight: Optional[any] = None,
+    num_nodes: Optional[int] = None,
+):
+    r"""A greedy clustering algorithm of picking an unmarked vertex and matching
+    it with one of its unmarked neighbors that maximizes its edge weight.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import graclus
+
+        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
+
+        cluster = graclus(edge_index, num_nodes=10)  # greedy matching of neighboring nodes
+        print(tuple(cluster.shape))  # (10,): a cluster id for every node
+        ```
+    """
+    edge_index_np = to_numpy(edge_index).astype(np.int64)
+    if num_nodes is None:
+        num_nodes = int(np.max(edge_index_np)) + 1 if edge_index_np.size > 0 else 0
+
+    if weight is not None:
+        weight_np = to_numpy(weight)
+    else:
+        weight_np = np.ones(edge_index_np.shape[1], dtype=np.float32)
+
+    # Build adjacency list
+    adj = [[] for _ in range(num_nodes)]
+    for idx in range(edge_index_np.shape[1]):
+        u = edge_index_np[0, idx]
+        v = edge_index_np[1, idx]
+        w = weight_np[idx]
+        adj[u].append((v, w))
+
+    cluster_np = np.full(num_nodes, -1, dtype=np.int64)
+    c = 0
+    for u in range(num_nodes):
+        if cluster_np[u] != -1:
+            continue
+
+        best_v = None
+        best_w = -float("inf")
+        for v, w in adj[u]:
+            if v != u and cluster_np[v] == -1:
+                if w > best_w:
+                    best_w = w
+                    best_v = v
+
+        if best_v is not None:
+            cluster_np[u] = c
+            cluster_np[best_v] = c
+        else:
+            cluster_np[u] = c
+        c += 1
+
+    return ops.convert_to_tensor(cluster_np, dtype="int64")
+

k3_node/layers/pool/knn.py

@@ -0,0 +1,253 @@
+from typing import NamedTuple, Optional
+from keras import ops
+import numpy as np
+from k3_node.ops.creation import repeat
+
+
+class KNNOutput(NamedTuple):
+    score: any
+    index: any
+
+
+class KNNIndex:
+    r"""A base class to perform k-nearest neighbor search."""
+    def __init__(
+        self,
+        index_factory: Optional[str] = None,
+        emb: Optional[any] = None,
+        reserve: Optional[int] = None,
+    ):
+        self.index_factory = index_factory
+        self.reserve = reserve
+        self.emb = None
+        if emb is not None:
+            self.add(emb)
+
+    @property
+    def numel(self) -> int:
+        if self.emb is None:
+            return 0
+        return ops.shape(self.emb)[0]
+
+    def add(self, emb):
+        if self.emb is None:
+            self.emb = emb
+        else:
+            self.emb = ops.concatenate([self.emb, emb], axis=0)
+
+    def get_emb(self):
+        return self.emb
+
+    def search(self, query, k: int) -> KNNOutput:
+        raise NotImplementedError
+
+
+class L2KNNIndex(KNNIndex):
+    r"""k-NN search using squared Euclidean (L2) distance."""
+    def __init__(self, emb: Optional[any] = None, **kwargs):
+        super().__init__(index_factory="IndexFlatL2", emb=emb, **kwargs)
+
+    def search(self, query, k: int) -> KNNOutput:
+        # query: [M, F], emb: [N, F]
+        q_exp = ops.expand_dims(query, axis=1)  # [M, 1, F]
+        e_exp = ops.expand_dims(self.emb, axis=0)  # [1, N, F]
+        dist = ops.sum(ops.power(q_exp - e_exp, 2), axis=-1)  # [M, N]
+
+        top_k_neg_score, top_k_idx = ops.top_k(-dist, k=k, sorted=True)
+        return KNNOutput(score=-top_k_neg_score, index=top_k_idx)
+
+
+class MIPSKNNIndex(KNNIndex):
+    r"""k-NN search using Maximum Inner Product Search (MIPS)."""
+    def __init__(self, emb: Optional[any] = None, **kwargs):
+        super().__init__(index_factory="IndexFlatIP", emb=emb, **kwargs)
+
+    def search(self, query, k: int) -> KNNOutput:
+        # query: [M, F], emb: [N, F]
+        score_mat = ops.matmul(query, ops.transpose(self.emb))  # [M, N]
+        top_k_score, top_k_idx = ops.top_k(score_mat, k=k, sorted=True)
+        return KNNOutput(score=top_k_score, index=top_k_idx)
+
+
+class ApproxL2KNNIndex(L2KNNIndex):
+    pass
+
+
+class ApproxMIPSKNNIndex(MIPSKNNIndex):
+    pass
+
+
+def knn(
+    x,
+    y,
+    k: int,
+    batch_x: Optional[any] = None,
+    batch_y: Optional[any] = None,
+    cosine: bool = False,
+    num_workers: int = 1,
+    batch_size: Optional[int] = None,
+):
+    r"""Finds for each element in `y` the `k` nearest points in `x`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import knn
+
+        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 = knn(pos, query, k=3)  # 3 nearest points in `pos` for every query point
+        print(tuple(assign.shape))  # (2, 12): [2, num_query * k]: (query index, point index)
+        ```
+    """
+    x = ops.convert_to_tensor(x)
+    y = ops.convert_to_tensor(y)
+    if len(ops.shape(x)) == 1:
+        x = ops.expand_dims(x, axis=-1)
+    if len(ops.shape(y)) == 1:
+        y = ops.expand_dims(y, axis=-1)
+
+    col = _knn_indices(x, y, k, batch_x, batch_y, cosine=cosine)  # [M, k] indices into x
+    M = ops.shape(y)[0]
+    row = repeat(ops.arange(0, M, dtype="int64"), k)
+    return ops.stack([row, ops.reshape(ops.cast(col, "int64"), (-1,))], axis=0)
+
+
+def _pairwise(y, x, cosine):
+    """Squared distances (or cosine distances) between the rows of ``y`` and ``x``, without
+    materializing ``[M, N, features]`` differences."""
+    if cosine:
+        y = y / ops.maximum(ops.norm(y, axis=-1, keepdims=True), 1e-12)
+        x = x / ops.maximum(ops.norm(x, axis=-1, keepdims=True), 1e-12)
+        return 1.0 - ops.matmul(y, ops.swapaxes(x, -1, -2))
+    y_sq = ops.sum(y * y, axis=-1, keepdims=True)
+    x_sq = ops.expand_dims(ops.sum(x * x, axis=-1), -2)
+    return ops.maximum(y_sq + x_sq - 2.0 * ops.matmul(y, ops.swapaxes(x, -1, -2)), 0.0)
+
+
+def _knn_indices(x, y, k, batch_x=None, batch_y=None, cosine=False, exclude_self=False):
+    """Indices ``[M, k]`` into ``x`` of the ``k`` nearest points of every row of ``y`` within the
+    same example. Eagerly, distances are computed per example (``[examples, n, n]``); when compiled,
+    in chunks of query rows masked by example."""
+    from k3_node.layers.aggr.base import to_dense_batch
+    from k3_node.layers.conv.utils import is_tracing
+
+    N, M = ops.shape(x)[0], ops.shape(y)[0]
+    if isinstance(M, int) and isinstance(N, int) and (M == 0 or N == 0):  # sizes are tensors while tracing
+        return ops.zeros((M, k), dtype="int32")
+    ids = [b for b in (batch_x, batch_y) if b is not None]
+    if not any(is_tracing(t) for t in [x, y] + ids) and isinstance(N, int) and isinstance(M, int):
+        try:
+            bx = np.zeros(N, np.int64) if batch_x is None else np.asarray(ops.convert_to_numpy(batch_x)).astype(np.int64)
+            by = np.zeros(M, np.int64) if batch_y is None else np.asarray(ops.convert_to_numpy(batch_y)).astype(np.int64)
+            num_examples = int(max(bx.max(initial=-1), by.max(initial=-1))) + 1
+            x_dense, x_mask = to_dense_batch(x, bx, dim_size=num_examples)
+            y_dense, _ = to_dense_batch(y, by, dim_size=num_examples)
+            dist = _pairwise(y_dense, x_dense, cosine)  # [examples, max_y, max_x]
+            dist = ops.where(ops.expand_dims(x_mask, 1), dist, 1e9)
+            x_start = np.concatenate([[0], np.cumsum(np.bincount(bx, minlength=num_examples))])[:-1]
+            y_start = np.concatenate([[0], np.cumsum(np.bincount(by, minlength=num_examples))])[:-1]
+            y_local = np.arange(M) - y_start[by]
+            if exclude_self:  # x and y are the same points
+                dist = ops.where(ops.expand_dims(ops.eye(dist.shape[1], dist.shape[2]), 0) > 0, 1e9, dist)
+            kk = min(k, int(dist.shape[2]))
+            _, idx = ops.top_k(-dist, k=kk, sorted=True)  # local indices [examples, max_y, kk]
+            flat = ops.reshape(idx, (-1, kk))
+            local = ops.take(flat, by * int(dist.shape[1]) + y_local, axis=0)  # [M, kk]
+            out = local + ops.convert_to_tensor(x_start[by][:, None].astype(np.int64), dtype=local.dtype)
+            if kk < k:  # fewer candidates than k: repeat the last one
+                out = ops.concatenate([out] + [out[:, -1:]] * (k - kk), axis=1)
+            return out
+        except (TypeError, ValueError, NotImplementedError, RuntimeError):
+            pass
+
+    if not isinstance(N, int) and y is x and batch_x is not None and batch_y is batch_x:
+        # compiled knn_graph with a point count known only at run time (TensorFlow): per example,
+        # since the chunks below would all be the full [N, N] distance matrix
+        return _knn_graph_indices_traced(x, k, ops.cast(batch_x, "int32"), cosine, exclude_self)
+
+    # Compiled: masked distances for chunks of query rows ([chunk, N] memory at a time)
+    bx = ops.zeros((N,), "int32") if batch_x is None else ops.cast(batch_x, "int32")
+    by = ops.zeros((M,), "int32") if batch_y is None else ops.cast(batch_y, "int32")
+    chunk = M if not isinstance(M, int) or not isinstance(N, int) else max(1, min(M, (1 << 24) // max(N, 1)))
+    outs = []
+    for start in range(0, M, chunk) if isinstance(M, int) else [0]:
+        stop = start + chunk if isinstance(M, int) else None
+        dist = _pairwise(y[start:stop], x, cosine)
+        mask = ops.expand_dims(by[start:stop], 1) != ops.expand_dims(bx, 0)
+        if exclude_self:
+            rows = ops.arange(ops.shape(dist)[0], dtype="int32") + start
+            mask = ops.logical_or(mask, ops.expand_dims(rows, 1) == ops.expand_dims(ops.arange(N, dtype="int32"), 0))
+        kk = min(k, N) if isinstance(N, int) else k
+        _, idx = ops.top_k(-ops.where(mask, 1e9, dist), k=kk, sorted=True)
+        if kk < k:  # fewer candidates than k: repeat the last one, as in the eager path
+            idx = ops.concatenate([idx] + [idx[:, -1:]] * (k - kk), axis=1)
+        outs.append(idx)
+    return outs[0] if len(outs) == 1 else ops.concatenate(outs, axis=0)
+
+
+def _knn_graph_indices_traced(x, k, batch, cosine, exclude_self):
+    """``_knn_indices(x, x, ...)`` for sizes only known at run time: distances within each example
+    (``[examples, max_points, max_points]``), like the eager path. Needs ``k`` <= points per example."""
+    from k3_node.layers.aggr.base import to_dense_batch
+    from k3_node.ops.segment import segment_sum
+
+    num_points = ops.shape(x)[0]
+    x_dense, mask = to_dense_batch(x, batch)
+    n = ops.shape(x_dense)[1]
+    dist = ops.where(ops.expand_dims(mask, 1), _pairwise(x_dense, x_dense, cosine), 1e9)
+    if exclude_self:
+        same = ops.expand_dims(ops.arange(n), 1) == ops.expand_dims(ops.arange(n), 0)
+        dist = ops.where(ops.expand_dims(same, 0), 1e9, dist)
+    kk = min(k, x_dense.shape[1]) if isinstance(x_dense.shape[1], int) else k
+    _, idx = ops.top_k(-dist, k=kk, sorted=True)  # local indices [examples, max_points, kk]
+    if kk < k:  # fewer candidates than k (e.g. Keras' shape inference): repeat the last one
+        idx = ops.concatenate([idx] + [idx[..., -1:]] * (k - kk), axis=-1)
+    counts = segment_sum(ops.ones_like(batch), batch, num_segments=num_points)
+    starts = ops.take(ops.cumsum(counts) - counts, batch, axis=0)  # first point of each point's example
+    local = ops.arange(num_points, dtype="int32") - starts
+    rows = ops.take(ops.reshape(idx, (-1, k)), batch * ops.cast(n, "int32") + local, axis=0)
+    return rows + ops.cast(ops.expand_dims(starts, 1), rows.dtype)
+
+
+def knn_graph(
+    x,
+    k: int,
+    batch: Optional[any] = None,
+    loop: bool = False,
+    flow: str = "source_to_target",
+    cosine: bool = False,
+    num_workers: int = 1,
+    batch_size: Optional[int] = None,
+):
+    r"""Computes graph edges to the nearest `k` points.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import knn_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 = knn_graph(pos, k=3)  # connect every point to its 3 nearest neighbors
+        print(tuple(edge_index.shape))  # (2, 30)
+        ```
+    """
+    assert flow in ["source_to_target", "target_to_source"]
+    x = ops.convert_to_tensor(x)
+    if len(ops.shape(x)) == 1:
+        x = ops.expand_dims(x, axis=-1)
+
+    N = ops.shape(x)[0]
+    col_indices = _knn_indices(x, x, k, batch, batch, cosine=cosine, exclude_self=not loop)
+    row = repeat(ops.arange(0, N, dtype="int64"), k)
+    col = ops.reshape(ops.cast(col_indices, "int64"), (-1,))
+
+    if flow == "source_to_target":
+        return ops.stack([col, row], axis=0)
+    else:
+        return ops.stack([row, col], axis=0)
+