Use one validate inputs function instead of 2

Author: coreyjadams

physicsnemo/nn/functional/neighbors/radius_search/_torch_impl.py

@@ -17,25 +17,7 @@
 
 import torch
 
-from .utils import format_returns
-
-
-def _validate_inputs(points: torch.Tensor, queries: torch.Tensor):
-    """Validate and normalize inputs to (B, N, 3) shape. Returns (points, queries, was_unbatched)."""
-    if points.ndim == 2 and queries.ndim == 2:
-        return points.unsqueeze(0), queries.unsqueeze(0), True
-    elif points.ndim == 3 and queries.ndim == 3:
-        if points.shape[0] != queries.shape[0]:
-            raise ValueError(
-                f"Batch dimensions must match: points has {points.shape[0]}, "
-                f"queries has {queries.shape[0]}"
-            )
-        return points, queries, False
-    else:
-        raise ValueError(
-            f"points and queries must be 2D (N, 3) or 3D (B, N, 3), "
-            f"got {points.ndim}D and {queries.ndim}D"
-        )
+from .utils import format_returns, validate_inputs
 
 
 def _radius_search_dynamic(
@@ -89,7 +71,7 @@ def radius_search_impl(
     This is a brute force implementation that is not memory efficient.
     """
 
-    points, queries, was_unbatched = _validate_inputs(points, queries)
+    points, queries, was_unbatched = validate_inputs(points, queries)
     B = points.shape[0]
 
     if max_points is None:

physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py

@@ -37,7 +37,7 @@
     scatter_add_batched,
     scatter_add_unlimited,
 )
-from .utils import format_returns
+from .utils import format_returns, validate_inputs
 
 wp.config.quiet = True
 
@@ -46,24 +46,6 @@
 BLOCK_DIM = 32
 
 
-def _validate_inputs(points: torch.Tensor, queries: torch.Tensor):
-    """Validate and normalize inputs. Returns (points, queries, was_unbatched)."""
-    if points.ndim == 2 and queries.ndim == 2:
-        return points.unsqueeze(0), queries.unsqueeze(0), True
-    elif points.ndim == 3 and queries.ndim == 3:
-        if points.shape[0] != queries.shape[0]:
-            raise ValueError(
-                f"Batch dimensions must match: points has {points.shape[0]}, "
-                f"queries has {queries.shape[0]}"
-            )
-        return points, queries, False
-    else:
-        raise ValueError(
-            f"points and queries must be 2D (N, 3) or 3D (B, N, 3), "
-            f"got {points.ndim}D and {queries.ndim}D"
-        )
-
-
 def count_neighbors(
     grid: wp.HashGrid,
     wp_points: wp.array(dtype=wp.vec3),
@@ -226,7 +208,7 @@ def radius_search_impl(
     if points.device != queries.device:
         raise ValueError("points and queries must be on the same device")
 
-    points, queries, was_unbatched = _validate_inputs(points, queries)
+    points, queries, was_unbatched = validate_inputs(points, queries)
     B = points.shape[0]
     N_queries = queries.shape[1]
 

physicsnemo/nn/functional/neighbors/radius_search/utils.py

@@ -17,6 +17,24 @@
 import torch
 
 
+def validate_inputs(points: torch.Tensor, queries: torch.Tensor):
+    """Validate and normalize inputs to (B, N, 3) shape. Returns (points, queries, was_unbatched)."""
+    if points.ndim == 2 and queries.ndim == 2:
+        return points.unsqueeze(0), queries.unsqueeze(0), True
+    elif points.ndim == 3 and queries.ndim == 3:
+        if points.shape[0] != queries.shape[0]:
+            raise ValueError(
+                f"Batch dimensions must match: points has {points.shape[0]}, "
+                f"queries has {queries.shape[0]}"
+            )
+        return points, queries, False
+    else:
+        raise ValueError(
+            f"points and queries must be 2D (N, 3) or 3D (B, N, 3), "
+            f"got {points.ndim}D and {queries.ndim}D"
+        )
+
+
 def format_returns(
     indices: torch.Tensor,
     points: torch.Tensor,