k3-node 1.0.0__py3-none-any.whl

k3_node/layers/pool/test_pool.py

@@ -0,0 +1,456 @@
+import keras
+import pytest
+import numpy as np
+from keras import ops
+
+from k3_node.layers.pool import (
+    ASAPooling,
+    ApproxL2KNNIndex,
+    ApproxMIPSKNNIndex,
+    ClusterPooling,
+    EdgePooling,
+    FilterEdges,
+    KNNIndex,
+    L2KNNIndex,
+    LEConv,
+    MIPSKNNIndex,
+    MemPooling,
+    PANPooling,
+    SAGPooling,
+    SelectTopK,
+    TopKPooling,
+    approx_knn,
+    approx_knn_graph,
+    avg_pool,
+    avg_pool_neighbor_x,
+    avg_pool_x,
+    consecutive_cluster,
+    decimation_indices,
+    filter_adj,
+    fps,
+    global_add_pool,
+    global_max_pool,
+    global_mean_pool,
+    graclus,
+    knn,
+    knn_graph,
+    max_pool,
+    max_pool_neighbor_x,
+    max_pool_x,
+    nearest,
+    pool_batch,
+    pool_edge,
+    pool_pos,
+    radius,
+    radius_graph,
+    topk,
+    voxel_grid,
+)
+from k3_node.layers.pool.select.base import SelectOutput
+
+
+def test_global_pooling():
+    x = ops.convert_to_tensor([
+        [1.0, 2.0],
+        [3.0, 4.0],
+        [5.0, 6.0],
+        [7.0, 8.0],
+    ], dtype="float32")
+    batch = ops.convert_to_tensor([0, 0, 1, 1], dtype="int32")
+
+    add = global_add_pool(x, batch)
+    assert np.allclose(ops.convert_to_numpy(add), [[4.0, 6.0], [12.0, 14.0]])
+
+    mean = global_mean_pool(x, batch)
+    assert np.allclose(ops.convert_to_numpy(mean), [[2.0, 3.0], [6.0, 7.0]])
+
+    max_p = global_max_pool(x, batch)
+    assert np.allclose(ops.convert_to_numpy(max_p), [[3.0, 4.0], [7.0, 8.0]])
+
+
+def test_consecutive_cluster():
+    src = ops.convert_to_tensor([1, 4, 4, 7, 7, 7, 10], dtype="int32")
+    out, perm = consecutive_cluster(src)
+    assert np.array_equal(ops.convert_to_numpy(out), [0, 1, 1, 2, 2, 2, 3])
+    assert ops.shape(perm)[0] == 4
+
+
+def test_filter_adj_and_layer():
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 0],
+        [1, 2, 3, 0, 2],
+    ], dtype="int32")
+    edge_attr = ops.convert_to_tensor([1.0, 2.0, 3.0, 4.0, 5.0], dtype="float32")
+    node_index = ops.convert_to_tensor([0, 1, 2], dtype="int32")
+
+    new_edge_index, new_edge_attr = filter_adj(edge_index, edge_attr, node_index=node_index)
+    assert ops.shape(new_edge_index)[0] == 2
+    assert ops.shape(new_edge_index)[1] == 3  # (0,1), (1,2), (0,2)
+
+    # FilterEdges
+    layer = FilterEdges()
+    select_out = SelectOutput(
+        node_index=node_index,
+        num_nodes=4,
+        cluster_index=ops.arange(3, dtype="int32"),
+        num_clusters=3,
+        weight=ops.ones((3,), dtype="float32"),
+    )
+    batch = ops.zeros((4,), dtype="int32")
+    connect_out = layer(select_out, edge_index, edge_attr, batch)
+    assert ops.shape(connect_out.edge_index)[1] == 3
+    assert ops.shape(connect_out.batch)[0] == 3
+
+
+def test_topk_and_select_topk():
+    x = ops.convert_to_tensor([2.0, 1.0, 5.0, 3.0, 8.0, 4.0], dtype="float32")
+    batch = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int32")
+
+    # Ratio test
+    perm = topk(x, ratio=0.5, batch=batch)
+    perm_np = ops.convert_to_numpy(perm)
+    assert 2 in perm_np[:2]  # from first graph (5.0, 2.0)
+    assert 4 in perm_np[1:]  # from second graph (8.0, 4.0)
+
+    # Min score test
+    perm_min = topk(x, ratio=None, batch=batch, min_score=3.5)
+    perm_min_np = ops.convert_to_numpy(perm_min)
+    assert 2 in perm_min_np  # 5.0
+    assert 4 in perm_min_np  # 8.0
+    assert 5 in perm_min_np  # 4.0
+
+    # SelectTopK layer
+    layer = SelectTopK(in_channels=4, ratio=0.5)
+    feat = ops.ones((6, 4), dtype="float32")
+    out = layer(feat, batch)
+    assert isinstance(out, SelectOutput)
+    assert ops.shape(out.node_index)[0] == 4  # 2 per graph
+
+
+def test_topk_pooling():
+    in_channels = 8
+    layer = TopKPooling(in_channels, ratio=0.5)
+
+    x = ops.ones((6, in_channels), dtype="float32")
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 4, 5],
+        [1, 2, 0, 4, 5, 3],
+    ], dtype="int32")
+    batch = ops.convert_to_tensor([0, 0, 0, 1, 1, 1], dtype="int32")
+
+    out_x, out_edge_index, out_edge_attr, out_batch, perm, score = layer(
+        x, edge_index, batch=batch
+    )
+    assert ops.shape(out_x)[0] == 4  # 2 nodes per graph
+    assert ops.shape(out_x)[1] == in_channels
+    assert ops.shape(perm)[0] == 4
+    assert ops.shape(out_batch)[0] == 4
+
+
+def test_sag_pooling():
+    in_channels = 8
+    layer = SAGPooling(in_channels, ratio=0.5)
+
+    x = ops.ones((4, in_channels), dtype="float32")
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3, 0],
+        [1, 2, 3, 0, 2],
+    ], dtype="int32")
+
+    out_x, out_edge_index, out_edge_attr, out_batch, perm, score = layer(x, edge_index)
+    assert ops.shape(out_x)[0] == 2
+    assert ops.shape(out_x)[1] == in_channels
+    assert ops.shape(perm)[0] == 2
+
+
+def test_edge_pooling():
+    in_channels = 8
+    layer = EdgePooling(in_channels)
+
+    x = ops.ones((4, in_channels), dtype="float32")
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [1, 2, 3, 0],
+    ], dtype="int32")
+    batch = ops.zeros((4,), dtype="int32")
+
+    new_x, new_edge_index, new_batch, unpool_info = layer(x, edge_index, batch)
+    assert ops.shape(new_x)[0] <= 4
+    assert ops.shape(new_x)[1] == in_channels
+
+    # Unpool
+    unpooled_x, _, _ = layer.unpool(new_x, unpool_info)
+    assert ops.shape(unpooled_x)[0] == 4
+    assert ops.shape(unpooled_x)[1] == in_channels
+
+
+def test_cluster_pooling():
+    in_channels = 8
+    layer = ClusterPooling(in_channels)
+
+    x = ops.ones((4, in_channels), dtype="float32")
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [1, 2, 3, 0],
+    ], dtype="int32")
+    batch = ops.zeros((4,), dtype="int32")
+
+    new_x, new_edge_index, new_batch, unpool_info = layer(x, edge_index, batch)
+    assert unpool_info.edge_index is not None
+    assert unpool_info.cluster is not None
+    assert unpool_info.batch is not None
+
+
+def test_mem_pooling():
+    in_channels, out_channels, heads, num_clusters = 8, 16, 2, 4
+    layer = MemPooling(in_channels, out_channels, heads=heads, num_clusters=num_clusters)
+
+    x = ops.ones((2, 5, in_channels), dtype="float32")
+    out_x, S = layer(x)
+    assert ops.shape(out_x) == (2, num_clusters, out_channels)
+    assert ops.shape(S) == (2, 5, num_clusters)
+
+    loss = MemPooling.kl_loss(S)
+    assert ops.shape(loss) == ()
+
+
+def test_asap_pooling():
+    in_channels = 8
+    leconv = LEConv(in_channels, in_channels)
+    x = ops.ones((4, in_channels), dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 0]], dtype="int32")
+    le_out = leconv(x, edge_index)
+    assert ops.shape(le_out) == (4, in_channels)
+
+    asap = ASAPooling(in_channels, ratio=0.5)
+    out_x, out_edge_index, _, _, perm = asap(x, edge_index)
+    assert ops.shape(out_x)[0] == 2
+    assert ops.shape(out_x)[1] == in_channels
+
+
+def test_pan_pooling():
+    in_channels = 8
+    pan = PANPooling(in_channels, ratio=0.5)
+    x = ops.ones((4, in_channels), dtype="float32")
+    M = ops.eye(4, dtype="float32")
+
+    out_x, out_edge_index, _, _, perm, score = pan(x, M)
+    assert ops.shape(out_x)[0] == 2
+    assert ops.shape(out_x)[1] == in_channels
+    assert ops.shape(out_edge_index)[0] == 2
+
+
+def test_max_and_avg_pool():
+    cluster = ops.convert_to_tensor([0, 0, 1, 1], dtype="int32")
+    x = ops.convert_to_tensor([
+        [1.0, 5.0],
+        [3.0, 2.0],
+        [4.0, 8.0],
+        [6.0, 7.0],
+    ], dtype="float32")
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [1, 2, 3, 0],
+    ], dtype="int32")
+    batch = ops.zeros((4,), dtype="int32")
+
+    # max_pool_x & avg_pool_x
+    mx, mb = max_pool_x(cluster, x, batch)
+    assert np.allclose(ops.convert_to_numpy(mx), [[3.0, 5.0], [6.0, 8.0]])
+    ax, ab = avg_pool_x(cluster, x, batch)
+    assert np.allclose(ops.convert_to_numpy(ax), [[2.0, 3.5], [5.0, 7.5]])
+
+    # neighbor_x with raw tensors
+    mn_x = max_pool_neighbor_x(x, edge_index=edge_index)
+    assert ops.shape(mn_x) == (4, 2)
+    an_x = avg_pool_neighbor_x(x, edge_index=edge_index)
+    assert ops.shape(an_x) == (4, 2)
+
+    # max_pool & avg_pool (on graph)
+    m_x, m_edge, m_batch = max_pool(cluster, x, edge_index, batch=batch)
+    assert ops.shape(m_x)[0] == 2
+    a_x, a_edge, a_batch = avg_pool(cluster, x, edge_index, batch=batch)
+    assert ops.shape(a_x)[0] == 2
+
+
+def test_pool_helpers():
+    cluster = ops.convert_to_tensor([0, 0, 1, 1], dtype="int32")
+    perm = ops.convert_to_tensor([0, 2], dtype="int32")
+    batch = ops.convert_to_tensor([0, 0, 0, 0], dtype="int32")
+    pos = ops.convert_to_tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0], [3.0, 3.0]], dtype="float32")
+    edge_index = ops.convert_to_tensor([[0, 1, 2], [1, 2, 3]], dtype="int32")
+    edge_attr = ops.convert_to_tensor([[1.0], [2.0], [3.0]], dtype="float32")
+
+    p_edge, p_attr = pool_edge(cluster, edge_index, edge_attr)
+    assert ops.shape(p_edge)[0] == 2
+    assert p_attr is not None
+
+    p_b = pool_batch(perm, batch)
+    assert ops.shape(p_b)[0] == 2
+
+    p_p = pool_pos(cluster, pos)
+    assert ops.shape(p_p)[0] == 2
+
+
+def test_voxel_grid():
+    pos = ops.convert_to_tensor([
+        [0.0, 0.0],
+        [0.5, 0.5],
+        [1.2, 1.2],
+        [2.1, 2.1],
+    ], dtype="float32")
+    size = ops.convert_to_tensor([1.0, 1.0], dtype="float32")
+
+    c = voxel_grid(pos, size)
+    c_np = ops.convert_to_numpy(c)
+    assert c_np[0] == c_np[1]
+    assert c_np[0] != c_np[2]
+    assert c_np[2] != c_np[3]
+
+
+def test_graclus():
+    edge_index = ops.convert_to_tensor([
+        [0, 1, 2, 3],
+        [1, 2, 3, 0],
+    ], dtype="int32")
+    c = graclus(edge_index, num_nodes=4)
+    c_np = ops.convert_to_numpy(c)
+    assert len(np.unique(c_np)) == 2
+
+
+def test_decimation():
+    ptr = ops.convert_to_tensor([0, 4, 10], dtype="int32")
+    dec_idx, dec_ptr = decimation_indices(ptr, decimation_factor=2)
+    assert ops.shape(dec_idx)[0] == 5  # 2 from first, 3 from second
+    assert ops.shape(dec_ptr)[0] == 3
+
+
+def test_knn_and_indices():
+    x = ops.convert_to_tensor([
+        [0.0, 0.0],
+        [0.1, 0.1],
+        [1.0, 1.0],
+        [1.1, 1.1],
+    ], dtype="float32")
+    y = ops.convert_to_tensor([
+        [0.0, 0.0],
+        [1.0, 1.0],
+    ], dtype="float32")
+
+    # KNNIndex / L2KNNIndex
+    idx_layer = L2KNNIndex(x)
+    out = idx_layer.search(y, k=2)
+    assert ops.shape(out.score) == (2, 2)
+    assert ops.shape(out.index) == (2, 2)
+
+    # MIPSKNNIndex
+    mips_layer = MIPSKNNIndex(x)
+    mips_out = mips_layer.search(y, k=2)
+    assert ops.shape(mips_out.score) == (2, 2)
+    assert ops.shape(mips_out.index) == (2, 2)
+
+    # knn function
+    res = knn(x, y, k=2)
+    assert ops.shape(res) == (2, 4)
+
+    # knn_graph
+    g = knn_graph(x, k=2)
+    assert ops.shape(g)[0] == 2
+
+    # approx knn
+    approx_res = approx_knn(x, y, k=2)
+    assert ops.shape(approx_res) == (2, 4)
+    approx_g = approx_knn_graph(x, k=2)
+    assert ops.shape(approx_g)[0] == 2
+
+
+def test_point_cloud():
+    pos = ops.convert_to_tensor([
+        [0.0, 0.0],
+        [0.1, 0.0],
+        [10.0, 10.0],
+        [10.1, 10.0],
+    ], dtype="float32")
+
+    # fps
+    fps_idx = fps(pos, ratio=0.5)
+    assert ops.shape(fps_idx)[0] == 2
+
+    # radius
+    rad_edge = radius(pos, pos, r=1.0)
+    assert ops.shape(rad_edge)[0] == 2
+    assert ops.shape(rad_edge)[1] >= 4  # at least self loops + close pairs
+
+    # radius_graph
+    rad_g = radius_graph(pos, r=1.0)
+    assert ops.shape(rad_g)[0] == 2
+
+    # nearest
+    near = nearest(pos, pos)
+    assert ops.shape(near)[0] == 4
+
+
+def test_global_pool_with_unsorted_batch():
+    # Pooling layers (e.g. EdgePooling) can return batch vectors that are not sorted by graph.
+    import numpy as np
+    from keras import ops
+    from k3_node.layers import global_add_pool, global_mean_pool
+
+    x = np.arange(8, dtype="float32").reshape(4, 2)
+    batch = np.array([1, 0, 1, 0], dtype="int32")
+    np.testing.assert_allclose(ops.convert_to_numpy(global_add_pool(x, batch)), [[8, 10], [4, 6]])
+    np.testing.assert_allclose(ops.convert_to_numpy(global_mean_pool(x, batch)), [[4, 5], [2, 3]])
+
+
+@pytest.mark.skipif(keras.backend.backend() != "jax", reason="jax.jit specific")
+def test_global_pool_under_jax_jit_requires_size():
+    import jax
+    import numpy as np
+    from k3_node.layers import global_add_pool
+
+    x = np.arange(8, dtype="float32").reshape(4, 2)
+    batch = np.array([0, 0, 1, 1], dtype="int32")
+    with pytest.raises(ValueError, match="size="):
+        jax.jit(lambda x, b: global_add_pool(x, b))(x, batch)
+    out = jax.jit(lambda x, b: global_add_pool(x, b, size=2))(x, batch)
+    np.testing.assert_allclose(np.asarray(out), [[2, 4], [10, 12]])
+
+
+@pytest.mark.skipif(keras.backend.backend() not in ("tensorflow", "jax"), reason="needs a compiling backend")
+@pytest.mark.parametrize("layer_name", ["EdgePooling", "ClusterPooling"])
+def test_host_side_pooling_refuses_compiled_execution(layer_name):
+    # These layers pick clusters on the host; inside tf.function / jax.jit they used to return
+    # their input unpooled without any error.
+    import numpy as np
+    import k3_node.layers as L
+
+    layer = getattr(L, layer_name)(4)
+    x = np.random.randn(6, 4).astype("float32")
+    edge_index = np.array([[0, 1, 2, 3, 4], [1, 2, 3, 4, 5]], dtype="int32")
+    batch = np.zeros(6, dtype="int32")
+    layer(x, edge_index, batch)  # eager works
+
+    if keras.backend.backend() == "jax":
+        import jax
+
+        compiled = jax.jit(lambda x: layer(x, edge_index, batch)[0])
+    else:
+        import tensorflow as tf
+
+        compiled = tf.function(lambda x: layer(x, edge_index, batch)[0])
+    with pytest.raises(Exception, match="run_eagerly"):
+        compiled(x)
+
+
+def test_mem_pooling_respects_batch():
+    import numpy as np
+    from k3_node.layers import MemPooling
+
+    x = np.random.rand(7, 4).astype("float32")
+    batch = np.array([0, 0, 0, 1, 1, 2, 2])  # three graphs
+    pool = MemPooling(4, 8, heads=2, num_clusters=3)
+    out, S = pool(x, batch)
+    assert tuple(out.shape) == (3, 3, 8)  # one pooled set per graph
+    # Pooling a graph alone gives the same result as pooling it within the batch
+    alone, _ = pool(x[3:5], np.zeros(2, dtype="int32"))
+    np.testing.assert_allclose(ops.convert_to_numpy(out)[1], ops.convert_to_numpy(alone)[0], rtol=1e-5, atol=1e-6)

k3_node/layers/pool/topk_pool.py

@@ -0,0 +1,103 @@
+from typing import Callable, Optional, Tuple, Union
+from keras import layers, ops
+
+from .connect.filter_edges import FilterEdges
+from .select.topk import SelectTopK
+
+
+class TopKPooling(layers.Layer):
+    r""":math:`\mathrm{top}_k` pooling operator from the `"Graph U-Nets"
+    <https://arxiv.org/abs/1905.05178>`_, `"Towards Sparse
+    Hierarchical Graph Classifiers" <https://arxiv.org/abs/1811.01287>`_
+    and `"Understanding Attention and Generalization in Graph Neural
+    Networks" <https://arxiv.org/abs/1905.02850>`_ papers.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import TopKPooling
+
+        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 = TopKPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
+        out = layer(x, edge_index, batch=batch)
+        x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
+        print(tuple(x_pool.shape))  # (6, 8)
+        ```
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: Union[int, 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.select = SelectTopK(in_channels, ratio, min_score, nonlinearity)
+        self.connect = FilterEdges()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.select.reset_parameters()
+
+    def build(self, input_shape=None):
+        if not self.select.built:
+            self.select.build(input_shape)
+        if hasattr(self.connect, "built") and not self.connect.built:
+            self.connect.build(None)
+        self.built = True
+
+    def call(
+        self,
+        x,
+        edge_index,
+        edge_attr: Optional[any] = None,
+        batch: Optional[any] = None,
+        attn: Optional[any] = None,
+    ) -> Tuple[any, any, Optional[any], Optional[any], any, any]:
+        r"""Forward pass."""
+        if batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int32")
+
+        attn_input = x if attn is None else attn
+        select_out = self.select(attn_input, batch)
+
+        perm = select_out.node_index
+        score = select_out.weight
+
+        x_pooled = ops.take(x, perm, axis=0) * ops.expand_dims(score, axis=-1)
+        if self.multiplier != 1.0:
+            x_pooled = x_pooled * self.multiplier
+
+        connect_out = self.connect.call(select_out, edge_index, edge_attr, batch)
+
+        return (
+            x_pooled,
+            connect_out.edge_index,
+            connect_out.edge_attr,
+            connect_out.batch,
+            perm,
+            score,
+        )
+
+    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/voxel_grid.py

@@ -0,0 +1,70 @@
+from typing import List, Optional, Union
+from keras import ops
+import numpy as np
+from k3_node.ops.host import to_numpy
+
+
+def voxel_grid(
+    pos,
+    size: Union[float, List[float], any],
+    batch: Optional[any] = None,
+    start: Optional[Union[float, List[float], any]] = None,
+    end: Optional[Union[float, List[float], any]] = None,
+):
+    r"""Voxel grid pooling that clusters points within the same voxel.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import voxel_grid
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        pos = np.random.rand(10, 3).astype("float32")  # 3D positions
+
+        cluster = voxel_grid(pos, size=0.5)  # voxel id of each point
+        print(tuple(cluster.shape))  # (10,)
+        ```
+    """
+    pos_np = to_numpy(pos)
+    if pos_np.ndim == 1:
+        pos_np = pos_np[:, None]
+    dim = pos_np.shape[1]
+
+    if isinstance(size, (int, float)):
+        size_np = np.full(dim, float(size))
+    else:
+        size_np = np.array(size, dtype=np.float64)
+
+    if start is None:
+        start_np = np.min(pos_np, axis=0)
+    elif isinstance(start, (int, float)):
+        start_np = np.full(dim, float(start))
+    else:
+        start_np = np.array(start, dtype=np.float64)
+
+    if end is None:
+        end_np = np.max(pos_np, axis=0)
+    elif isinstance(end, (int, float)):
+        end_np = np.full(dim, float(end))
+    else:
+        end_np = np.array(end, dtype=np.float64)
+
+    # Number of bins per dimension
+    num_bins = np.floor((end_np - start_np) / size_np).astype(np.int64) + 1
+
+    # Discretized coordinates
+    c = np.floor((pos_np - start_np) / size_np).astype(np.int64)
+
+    # Compute flat index
+    cluster_np = np.zeros(len(pos_np), dtype=np.int64)
+    multiplier = 1
+    for d in range(dim):
+        cluster_np += c[:, d] * multiplier
+        multiplier *= num_bins[d]
+
+    if batch is not None:
+        batch_np = to_numpy(batch).astype(np.int64)
+        cluster_np += batch_np * multiplier
+
+    return ops.convert_to_tensor(cluster_np, dtype="int64")
+

k3_node/layers/unpool/__init__.py

@@ -0,0 +1,9 @@
+r"""Unpooling package."""
+
+from .knn_interpolate import knn_interpolate
+
+__all__ = [
+    "knn_interpolate",
+]
+
+classes = __all__

k3_node/layers/unpool/knn_interpolate.py

@@ -0,0 +1,57 @@
+from keras import ops
+
+from k3_node.layers.conv.utils import scatter
+from k3_node.layers.pool.knn import knn
+
+
+def knn_interpolate(x, pos_x, pos_y, batch_x=None, batch_y=None, k: int = 3, num_workers: int = 1):
+    r"""The k-NN interpolation from the `"PointNet++: Deep Hierarchical
+    Feature Learning on Point Sets in a Metric Space"
+    <https://arxiv.org/abs/1706.02413>`_ paper.
+
+    For each point :math:`y` with position :math:`\mathbf{p}(y)`, its
+    interpolated features :math:`\mathbf{f}(y)` are given by
+
+    .. math::
+        \mathbf{f}(y) = \frac{\sum_{i=1}^k w(x_i) \mathbf{f}(x_i)}{\sum_{i=1}^k
+        w(x_i)} \textrm{, where } w(x_i) = \frac{1}{d(\mathbf{p}(y),
+        \mathbf{p}(x_i))^2}
+
+    and :math:`\{ x_1, \ldots, x_k \}` denoting the :math:`k` nearest points
+    to :math:`y`.
+
+    Args:
+        x: Node feature matrix :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        pos_x: Node position matrix :math:`\in \mathbb{R}^{N \times d}`.
+        pos_y: Upsampled node position matrix :math:`\in \mathbb{R}^{M \times d}`.
+        batch_x: Batch vector assigning each node from :math:`\mathbf{X}` to
+            a specific example. (default: :obj:`None`)
+        batch_y: Batch vector assigning each node from :math:`\mathbf{Y}` to
+            a specific example. (default: :obj:`None`)
+        k (int, optional): Number of neighbors. (default: :obj:`3`)
+        num_workers (int, optional): Unused, kept for API compatibility.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import knn_interpolate
+
+        x = np.random.rand(6, 8).astype("float32")  # features of 6 coarse points
+        pos_x = np.random.rand(6, 3).astype("float32")
+        pos_y = np.random.rand(10, 3).astype("float32")  # 10 fine points to interpolate onto
+        out = knn_interpolate(x, pos_x, pos_y, k=3)
+        print(tuple(out.shape))  # (10, 8)
+        ```
+    """
+    assign_index = knn(pos_x, pos_y, k, batch_x=batch_x, batch_y=batch_y, num_workers=num_workers)
+    y_idx, x_idx = assign_index[0], assign_index[1]
+
+    diff = ops.take(pos_x, x_idx, axis=0) - ops.take(pos_y, y_idx, axis=0)
+    squared_distance = ops.sum(diff * diff, axis=-1, keepdims=True)
+    weights = 1.0 / ops.maximum(squared_distance, 1e-16)
+
+    num_y = ops.shape(pos_y)[0]
+    y = scatter(ops.take(x, x_idx, axis=0) * weights, y_idx, dim=0, dim_size=num_y, reduce="sum")
+    y = y / scatter(weights, y_idx, dim=0, dim_size=num_y, reduce="sum")
+
+    return y