mps-pointops 0.3.0__py3-none-any.whl

mps_pointops/__init__.py

@@ -0,0 +1,7 @@
+"""Point cloud ops for PyTorch on Apple Silicon (MPS)."""
+
+from . import flat, reference
+from .ops import ball_query, furthest_point_sample, knn
+
+__all__ = ["ball_query", "flat", "furthest_point_sample", "knn", "reference"]
+__version__ = "0.3.0"

mps_pointops/_ball_query_mps.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: MIT
+"""Deterministic radius-neighbor search on Apple GPUs.
+
+The Metal source is original to this project. The public contract follows the
+mathematical definition of radius search and the documented first-K convention
+used by PyTorch3D; no third-party implementation is copied here.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from importlib.resources import files
+import math
+import struct
+from typing import NamedTuple
+
+import torch
+
+
+_MIN_SUPPORTED_RADIUS_F32 = 2.0**-112
+
+
+class BallQueryResult(NamedTuple):
+    """Squared distances and indices, both shaped ``(B, Q, K)``."""
+
+    distances: torch.Tensor
+    indices: torch.Tensor
+
+
+def _require_compile_shader() -> None:
+    if not hasattr(torch.mps, "compile_shader"):
+        raise RuntimeError(
+            f"mps_pointops needs PyTorch 2.7 or later (torch.mps.compile_shader); found {torch.__version__}"
+        )
+
+
+@lru_cache(maxsize=1)
+def _shader_library():
+    _require_compile_shader()
+    source = files(__package__).joinpath("kernels", "ball_query.metal").read_text()
+    return torch.mps.compile_shader(source)
+
+
+def _lengths_or_full(
+    lengths: torch.Tensor | None, *, batch: int, count: int, device: torch.device, name: str
+) -> torch.Tensor:
+    if lengths is None:
+        return torch.full((batch,), count, device=device, dtype=torch.int64)
+    if not isinstance(lengths, torch.Tensor):
+        raise TypeError(f"{name} must be a torch.Tensor or None")
+    if lengths.shape != (batch,) or lengths.dtype != torch.int64 or lengths.device != device:
+        raise ValueError(f"{name} must have shape ({batch},), int64 dtype and device {device}")
+    if torch.any((lengths < 0) | (lengths > count)).item():
+        raise ValueError(f"{name} values must lie in [0, {count}]")
+    return lengths.contiguous()
+
+
+def _checked_radius_and_k(radius: float, k: int) -> tuple[float, float]:
+    """Validate ``radius`` and ``k`` for every ball query path, on any device.
+
+    Returns ``(radius_f32, radius_squared)``: the radius rounded to float32 and
+    fl32(fl32(radius) * fl32(radius)).
+    """
+    if not isinstance(k, int) or isinstance(k, bool) or k < 0 or k > (1 << 63) - 1:
+        raise ValueError("k must be a non-negative integer")
+    if not isinstance(radius, (int, float)) or isinstance(radius, bool):
+        raise ValueError("radius must be a finite non-negative number")
+    try:
+        radius = float(radius)
+    except OverflowError as exc:
+        raise ValueError("radius must be a finite non-negative number") from exc
+    if not math.isfinite(radius) or radius < 0:
+        raise ValueError("radius must be a finite non-negative number")
+
+    try:
+        radius_f32 = struct.unpack("f", struct.pack("f", radius))[0]
+    except OverflowError as exc:
+        raise ValueError("radius must fit in float32") from exc
+    if not math.isfinite(radius_f32) or (radius > 0 and radius_f32 < _MIN_SUPPORTED_RADIUS_F32):
+        raise ValueError("positive radius must round to a normal float32 value at least 2**-112")
+
+    # This is fl32(fl32(radius) * fl32(radius)), as in PyTorch3D's float
+    # radius2 = radius * radius. Two binary32 operands multiply exactly in
+    # binary64 (at most 48 significant bits); struct.pack rounds once to f32.
+    radius_squared = float(radius_f32) * float(radius_f32)
+    try:
+        radius_squared = struct.unpack("f", struct.pack("f", radius_squared))[0]
+    except OverflowError as exc:
+        radius_squared = math.inf
+    return radius_f32, radius_squared
+
+
+class _BallQuery(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        queries: torch.Tensor,
+        points: torch.Tensor,
+        query_lengths: torch.Tensor,
+        point_lengths: torch.Tensor,
+        radius_sq: float,
+        radius_f32: float,
+        k: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, query_count, _ = queries.shape
+        point_count = points.shape[1]
+        indices = torch.empty((batch, query_count, k), device=queries.device, dtype=torch.int64)
+        distances = torch.empty((batch, query_count, k), device=queries.device, dtype=torch.float32)
+
+        if indices.numel():
+            kernel = (
+                _shader_library().ball_query_f32
+                if queries.dtype == torch.float32
+                else _shader_library().ball_query_f16
+            )
+            kernel(
+                queries,
+                points,
+                query_lengths,
+                point_lengths,
+                indices,
+                distances,
+                batch,
+                query_count,
+                point_count,
+                k,
+                radius_sq,
+                radius_f32,
+                threads=[batch * query_count, 1, 1],
+                group_size=[256, 1, 1],
+            )
+
+        ctx.save_for_backward(queries, points, indices)
+        ctx.mark_non_differentiable(indices)
+        return distances, indices
+
+    @staticmethod
+    def backward(ctx, grad_distances: torch.Tensor | None, grad_indices: None):
+        if grad_distances is None:
+            return None, None, None, None, None, None, None
+
+        queries, points, indices = ctx.saved_tensors
+        batch, query_count, k = indices.shape
+        if not indices.numel() or points.shape[1] == 0:
+            return torch.zeros_like(queries), torch.zeros_like(points), None, None, None, None, None
+
+        valid = indices >= 0
+        safe_indices = indices.clamp_min(0)
+        gathered = points.float().gather(
+            1, safe_indices.reshape(batch, -1, 1).expand(-1, -1, 3)
+        ).reshape(batch, query_count, k, 3)
+        delta = queries.float().unsqueeze(2) - gathered
+        # Padded slots may gather a non-finite point 0. Mask before multiplying:
+        # IEEE-754 makes 0 * NaN equal NaN, not zero.
+        delta = torch.where(valid.unsqueeze(-1), delta, torch.zeros_like(delta))
+        valid_grad = torch.where(valid, grad_distances, torch.zeros_like(grad_distances))
+        coeff = (2.0 * valid_grad).unsqueeze(-1)
+        contributions = coeff * delta
+
+        grad_queries = (
+            contributions.sum(dim=2).to(queries.dtype) if ctx.needs_input_grad[0] else None
+        )
+        grad_points = None
+        if ctx.needs_input_grad[1]:
+            grad_points = torch.zeros_like(points, dtype=torch.float32)
+            grad_points.scatter_add_(
+                1,
+                safe_indices.reshape(batch, -1, 1).expand(-1, -1, 3),
+                -contributions.reshape(batch, -1, 3),
+            )
+            grad_points = grad_points.to(points.dtype)
+        return grad_queries, grad_points, None, None, None, None, None
+
+
+def ball_query(
+    queries: torch.Tensor,
+    points: torch.Tensor,
+    *,
+    radius: float,
+    k: int,
+    query_lengths: torch.Tensor | None = None,
+    point_lengths: torch.Tensor | None = None,
+) -> BallQueryResult:
+    """Find the first ``k`` points strictly inside each radius.
+
+    ``queries`` and ``points`` are float32 or float16 MPS tensors with shapes ``(B,Q,3)``
+    and ``(B,P,3)``. The output is ordered by input point index, not distance.
+    Missing neighbors use index ``-1`` and squared distance ``0``. Under the
+    supported Safe math mode, non-finite coordinates never match. Gradients of
+    valid squared distances propagate to both coordinate tensors; selection
+    indices are non-differentiable. ``radius`` is a Python number, not a
+    differentiable Tensor.
+
+    Coordinates remain on MPS. Supplying lengths checks their bounds with a
+    scalar MPS-to-CPU synchronization. Non-contiguous coordinate tensors are
+    copied to contiguous MPS storage before dispatch. The radius is rounded
+    to float32 first and must then be zero or at least 2**-112. Smaller
+    positive values are rejected because Metal may flush coordinate differences
+    that matter to the radius decision. Boundary decisions use float32 arithmetic.
+    """
+    if not isinstance(queries, torch.Tensor) or not isinstance(points, torch.Tensor):
+        raise TypeError("queries and points must be torch.Tensor values")
+    if queries.ndim != 3 or points.ndim != 3 or queries.shape[-1] != 3 or points.shape[-1] != 3:
+        raise ValueError("queries and points must have shapes (B, Q, 3) and (B, P, 3)")
+    if queries.shape[0] != points.shape[0]:
+        raise ValueError("queries and points must have the same batch size")
+    if queries.device.type != "mps" or points.device != queries.device:
+        raise ValueError("queries and points must be on the same MPS device")
+    if queries.dtype not in (torch.float32, torch.float16) or points.dtype != queries.dtype:
+        raise TypeError("queries and points must have the same float32 or float16 dtype")
+    radius_f32, radius_squared = _checked_radius_and_k(radius, k)
+
+    batch, query_count, _ = queries.shape
+    point_count = points.shape[1]
+    if batch * query_count > (1 << 32) - 1:
+        raise ValueError("B * Q exceeds the Metal 1-D dispatch limit")
+    if batch * query_count * 3 > (1 << 63) - 1 or batch * point_count * 3 > (1 << 63) - 1:
+        raise ValueError("coordinate element count exceeds int64 indexing")
+    if batch * query_count * k > (1 << 63) - 1:
+        raise ValueError("output element count exceeds int64 indexing")
+
+    if not torch.backends.mps.is_available():
+        raise RuntimeError("PyTorch MPS is unavailable on this host")
+    query_lengths = _lengths_or_full(
+        query_lengths, batch=batch, count=query_count, device=queries.device, name="query_lengths"
+    )
+    point_lengths = _lengths_or_full(
+        point_lengths, batch=batch, count=point_count, device=queries.device, name="point_lengths"
+    )
+
+    distances, indices = _BallQuery.apply(
+        queries.contiguous(),
+        points.contiguous(),
+        query_lengths,
+        point_lengths,
+        radius_squared,
+        radius_f32,
+        k,
+    )
+    return BallQueryResult(distances=distances, indices=indices)

mps_pointops/_flat_fps_mps.py

@@ -0,0 +1,96 @@
+"""Farthest point sampling for flat, variable-length point clouds on MPS."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from importlib.resources import files
+
+import torch
+from torch import Tensor
+
+from ._ball_query_mps import _require_compile_shader
+
+
+_THREADS = 1024
+_MAX_LOCAL_POINTS = (1 << 32) - 1
+
+
+@lru_cache(maxsize=1)
+def _shader_library():
+    _require_compile_shader()
+    width_shader = torch.mps.compile_shader(
+        "kernel void width(device long* out, uint w [[threads_per_simdgroup]]) { out[0] = w; }"
+    )
+    width = torch.empty(1, dtype=torch.long, device="mps")
+    width_shader.width(width, threads=1, group_size=1)
+    actual_width = int(width.item())
+    if actual_width != 32:
+        raise RuntimeError(
+            f"mps_pointops kernels need 32-wide simdgroups, this GPU uses {actual_width}"
+        )
+    source = files(__package__).joinpath("kernels", "fps_flat.metal").read_text()
+    return torch.mps.compile_shader(source)
+
+
+def fps_flat(x: Tensor, ptr: Tensor, out_ptr: Tensor, starts: Tensor) -> Tensor:
+    """Sample flat ``x`` by segment and return global int64 point indices.
+
+    ``ptr[b:b+2]`` bounds input cloud ``b``; ``out_ptr[b:b+2]`` bounds its
+    output. ``starts[b]`` is a *local* index in that cloud. A cloud with zero
+    requested samples may be empty. As in the existing FPS kernel, requests
+    longer than a cloud repeat its local index zero after exhaustion.
+    """
+    if not all(isinstance(t, Tensor) for t in (x, ptr, out_ptr, starts)):
+        raise TypeError("x, ptr, out_ptr, and starts must be torch.Tensor values")
+    if x.ndim != 2 or x.shape[1] != 3:
+        raise ValueError(f"x must have shape (Total_N, 3), got {tuple(x.shape)}")
+    if ptr.ndim != 1 or ptr.numel() == 0:
+        raise ValueError("ptr must be a nonempty one-dimensional tensor")
+    batch = ptr.numel() - 1
+    if out_ptr.shape != (batch + 1,) or starts.shape != (batch,):
+        raise ValueError("out_ptr must have shape (B+1,) and starts must have shape (B,)")
+    if x.device.type != "mps" or any(t.device != x.device for t in (ptr, out_ptr, starts)):
+        raise ValueError("x, ptr, out_ptr, and starts must be on the same MPS device")
+    if x.dtype != torch.float32:
+        raise TypeError(f"x must be float32 on MPS, got {x.dtype}")
+    if any(t.dtype != torch.int64 for t in (ptr, out_ptr, starts)):
+        raise TypeError("ptr, out_ptr, and starts must be int64")
+    if batch > (1 << 32) - 1:
+        raise ValueError("B exceeds the Metal grid limit")
+
+    x = x.contiguous()
+    ptr = ptr.contiguous()
+    out_ptr = out_ptr.contiguous()
+    starts = starts.contiguous()
+    lengths = ptr[1:] - ptr[:-1]
+    counts = out_ptr[1:] - out_ptr[:-1]
+    invalid = (
+        (ptr[0] != 0)
+        | (ptr[-1] != x.shape[0])
+        | (out_ptr[0] != 0)
+        | torch.any(lengths < 0)
+        | torch.any(lengths > _MAX_LOCAL_POINTS)
+        | torch.any(counts < 0)
+        | torch.any((counts > 0) & ((lengths == 0) | (starts < 0) | (starts >= lengths)))
+    )
+    # Output shape is data-dependent. Transfer only its total size and the
+    # validation flag; coordinates and segment offsets remain on the GPU.
+    output_size, is_invalid = torch.stack((out_ptr[-1], invalid.to(torch.int64))).cpu().tolist()
+    if is_invalid:
+        raise ValueError("invalid ptr, out_ptr, or starts for the input point clouds")
+
+    out = torch.empty(output_size, dtype=torch.int64, device=x.device)
+    if batch == 0 or output_size == 0:
+        return out
+    min_d2 = torch.empty(x.shape[0], dtype=torch.float32, device=x.device)
+    _shader_library().fps_flat(
+        x,
+        ptr,
+        out_ptr,
+        starts,
+        min_d2,
+        out,
+        threads=(_THREADS, batch),
+        group_size=(_THREADS, 1),
+    )
+    return out

mps_pointops/_flat_search_mps.py

@@ -0,0 +1,126 @@
+"""Metal searches over sorted, flattened point-cloud batches.
+
+These private kernels produce rectangular global-index arrays with ``-1``
+padding. The public flat API removes padding to build torch_cluster-style
+``[query, reference]`` edge indices.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+from importlib import resources
+import math
+import struct
+
+import torch
+from torch import Tensor
+
+from ._ball_query_mps import _checked_radius_and_k, _require_compile_shader
+from .ops import _check_simd_width
+
+
+_MAX_K = 256
+_QUERIES_PER_GROUP = 8
+_UINT_MAX = (1 << 32) - 1
+
+
+@cache
+def _library():
+    _require_compile_shader()
+    source = resources.files(__package__).joinpath("kernels", "flat_search.metal").read_text()
+    return torch.mps.compile_shader(source)
+
+
+def _check_inputs(x: Tensor, y: Tensor, ptr_x: Tensor, ptr_y: Tensor) -> tuple[int, int]:
+    if x.ndim != 2 or y.ndim != 2 or x.shape[1] != 3 or y.shape[1] != 3:
+        raise ValueError("x and y must have shapes (N, 3) and (M, 3)")
+    if x.device.type != "mps" or y.device != x.device:
+        raise ValueError("x and y must be on the same MPS device")
+    if ptr_x.ndim != 1 or ptr_y.ndim != 1 or ptr_x.shape != ptr_y.shape or not ptr_x.numel():
+        raise ValueError("ptr_x and ptr_y must have the same nonempty 1-D shape")
+    for name, ptr, count in (("ptr_x", ptr_x, x.shape[0]), ("ptr_y", ptr_y, y.shape[0])):
+        if ptr.dtype != torch.int64 or ptr.device != x.device:
+            raise ValueError(f"{name} must be an int64 tensor on {x.device}")
+        if int(ptr[0].item()) != 0 or int(ptr[-1].item()) != count:
+            raise ValueError(f"{name} must start at 0 and end at {count}")
+        lengths = ptr[1:] - ptr[:-1]
+        if bool(torch.any((lengths < 0) | (lengths > _UINT_MAX)).item()):
+            raise ValueError(f"{name} must be nondecreasing with per-batch size <= 2**32-1")
+    if y.shape[0] > _UINT_MAX:
+        raise ValueError("flat query count exceeds the Metal 1-D dispatch limit")
+    return int(ptr_x.numel() - 1), int(y.shape[0])
+
+
+def _torch_cluster_radius_sq(r: float) -> float:
+    """The threshold torch_cluster compares against: fl32(r * r), r * r in double.
+
+    torch_cluster takes ``r`` as a double and passes ``r * r`` to a float32
+    kernel. PyTorch3D-style Ball Query instead squares the float32-rounded
+    radius; the two can differ by one float32 ULP (for example at r = 0.21).
+    """
+    try:
+        return struct.unpack("f", struct.pack("f", r * r))[0]
+    except OverflowError:
+        return math.inf
+
+
+def knn_indices(x: Tensor, y: Tensor, ptr_x: Tensor, ptr_y: Tensor, k: int) -> Tensor:
+    """Return global x indices, sorted by (squared distance, x index)."""
+    batch_count, query_count = _check_inputs(x, y, ptr_x, ptr_y)
+    if x.dtype != torch.float32 or y.dtype != torch.float32:
+        raise TypeError("flat Metal kNN requires float32 x and y")
+    if not isinstance(k, int) or isinstance(k, bool) or not 0 <= k <= _MAX_K:
+        raise ValueError(f"k must be an integer in [0, {_MAX_K}] for the flat Metal kernel")
+    if query_count * k > (1 << 63) - 1:
+        raise ValueError("kNN output size exceeds int64 indexing")
+    out = torch.empty((query_count, k), dtype=torch.int64, device=x.device)
+    if query_count == 0 or k == 0:
+        return out
+    _check_simd_width()
+    group = 32 * _QUERIES_PER_GROUP
+    groups = (query_count + _QUERIES_PER_GROUP - 1) // _QUERIES_PER_GROUP
+    _library().flat_knn_indices(
+        y.contiguous(), x.contiguous(), ptr_y.contiguous(), ptr_x.contiguous(), out,
+        query_count, batch_count, k,
+        threads=[groups * group, 1, 1], group_size=[group, 1, 1],
+    )
+    return out
+
+
+def radius_indices(
+    x: Tensor, y: Tensor, ptr_x: Tensor, ptr_y: Tensor, r: float,
+    max_num_neighbors: int, ignore_same_index: bool = False,
+) -> Tensor:
+    """Return first valid global x indices per query, with ``-1`` padding.
+
+    When requested, equal global x/y index numbers are excluded before the
+    neighbor cap is applied, matching pyg-lib's radius operator.
+    """
+    batch_count, query_count = _check_inputs(x, y, ptr_x, ptr_y)
+    if x.dtype not in (torch.float32, torch.float16) or y.dtype != x.dtype:
+        raise TypeError("flat Metal radius requires matching float32 or float16 x and y")
+    radius_f32, _ = _checked_radius_and_k(r, max_num_neighbors)
+    radius_sq = _torch_cluster_radius_sq(float(r))
+    if query_count * max_num_neighbors > (1 << 63) - 1:
+        raise ValueError("radius output size exceeds int64 indexing")
+    out = torch.empty((query_count, max_num_neighbors), dtype=torch.int64, device=x.device)
+    if query_count == 0 or max_num_neighbors == 0:
+        return out
+    if radius_f32 == 0.0:
+        out.fill_(-1)
+        return out
+    _check_simd_width()
+    kernel = (
+        _library().flat_radius_indices_f32
+        if x.dtype == torch.float32
+        else _library().flat_radius_indices_f16
+    )
+    group = 32 * _QUERIES_PER_GROUP
+    groups = (query_count + _QUERIES_PER_GROUP - 1) // _QUERIES_PER_GROUP
+    kernel(
+        y.contiguous(), x.contiguous(), ptr_y.contiguous(), ptr_x.contiguous(), out,
+        query_count, batch_count, max_num_neighbors, radius_sq, radius_f32,
+        int(ignore_same_index),
+        threads=[groups * group, 1, 1], group_size=[group, 1, 1],
+    )
+    return out

mps_pointops/compat.py

@@ -0,0 +1,166 @@
+"""Optional stand-ins for ``pointnet2_ops``, ``knn_cuda`` and ``torch_cluster``.
+
+Call ``install()`` before importing code that uses them::
+
+    import mps_pointops.compat
+    mps_pointops.compat.install()
+
+    from pointnet2_ops import pointnet2_utils  # served by mps_pointops
+    from knn_cuda import KNN
+    from torch_cluster import fps, knn, radius
+
+``install()`` registers ``pointnet2_ops``, ``pointnet2_ops.pointnet2_utils``,
+``knn_cuda`` and ``torch_cluster`` in ``sys.modules``. A name that is already
+importable, such as the real package on a CUDA machine, is left alone unless
+``force=True``. Call it before importing the packages that use these names.
+
+Covered: ``furthest_point_sample``, ``gather_operation``,
+``grouping_operation`` and ``ball_query`` from ``pointnet2_utils``, and
+``KNN`` from ``knn_cuda``. Differences from the CUDA versions:
+
+- Near ties can resolve differently. The kernels round each squared
+  distance without FMA and break ties by the smaller index; the CUDA
+  kernels use their own reduction order.
+- ``ball_query`` uses the Metal kernel for MPS inputs and pads in the
+  ``pointnet2_ops`` convention.
+- The ``torch_cluster`` shim exposes ``fps``, ``knn``, ``radius`` and their
+  same-set graph wrappers for flat three-dimensional point coordinates.
+  Other names needed for PyG 2.7 package import raise ``NotImplementedError``
+  when called. It does not register PyG's separate ``torch.ops.pyg`` operators.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from importlib.machinery import ModuleSpec
+import sys
+import types
+
+import torch
+from torch import Tensor
+
+from . import flat, ops
+
+
+# ---------------------------------------------------------------- pointnet2_ops
+
+
+def furthest_point_sample(xyz: Tensor, npoint: int) -> Tensor:
+    """Like ``pointnet2_ops.pointnet2_utils.furthest_point_sample``.
+
+    (B, N, 3) float32 -> (B, npoint) int32. Starts at index 0 and never picks
+    points with x^2 + y^2 + z^2 <= 1e-3, as pointnet2_ops does.
+    """
+    return ops.furthest_point_sample(xyz, npoint, skip_near_origin=True).int()
+
+
+def gather_operation(features: Tensor, idx: Tensor) -> Tensor:
+    """Like ``pointnet2_utils.gather_operation``: (B, C, N), (B, npoint) -> (B, C, npoint)."""
+    B, C, _ = features.shape
+    return features.gather(2, idx.long().unsqueeze(1).expand(B, C, -1))
+
+
+def grouping_operation(features: Tensor, idx: Tensor) -> Tensor:
+    """Like ``pointnet2_utils.grouping_operation``: (B, C, N), (B, npoint, nsample) -> (B, C, npoint, nsample)."""
+    B, C, _ = features.shape
+    _, npoint, nsample = idx.shape
+    flat = idx.long().reshape(B, 1, npoint * nsample).expand(B, C, -1)
+    return features.gather(2, flat).reshape(B, C, npoint, nsample)
+
+
+def ball_query(radius: float, nsample: int, xyz: Tensor, new_xyz: Tensor) -> Tensor:
+    """Like ``pointnet2_utils.ball_query``: (B, N, 3), (B, npoint, 3) -> (B, npoint, nsample) int32.
+
+    Takes the first ``nsample`` points in input order with squared distance
+    ``< radius**2``. Empty slots repeat the first neighbor, and a query with no
+    neighbor gets all zeros, as in pointnet2_ops.
+    """
+    _, idx = ops.ball_query(new_xyz, xyz, radius, nsample)
+    first = idx[..., :1].clamp(min=0)
+    return torch.where(idx >= 0, idx, first).int()
+
+
+# ---------------------------------------------------------------- knn_cuda
+
+
+class KNN(torch.nn.Module):
+    """Like ``knn_cuda.KNN``.
+
+    With ``transpose_mode=True``, ``ref`` is (B, N, 3) and ``query`` is
+    (B, M, 3), and ``dist`` and ``idx`` are (B, M, k). Otherwise ``ref`` is
+    (B, 3, N), ``query`` is (B, 3, M), and the outputs are (B, k, M). ``dist``
+    is Euclidean. No gradients flow, as in knn_cuda.
+    """
+
+    def __init__(self, k: int, transpose_mode: bool = False):
+        super().__init__()
+        self.k = k
+        self.transpose_mode = transpose_mode
+
+    def forward(self, ref: Tensor, query: Tensor) -> tuple[Tensor, Tensor]:
+        if not self.transpose_mode:
+            ref, query = ref.transpose(1, 2), query.transpose(1, 2)
+        with torch.no_grad():
+            dist, idx = ops.knn(query.float(), ref.float(), self.k)
+        if not self.transpose_mode:
+            dist, idx = dist.transpose(1, 2), idx.transpose(1, 2)
+        return dist, idx
+
+
+# ---------------------------------------------------------------- install
+
+
+def _module(name: str, **attrs) -> types.ModuleType:
+    module = types.ModuleType(name)
+    module.__dict__.update(attrs)
+    module.__spec__ = ModuleSpec(name, loader=None, is_package="__path__" in attrs)
+    return module
+
+
+def _unsupported_torch_cluster(name: str):
+    def unsupported(*args, **kwargs):
+        raise NotImplementedError(
+            f"torch_cluster.{name} is outside the mps_pointops point-cloud subset"
+        )
+
+    unsupported.__name__ = name
+    return unsupported
+
+
+def install(force: bool = False) -> list[str]:
+    """Register the stand-in modules. Returns the names that were installed."""
+    pointnet2_utils = _module(
+        "pointnet2_ops.pointnet2_utils",
+        furthest_point_sample=furthest_point_sample,
+        gather_operation=gather_operation,
+        grouping_operation=grouping_operation,
+        ball_query=ball_query,
+    )
+    packages = {
+        "pointnet2_ops": {
+            "pointnet2_ops": _module("pointnet2_ops", pointnet2_utils=pointnet2_utils, __path__=[]),
+            "pointnet2_ops.pointnet2_utils": pointnet2_utils,
+        },
+        "knn_cuda": {"knn_cuda": _module("knn_cuda", KNN=KNN)},
+        "torch_cluster": {
+            "torch_cluster": _module(
+                "torch_cluster",
+                fps=flat.fps,
+                knn=flat.knn,
+                radius=flat.radius,
+                knn_graph=flat.knn_graph,
+                radius_graph=flat.radius_graph,
+                grid_cluster=_unsupported_torch_cluster("grid_cluster"),
+                graclus_cluster=_unsupported_torch_cluster("graclus_cluster"),
+                random_walk=_unsupported_torch_cluster("random_walk"),
+                nearest=_unsupported_torch_cluster("nearest"),
+            )
+        },
+    }
+    installed = []
+    for top, modules in packages.items():
+        if not force and (top in sys.modules or importlib.util.find_spec(top) is not None):
+            continue
+        sys.modules.update(modules)
+        installed.extend(modules)
+    return installed