fix(torch): densify ragged knn_search output in PointTransformer's knn_batch

[dripston]

Problem

Training PointTransformer on a point cloud (or a batch region) with fewer candidate points than k crashes:

RuntimeError: shape '[-1, 16]' is invalid for input of size 345

Root cause

open3d.ml.torch.ops.knn_search's own docstring states it "supports returning less than k neighbors if there are less than k points" and that its output format is "compatible with the radius_search and fixed_radius_search ops" — i.e. neighbors_index is ragged: its total length is not guaranteed to be a multiple of k, and neighbors_row_splits marks each query's actual sub-range.

knn_batch in ml3d/torch/models/point_transformer.py ignored this and called ans.neighbors_index.reshape(-1, k) directly. When every query has exactly k neighbors this happens to work; the moment any query has fewer, the reshape either raises (this issue) or, if the ragged total happens to still divide evenly by k, silently misaligns which neighbors belong to which query.

Fix

ml3d/torch/models/kpconv.py's batch_neighbors already solves the identical ragged-output problem for FixedRadiusSearch using open3d.ml.torch.ops.ragged_to_dense (values, row_splits, out_col_size, default_value). I applied the same op here, using neighbors_row_splits to densify to (num_queries, k).

One difference from kpconv.py's usage: knn_batch's result is used by queryandgroup to index directly into points/feat (points[idx.view(-1).long(), :]), so I can't reuse ragged_to_dense's out-of-range default_value (points.shape[0]) as the final padding — that would just trade the reshape crash for an IndexError. After densifying with that sentinel, padded slots are replaced with that query's own first (real) neighbor, which keeps every returned index valid and contributes a harmless, already-nearby duplicate to the fixed-size attention window instead of an arbitrary or out-of-range one. The one further edge case — a query with zero neighbors at all — falls back to point index 0. Distances get the same treatment, defaulting to 0.0 in that same all-empty case so no inf reaches the downstream attention/softmax.

Testing

• Added test_pointtransformer_knn_batch_ragged_neighbors in tests/test_models_torch.py, using a 5-point cloud with k=16 (guaranteeing every query gets fewer than k neighbors, matching the issue's actual trigger condition) and asserting the returned idx/dist have the correct (n_points, k) shape with every index in [0, n_points).
• Verified the padding logic in isolation with a pure-Python/torch reimplementation of ragged_to_dense (this environment has no compiled Open3D ML ops), covering: a full row (unaffected), a short row (padded with its own first neighbor), and a fully empty row (falls back to index 0) — all three produced correctly shaped, in-range output.
• Confirmed via knn_search's own docstring (extracted from the open3d==0.20.0 wheel) that its ragged output shares neighbors_row_splits with FixedRadiusSearch's result, validating that the kpconv.py densification pattern applies unchanged.
• python3 -m py_compile passes on both changed files.
• Could not run the added test or the existing test_pointtransformer_torch here — this environment doesn't have the compiled open3d/open3d.ml.torch package with GPU ops. Please run pytest tests/test_models_torch.py -k pointtransformer in CI/a configured dev environment to confirm.

Fixes #683

🤖 Generated with Claude Code