diff --git a/physicsnemo/mesh/README.md b/physicsnemo/mesh/README.md
index 601e18753c..5b555eaa3b 100644
--- a/physicsnemo/mesh/README.md
+++ b/physicsnemo/mesh/README.md
@@ -422,7 +422,7 @@ neighbors of mesh elements (i.e., based on the mesh connectivity,as opposed to
 Note that these use an efficient sparse (`indices`, `offsets`) encoding of the
 adjacency relationships, which is used internally for all computations. (See the
 dedicated
-[`physicsnemo.mesh.neighbors._adjacency.py`](physicsnemo/mesh/neighbors/_adjacency.py)
+[`physicsnemo.mesh.neighbors._adjacency.py`](./neighbors/_adjacency.py)
 module.) You can convert these to a typical ragged list-of-lists representation
 with `.to_list()`, which is useful for debugging or interoperability, at the
 cost of performance:
@@ -540,35 +540,35 @@ Key design decisions enable these principles:
 
 ## Documentation & Resources
 
-- **Examples**: See [`examples/`](examples/) directory for runnable demonstrations
-- **Tests**: See [`test/`](test/) directory for comprehensive test suite showing usage
+- **Examples**: See [`examples/`](../../examples/) directory for runnable demonstrations
+- **Tests**: See [`test/`](../../test/mesh/) directory for comprehensive test suite showing usage
   patterns
-- **Source**: Explore [`physicsnemo/mesh/`](physicsnemo/mesh/) for implementation details
+- **Source**: Explore [`physicsnemo/mesh/`](./) for implementation details
 
 **Module Organization:**
 
-- [`physicsnemo.mesh.calculus`](physicsnemo/mesh/calculus/) - Discrete differential
+- [`physicsnemo.mesh.calculus`](./calculus/) - Discrete differential
   operators
-- [`physicsnemo.mesh.curvature`](physicsnemo/mesh/curvature/) - Gaussian and mean
+- [`physicsnemo.mesh.curvature`](./curvature/) - Gaussian and mean
   curvature
-- [`physicsnemo.mesh.subdivision`](physicsnemo/mesh/subdivision/) - Mesh refinement
+- [`physicsnemo.mesh.subdivision`](./subdivision/) - Mesh refinement
   schemes
-- [`physicsnemo.mesh.boundaries`](physicsnemo/mesh/boundaries/) - Boundary detection
+- [`physicsnemo.mesh.boundaries`](./boundaries/) - Boundary detection
   and facet extraction
-- [`physicsnemo.mesh.neighbors`](physicsnemo/mesh/neighbors/) - Adjacency computations
-- [`physicsnemo.mesh.spatial`](physicsnemo/mesh/spatial/) - BVH and spatial queries
-- [`physicsnemo.mesh.sampling`](physicsnemo/mesh/sampling/) - Point sampling and
+- [`physicsnemo.mesh.neighbors`](./neighbors/) - Adjacency computations
+- [`physicsnemo.mesh.spatial`](./spatial/) - BVH and spatial queries
+- [`physicsnemo.mesh.sampling`](./sampling/) - Point sampling and
   interpolation
-- [`physicsnemo.mesh.transformations`](physicsnemo/mesh/transformations/) - Geometric
+- [`physicsnemo.mesh.transformations`](./transformations/) - Geometric
   operations
-- [`physicsnemo.mesh.repair`](physicsnemo/mesh/repair/) - Mesh cleaning and topology
+- [`physicsnemo.mesh.repair`](./repair/) - Mesh cleaning and topology
   repair
-- [`physicsnemo.mesh.validation`](physicsnemo/mesh/validation/) - Quality metrics
+- [`physicsnemo.mesh.validation`](./validation/) - Quality metrics
   and statistics
-- [`physicsnemo.mesh.visualization`](physicsnemo/mesh/visualization/) - Matplotlib
+- [`physicsnemo.mesh.visualization`](./visualization/) - Matplotlib
   and PyVista backends
-- [`physicsnemo.mesh.io`](physicsnemo/mesh/io/) - PyVista import/export
-- [`physicsnemo.mesh.examples`](physicsnemo/mesh/examples/) - Example mesh generators
+- [`physicsnemo.mesh.io`](./io/) - PyVista import/export
+- [`physicsnemo.mesh.examples`](./examples/) - Example mesh generators
 
 ---
 
diff --git a/physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py b/physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py
index 6698be028d..4afc17913f 100644
--- a/physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py
+++ b/physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py
@@ -53,7 +53,7 @@ def count_neighbors(
     grid: wp.HashGrid,
     wp_points: wp.array(dtype=wp.vec3),
     wp_queries: wp.array(dtype=wp.vec3),
-    wp_launch_device: wp.context.Device | None,
+    wp_launch_device: wp.Device | None,
     wp_launch_stream: wp.Stream | None,
     radius: float,
     N_queries: int,
@@ -65,7 +65,7 @@ def count_neighbors(
         grid (wp.HashGrid): The hash grid to use for the search.
         wp_points (wp.array): The points to search in, as a warp array.
         wp_queries (wp.array): The queries to search for, as a warp array.
-        wp_launch_device (wp.context.Device | None): The device to launch the kernel on.
+        wp_launch_device (wp.Device | None): The device to launch the kernel on.
         wp_launch_stream (wp.Stream | None): The stream to launch the kernel on.
         radius (float): The radius that bounds the search.
         N_queries (int): Total number of query points.
@@ -114,7 +114,7 @@ def gather_neighbors(
     wp_points: wp.array(dtype=wp.vec3),
     wp_queries: wp.array(dtype=wp.vec3),
     wp_offset: wp.array(dtype=wp.int32),
-    wp_launch_device: wp.context.Device | None,
+    wp_launch_device: wp.Device | None,
     wp_launch_stream: wp.Stream | None,
     radius: float,
     N_queries: int,
@@ -131,7 +131,7 @@ def gather_neighbors(
         wp_points (wp.array): The points to search in, as a warp array.
         wp_queries (wp.array): The queries to search for, as a warp array.
         wp_offset (wp.array): The offset in output for each input point, as a warp array.
-        wp_launch_device (wp.context.Device | None): The device to launch the kernel on.
+        wp_launch_device (wp.Device | None): The device to launch the kernel on.
         wp_launch_stream (wp.Stream | None): The stream to launch the kernel on.
         radius (float): The radius that bounds the search.
         N_queries (int): Total number of query points.