k3-node 1.0.0__py3-none-any.whl

k3_node/layers/pool/avg_pool.py

@@ -0,0 +1,165 @@
+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_sum
+from k3_node.ops.host import to_numpy
+
+
+def _avg_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
+    sum_x = segment_sum(x, cluster, num_segments=size)
+    ones = ops.ones_like(x)
+    count = segment_sum(ones, cluster, num_segments=size)
+    return sum_x / ops.maximum(count, 1.0)
+
+
+def avg_pool_x(
+    cluster,
+    x,
+    batch,
+    batch_size: Optional[int] = None,
+    size: Optional[int] = None,
+) -> Tuple[any, Optional[any]]:
+    r"""Average-pools node features according to the clustering defined in `cluster`.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import avg_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 = avg_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 _avg_pool_x(cluster, x, batch_size * size), None
+
+    cluster, perm = consecutive_cluster(cluster)
+    x = _avg_pool_x(cluster, x)
+    batch = pool_batch(perm, batch)
+    return x, batch
+
+
+def avg_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` using averaging.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import avg_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 = avg_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 = _avg_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, reduce="mean")
+
+        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 = _avg_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, reduce="mean")
+    pooled_batch = pool_batch(perm, batch) if batch is not None else None
+
+    return pooled_x, pooled_edge_index, pooled_batch
+
+
+def avg_pool_neighbor_x(
+    data,
+    edge_index=None,
+    flow: str = "source_to_target",
+):
+    r"""Average-pools neighboring node features.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import avg_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 = avg_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)
+    sum_x = segment_sum(x_src, col, num_segments=num_nodes)
+    ones = ops.ones_like(x_src)
+    count = segment_sum(ones, col, num_segments=num_nodes)
+    out_x = sum_x / ops.maximum(count, 1.0)
+    if is_data_obj:
+        data.x = out_x
+        return data
+    return out_x

k3_node/layers/pool/cluster_pool.py

@@ -0,0 +1,168 @@
+from typing import NamedTuple, Optional, Tuple
+from keras import layers, ops
+import numpy as np
+from k3_node.ops.segment import segment_sum
+
+
+class UnpoolInfo(NamedTuple):
+    edge_index: any
+    cluster: any
+    batch: any
+
+
+class ClusterPooling(layers.Layer):
+    r"""The cluster pooling operator from the `"Edge-Based Graph Component
+    Pooling" <https://arxiv.org/abs/2409.11856>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import ClusterPooling
+
+        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 = ClusterPooling(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: str = "tanh",
+        dropout: float = 0.0,
+        threshold: Optional[float] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        assert edge_score_method in ["tanh", "sigmoid", "log_softmax"]
+
+        if threshold is None:
+            threshold = 0.5 if edge_score_method == "sigmoid" else 0.0
+
+        self.in_channels = in_channels
+        self.edge_score_method = edge_score_method
+        self.dropout_rate = dropout
+        self.drop = layers.Dropout(dropout) if dropout > 0.0 else None
+        self.threshold = threshold
+
+        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()
+
+    def call(
+        self,
+        x,
+        edge_index,
+        batch,
+        training: bool = False,
+    ) -> Tuple[any, any, any, UnpoolInfo]:
+        r"""Forward pass."""
+        from k3_node.layers.conv.utils import eager_only_placeholder
+        if eager_only_placeholder("ClusterPooling", x, edge_index):
+            num_nodes = ops.shape(x)[0]
+            unpool_info = UnpoolInfo(edge_index, ops.arange(num_nodes, dtype="int32"), batch)
+            return x, edge_index, batch, unpool_info
+
+        edge_index_np = ops.convert_to_numpy(edge_index).astype(np.int64)
+        mask = edge_index_np[0] != edge_index_np[1]
+        edge_index_filtered = edge_index_np[:, mask]
+
+        row = ops.convert_to_tensor(edge_index_filtered[0], dtype="int32")
+        col = ops.convert_to_tensor(edge_index_filtered[1], dtype="int32")
+
+        edge_attr = ops.concatenate([ops.take(x, row, axis=0), ops.take(x, col, axis=0)], axis=-1)
+        edge_score = ops.reshape(self.lin(edge_attr), (-1,))
+        if self.drop is not None:
+            edge_score = self.drop(edge_score, training=training)
+
+        if self.edge_score_method == "tanh":
+            edge_score = ops.tanh(edge_score)
+        elif self.edge_score_method == "sigmoid":
+            edge_score = ops.sigmoid(edge_score)
+        else:
+            edge_score = ops.log_softmax(edge_score, axis=0)
+
+        edge_index_tensor = ops.convert_to_tensor(edge_index_filtered, dtype=edge_index.dtype)
+        return self._merge_edges(x, edge_index_tensor, batch, edge_score)
+
+    def _merge_edges(
+        self,
+        x,
+        edge_index,
+        batch,
+        edge_score,
+    ) -> Tuple[any, any, any, UnpoolInfo]:
+        from scipy.sparse import coo_matrix
+        from scipy.sparse.csgraph import connected_components
+
+        from k3_node.layers.conv.utils import host_callback
+
+        num_nodes = int(ops.shape(x)[0])
+        num_edges = int(ops.shape(edge_index)[1])
+        threshold = self.threshold
+
+        def contract(edge_index_np, edge_score_np, batch_np):
+            # Clusters are the weakly connected components of the edges scoring above the threshold.
+            edge_index_np = edge_index_np.astype(np.int64)
+            edge_contract = edge_index_np[:, edge_score_np > threshold]
+            if edge_contract.shape[1] > 0:
+                adj = coo_matrix(
+                    (np.ones(edge_contract.shape[1]), (edge_contract[0], edge_contract[1])),
+                    shape=(num_nodes, num_nodes),
+                )
+                _, cluster_np = connected_components(adj, directed=True, connection="weak")
+            else:
+                cluster_np = np.arange(num_nodes)
+            num_clusters = int(np.max(cluster_np)) + 1 if num_nodes > 0 else 0
+
+            # Nodes without any contracted edge keep their own features (unit diagonal score).
+            single = np.ones(num_nodes, dtype=bool)
+            single[edge_contract[0]] = False
+            single[edge_contract[1]] = False
+
+            # Coarsened edges between distinct clusters, in (row, col) order.
+            pairs = cluster_np[edge_index_np]
+            pairs = np.unique(pairs[:, pairs[0] != pairs[1]], axis=1)
+            edges_pad = np.zeros((2, num_edges), dtype=np.int64)
+            edges_pad[:, : pairs.shape[1]] = pairs
+
+            batch_pad = np.zeros(num_nodes, dtype=np.int64)
+            batch_pad[cluster_np] = batch_np
+            return cluster_np, num_clusters, single, edges_pad, pairs.shape[1], batch_pad
+
+        cluster, num_clusters, single, edges_pad, num_new_edges, batch_pad = host_callback(
+            contract,
+            [((num_nodes,), "int32"), ((), "int32"), ((num_nodes,), "float32"),
+             ((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)
+
+        # x_out = (S @ C)^T @ x, computed sparsely: every edge (row -> col) adds score * x[row] to
+        # cluster(col), and every unmatched node adds its own features to its cluster. The score
+        # enters as a tensor so gradients reach the scoring layer.
+        row = ops.cast(edge_index[0], "int32")
+        col = ops.cast(edge_index[1], "int32")
+        msgs = ops.expand_dims(ops.cast(edge_score, x.dtype), -1) * ops.take(x, row, axis=0)
+        x_out = segment_sum(msgs, ops.take(cluster, col, axis=0), num_segments=num_clusters)
+        x_out = x_out + segment_sum(
+            x * ops.expand_dims(ops.cast(single, x.dtype), -1), cluster, num_segments=num_clusters
+        )
+
+        edge_index_out = ops.cast(edges_pad[:, :num_new_edges], edge_index.dtype)
+        batch_out = ops.cast(batch_pad[:num_clusters], batch.dtype)
+
+        unpool_info = UnpoolInfo(edge_index, cluster, batch)
+        return x_out, edge_index_out, batch_out, unpool_info
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.in_channels})"

k3_node/layers/pool/connect/__init__.py

@@ -0,0 +1,10 @@
+from .base import Connect, ConnectOutput
+from .filter_edges import FilterEdges, filter_adj
+
+__all__ = [
+    "ConnectOutput",
+    "Connect",
+    "filter_adj",
+    "FilterEdges",
+]
+

k3_node/layers/pool/connect/base.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from typing import Optional
+from keras import layers, ops
+
+from ..select.base import SelectOutput
+
+
+@dataclass
+class ConnectOutput:
+    r"""The output of the :class:`Connect` method, which holds the coarsened
+    graph structure, and optional pooled edge features and batch vectors.
+
+    Args:
+        edge_index: The edge indices of the coarsened graph.
+        edge_attr: The pooled edge features of the coarsened graph. (default: None)
+        batch: The pooled batch vector of the coarsened graph. (default: None)
+    """
+    edge_index: any
+    edge_attr: Optional[any] = None
+    batch: Optional[any] = None
+
+    def __post_init__(self):
+        shape_edge = getattr(self.edge_index, "shape", None)
+        if shape_edge is not None:
+            if len(shape_edge) != 2:
+                raise ValueError(
+                    f"Expected 'edge_index' to be two-dimensional "
+                    f"(got {len(shape_edge)} dimensions)"
+                )
+            if shape_edge[0] is not None and shape_edge[0] != 2:
+                raise ValueError(
+                    f"Expected 'edge_index' to have size '2' in the first dimension "
+                    f"(got '{shape_edge[0]}')"
+                )
+        if self.edge_attr is not None:
+            shape_attr = getattr(self.edge_attr, "shape", None)
+            if (
+                shape_edge is not None
+                and shape_attr is not None
+                and len(shape_edge) == 2
+                and len(shape_attr) >= 1
+                and shape_edge[1] is not None
+                and shape_attr[0] is not None
+                and shape_attr[0] != shape_edge[1]
+            ):
+                raise ValueError(
+                    f"Expected 'edge_index' and 'edge_attr' to hold the same number "
+                    f"of edges (got {shape_edge[1]} and {shape_attr[0]} edges)"
+                )
+
+
+import keras
+
+if keras.config.backend() == "jax":
+    try:
+        import jax
+        from jax.tree_util import register_pytree_node
+
+        register_pytree_node(
+            ConnectOutput,
+            lambda c: ((c.edge_index, c.edge_attr, c.batch), ()),
+            lambda aux, children: ConnectOutput(children[0], children[1], children[2]),
+        )
+    except Exception:
+        pass
+
+
+
+class Connect(layers.Layer):
+    r"""An abstract base class for implementing custom edge connection
+    operators as described in the `"Understanding Pooling in Graph Neural
+    Networks" <https://arxiv.org/abs/1905.05178>`_ paper.
+    """
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        pass
+
+    def __call__(self, *args, **kwargs):
+        if len(args) > 0 and isinstance(args[0], SelectOutput):
+            return self.call(*args, **kwargs)
+        return super().__call__(*args, **kwargs)
+
+    def call(
+        self,
+        select_output: SelectOutput,
+        edge_index,
+        edge_attr: Optional[any] = None,
+        batch: Optional[any] = None,
+    ) -> ConnectOutput:
+        raise NotImplementedError
+
+    @staticmethod
+    def get_pooled_batch(
+        select_output: SelectOutput,
+        batch: Optional[any],
+    ) -> Optional[any]:
+        r"""Returns the batch vector of the coarsened graph."""
+        if batch is None:
+            return None
+        return ops.take(batch, select_output.node_index, axis=0)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'

k3_node/layers/pool/connect/filter_edges.py

@@ -0,0 +1,113 @@
+from typing import Optional, Tuple
+from keras import ops
+import numpy as np
+
+from .base import Connect, ConnectOutput
+from ..select.base import SelectOutput
+
+
+from k3_node.layers.conv.utils import is_tracing
+from k3_node.ops.creation import full
+
+
+def filter_adj(
+    edge_index,
+    edge_attr: Optional[any] = None,
+    node_index=None,
+    cluster_index: Optional[any] = None,
+    num_nodes: Optional[int] = None,
+) -> Tuple[any, Optional[any]]:
+    r"""Filters out edges if their incident nodes are not in any cluster.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import filter_adj
+
+        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
+
+        kept_nodes = np.array([0, 2, 4, 6, 8])
+        edge_index_new, _ = filter_adj(edge_index, node_index=kept_nodes)  # edges between kept nodes, relabeled
+        print(edge_index_new.shape[0])  # 2
+        ```
+    """
+    if is_tracing(edge_index) or (node_index is not None and is_tracing(node_index)):
+        return edge_index, edge_attr
+
+    if node_index is None:
+        return edge_index, edge_attr
+
+    edge_index = ops.cast(edge_index, "int32")
+    node_index = ops.cast(node_index, "int32")
+    if cluster_index is None:
+        cluster_index = ops.arange(ops.shape(node_index)[0], dtype="int32")
+    else:
+        cluster_index = ops.cast(cluster_index, "int32")
+
+    if num_nodes is None:
+        num_nodes = ops.max(node_index) + 1 if ops.shape(node_index)[0] > 0 else 0
+        if ops.shape(edge_index)[1] > 0:
+            num_nodes = ops.maximum(num_nodes, ops.max(edge_index) + 1)
+    try:
+        num_nodes = int(num_nodes)
+    except (TypeError, ValueError):
+        pass
+
+    mapping = full((num_nodes,), -1, dtype="int32")
+    mapping = ops.scatter_update(mapping, ops.expand_dims(node_index, -1), cluster_index)
+
+    row = ops.take(mapping, edge_index[0], axis=0)
+    col = ops.take(mapping, edge_index[1], axis=0)
+    mask = (row >= 0) & (col >= 0)
+    valid_idx = ops.where(mask)
+    if isinstance(valid_idx, (tuple, list)):
+        valid_idx = valid_idx[0]
+    valid_idx = ops.reshape(valid_idx, (-1,))
+
+    new_edge_index = ops.stack(
+        [ops.take(row, valid_idx, axis=0), ops.take(col, valid_idx, axis=0)], axis=0
+    )
+
+    new_edge_attr = None
+    if edge_attr is not None:
+        new_edge_attr = ops.take(edge_attr, valid_idx, axis=0)
+
+    return new_edge_index, new_edge_attr
+
+
+class FilterEdges(Connect):
+    r"""Filters out edges if their incident nodes are not in any cluster.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import FilterEdges, SelectTopK
+
+        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
+
+        select_output = SelectTopK(in_channels=8, ratio=0.5)(x, batch)
+        connect = FilterEdges()
+        out = connect(select_output, edge_index, batch=batch)  # keep edges between kept nodes
+        print(out.edge_index.shape[0])  # 2
+        ```
+    """
+    def call(
+        self,
+        select_output: SelectOutput,
+        edge_index,
+        edge_attr: Optional[any] = None,
+        batch: Optional[any] = None,
+    ) -> ConnectOutput:
+        new_edge_index, new_edge_attr = filter_adj(
+            edge_index,
+            edge_attr,
+            select_output.node_index,
+            select_output.cluster_index,
+            num_nodes=select_output.num_nodes,
+        )
+        new_batch = self.get_pooled_batch(select_output, batch)
+        return ConnectOutput(new_edge_index, new_edge_attr, new_batch)
+

k3_node/layers/pool/consecutive.py

@@ -0,0 +1,30 @@
+from typing import Tuple
+from keras import ops
+import numpy as np
+from k3_node.ops.host import to_numpy
+
+
+def consecutive_cluster(src) -> Tuple[any, any]:
+    r"""Maps elements in `src` to consecutive integers starting from 0,
+    returning the mapped indices and a permutation array of representative indices.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import consecutive_cluster
+
+        cluster = np.array([4, 4, 9, 2, 9])
+        new_cluster, perm = consecutive_cluster(cluster)  # relabel cluster ids to 0..num_clusters-1
+        print(tuple(new_cluster.shape), tuple(perm.shape))  # (5,) (3,)
+        ```
+    """
+    src_np = to_numpy(src)
+    unique, inv = np.unique(src_np, return_inverse=True)
+    perm = np.empty(len(unique), dtype=inv.dtype)
+    arange = np.arange(len(inv), dtype=inv.dtype)
+    perm[inv] = arange
+
+    inv_tensor = ops.convert_to_tensor(inv, dtype=src.dtype)
+    perm_tensor = ops.convert_to_tensor(perm, dtype=src.dtype)
+    return inv_tensor, perm_tensor
+

k3_node/layers/pool/decimation.py

@@ -0,0 +1,48 @@
+from typing import Tuple, Union
+from keras import ops
+import numpy as np
+from k3_node.ops.host import _in_shape_inference, to_numpy
+
+
+def decimation_indices(
+    ptr,
+    decimation_factor: Union[int, float],
+) -> Tuple[any, any]:
+    r"""Gets indices which downsample each point cloud by a decimation factor.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.layers import decimation_indices
+
+        ptr = np.array([0, 4, 10])  # two graphs with 4 and 6 nodes
+        index, new_ptr = decimation_indices(ptr, decimation_factor=2)  # keep every 2nd node per graph
+        print(tuple(index.shape), tuple(new_ptr.shape))  # (5,) (3,)
+        ```
+    """
+    if decimation_factor < 1:
+        raise ValueError(
+            f"The argument `decimation_factor` should be higher than (or "
+            f"equal to) 1 for downsampling. (got {decimation_factor})"
+        )
+
+    ptr_np = to_numpy(ptr)
+    batch_size = len(ptr_np) - 1
+    count = ptr_np[1:] - ptr_np[:-1]
+    if _in_shape_inference():  # `ptr` is placeholder zeros: keep one (valid) node per graph
+        count = np.maximum(count, 1)
+    decim_count = np.maximum(count // int(decimation_factor), 1)
+
+    decim_indices_list = []
+    for i in range(batch_size):
+        perm = np.random.permutation(count[i])[:decim_count[i]]
+        decim_indices_list.append(ptr_np[i] + perm)
+
+    decim_indices = np.concatenate(decim_indices_list, axis=0)
+    decim_ptr = np.concatenate([[0], np.cumsum(decim_count)])
+
+    return (
+        ops.convert_to_tensor(decim_indices, dtype=ptr.dtype),
+        ops.convert_to_tensor(decim_ptr, dtype=ptr.dtype),
+    )
+