mujoco_lidar 0.3.0__tar.gz

mujoco_lidar-0.3.0/LICENSE

@@ -0,0 +1,23 @@
+MIT License
+
+Copyright (c) 2025 Yufei Jia
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+

mujoco_lidar-0.3.0/PKG-INFO

@@ -0,0 +1,175 @@
+Metadata-Version: 2.4
+Name: mujoco_lidar
+Version: 0.3.0
+Summary: A LiDAR sensor designed for MuJoCo
+Author-email: Yufei Jia <jyf23@mails.tsinghua.edu.cn>
+License: MIT
+Project-URL: Homepage, https://github.com/TATP-233/MuJoCo-LiDAR.git
+Project-URL: Repository, https://github.com/TATP-233/MuJoCo-LiDAR.git
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: mujoco>=3.2.0
+Requires-Dist: scipy
+Requires-Dist: pynput
+Requires-Dist: matplotlib
+Requires-Dist: zhplot
+Provides-Extra: taichi
+Requires-Dist: taichi>=1.6.0; extra == "taichi"
+Requires-Dist: tibvh>=0.1.2; extra == "taichi"
+Provides-Extra: jax
+Requires-Dist: jax[cuda12]; extra == "jax"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.3.0; extra == "dev"
+Provides-Extra: examples
+Requires-Dist: onnxruntime>=1.19.2; extra == "examples"
+Requires-Dist: etils[epath]>=1.13.0; extra == "examples"
+Dynamic: license-file
+
+# MuJoCo-LiDAR: High-Performance LiDAR Simulation
+
+[![PyPI](https://img.shields.io/pypi/v/mujoco-lidar)](https://pypi.org/project/mujoco-lidar/)
+[![Python](https://img.shields.io/pypi/pyversions/mujoco-lidar)](https://pypi.org/project/mujoco-lidar/)
+
+High-performance LiDAR simulation for MuJoCo with CPU, Taichi, and JAX backends.
+
+<p align="center">
+  <img src="./assets/go2.png" width="49%" />
+  <img src="./assets/g1.png" width="49%" />
+</p>
+<p align="center">
+  <img src="./assets/g1_native.png" width="32%" />
+  <img src="./assets/go2_native.png" width="32%" />
+  <img src="./assets/lidar_rviz.png" width="33%" />
+</p>
+
+[中文文档](README_zh.md) | [Installation](docs/en/INSTALLATION.md) | [Usage Guide](docs/en/USAGE.md) | [Development](docs/en/DEVELOPMENT.md) | [Contributing](CONTRIBUTING.md)
+
+## Features
+
+- **Multi-Backend Support**:
+  - **CPU**: MuJoCo native `mj_multiRay`, no GPU required
+  - **Taichi**: GPU parallel computing, supports Mesh and Hfield
+  - **JAX**: GPU + MJX integration, batch simulation support
+- **High Performance**: 1M+ rays/sec on GPU, real-time BVH construction
+- **Multiple LiDAR Models**: Velodyne (HDL-64E, VLP-32C), Livox (mid360, avia), Ouster (OS-128), custom patterns
+- **ROS Integration**: Ready-to-use ROS1/ROS2 examples
+
+## Quick Start
+
+### Installation
+
+**From PyPI:**
+
+```bash
+# Basic (CPU only)
+uv add mujoco-lidar
+
+# With Taichi backend (GPU)
+uv add "mujoco-lidar[taichi]"
+
+# With JAX backend (GPU + batch)
+uv add "mujoco-lidar[jax]"
+```
+
+**From Source:**
+
+```bash
+git clone https://github.com/TATP-233/MuJoCo-LiDAR.git
+cd MuJoCo-LiDAR
+
+uv sync --extra dev                        # CPU only
+uv sync --extra dev --extra taichi         # with Taichi backend
+uv sync --extra dev --extra taichi --extra jax  # all backends
+```
+
+See [Installation Guide](docs/en/INSTALLATION.md) for details.
+
+### Basic Usage
+
+```python
+import mujoco
+from mujoco_lidar import MjLidarWrapper, scan_gen
+
+# Load model
+model = mujoco.MjModel.from_xml_path("scene.xml")
+data = mujoco.MjData(model)
+
+# Create LiDAR
+lidar = MjLidarWrapper(
+    model,
+    site_name="lidar_site",
+    backend="cpu",  # or "taichi", "jax"
+    cutoff_dist=50.0
+)
+
+# Generate scan pattern
+theta, phi = scan_gen.generate_HDL64()
+
+# Trace rays
+ranges = lidar.trace_rays(data, theta, phi)
+```
+
+See [Usage Guide](docs/en/USAGE.md) for more examples.
+
+## Performance
+
+| Backend | Rays/sec | Hardware | Batch Support |
+|---------|----------|----------|---------------|
+| CPU     | ~9M      | Native   | No            |
+| Taichi  | ~62M     | GPU      | Yes           |
+| JAX     | ~231M    | GPU      | Yes           |
+
+Run benchmarks: `make benchmark`
+
+## Documentation
+
+- [Installation Guide](docs/en/INSTALLATION.md) - Detailed installation instructions
+- [Usage Guide](docs/en/USAGE.md) - Examples and tutorials
+- [API Reference](docs/en/API.md) - Complete API documentation
+- [Development Guide](docs/en/DEVELOPMENT.md) - Contributing and testing
+- [Project Structure](docs/en/PROJECT_STRUCTURE.md) - Codebase organization
+
+## Examples
+
+- [examples/example_native.py](examples/example_native.py) - CPU backend
+- [examples/example_taichi.py](examples/example_taichi.py) - Taichi backend
+- [examples/lidar_vis_ros2.py](examples/lidar_vis_ros2.py) - ROS2 integration
+- [examples/unitree_go2_ros2.py](examples/unitree_go2_ros2.py) - Unitree Go2 robot
+
+## Development
+
+```bash
+git clone https://github.com/TATP-233/MuJoCo-LiDAR.git
+cd MuJoCo-LiDAR
+uv sync --extra dev
+
+make test      # Run tests
+make lint      # Check code quality
+make benchmark # Run performance tests
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines.
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## Citation
+
+If you use this project in your research, please cite:
+
+```bibtex
+@software{mujoco_lidar,
+  title = {MuJoCo-LiDAR: High-Performance LiDAR Simulation},
+  author = {Yufei Jia},
+  year = {2024},
+  url = {https://github.com/TATP-233/MuJoCo-LiDAR}
+}
+```

mujoco_lidar-0.3.0/README.md

@@ -0,0 +1,141 @@
+# MuJoCo-LiDAR: High-Performance LiDAR Simulation
+
+[![PyPI](https://img.shields.io/pypi/v/mujoco-lidar)](https://pypi.org/project/mujoco-lidar/)
+[![Python](https://img.shields.io/pypi/pyversions/mujoco-lidar)](https://pypi.org/project/mujoco-lidar/)
+
+High-performance LiDAR simulation for MuJoCo with CPU, Taichi, and JAX backends.
+
+<p align="center">
+  <img src="./assets/go2.png" width="49%" />
+  <img src="./assets/g1.png" width="49%" />
+</p>
+<p align="center">
+  <img src="./assets/g1_native.png" width="32%" />
+  <img src="./assets/go2_native.png" width="32%" />
+  <img src="./assets/lidar_rviz.png" width="33%" />
+</p>
+
+[中文文档](README_zh.md) | [Installation](docs/en/INSTALLATION.md) | [Usage Guide](docs/en/USAGE.md) | [Development](docs/en/DEVELOPMENT.md) | [Contributing](CONTRIBUTING.md)
+
+## Features
+
+- **Multi-Backend Support**:
+  - **CPU**: MuJoCo native `mj_multiRay`, no GPU required
+  - **Taichi**: GPU parallel computing, supports Mesh and Hfield
+  - **JAX**: GPU + MJX integration, batch simulation support
+- **High Performance**: 1M+ rays/sec on GPU, real-time BVH construction
+- **Multiple LiDAR Models**: Velodyne (HDL-64E, VLP-32C), Livox (mid360, avia), Ouster (OS-128), custom patterns
+- **ROS Integration**: Ready-to-use ROS1/ROS2 examples
+
+## Quick Start
+
+### Installation
+
+**From PyPI:**
+
+```bash
+# Basic (CPU only)
+uv add mujoco-lidar
+
+# With Taichi backend (GPU)
+uv add "mujoco-lidar[taichi]"
+
+# With JAX backend (GPU + batch)
+uv add "mujoco-lidar[jax]"
+```
+
+**From Source:**
+
+```bash
+git clone https://github.com/TATP-233/MuJoCo-LiDAR.git
+cd MuJoCo-LiDAR
+
+uv sync --extra dev                        # CPU only
+uv sync --extra dev --extra taichi         # with Taichi backend
+uv sync --extra dev --extra taichi --extra jax  # all backends
+```
+
+See [Installation Guide](docs/en/INSTALLATION.md) for details.
+
+### Basic Usage
+
+```python
+import mujoco
+from mujoco_lidar import MjLidarWrapper, scan_gen
+
+# Load model
+model = mujoco.MjModel.from_xml_path("scene.xml")
+data = mujoco.MjData(model)
+
+# Create LiDAR
+lidar = MjLidarWrapper(
+    model,
+    site_name="lidar_site",
+    backend="cpu",  # or "taichi", "jax"
+    cutoff_dist=50.0
+)
+
+# Generate scan pattern
+theta, phi = scan_gen.generate_HDL64()
+
+# Trace rays
+ranges = lidar.trace_rays(data, theta, phi)
+```
+
+See [Usage Guide](docs/en/USAGE.md) for more examples.
+
+## Performance
+
+| Backend | Rays/sec | Hardware | Batch Support |
+|---------|----------|----------|---------------|
+| CPU     | ~9M      | Native   | No            |
+| Taichi  | ~62M     | GPU      | Yes           |
+| JAX     | ~231M    | GPU      | Yes           |
+
+Run benchmarks: `make benchmark`
+
+## Documentation
+
+- [Installation Guide](docs/en/INSTALLATION.md) - Detailed installation instructions
+- [Usage Guide](docs/en/USAGE.md) - Examples and tutorials
+- [API Reference](docs/en/API.md) - Complete API documentation
+- [Development Guide](docs/en/DEVELOPMENT.md) - Contributing and testing
+- [Project Structure](docs/en/PROJECT_STRUCTURE.md) - Codebase organization
+
+## Examples
+
+- [examples/example_native.py](examples/example_native.py) - CPU backend
+- [examples/example_taichi.py](examples/example_taichi.py) - Taichi backend
+- [examples/lidar_vis_ros2.py](examples/lidar_vis_ros2.py) - ROS2 integration
+- [examples/unitree_go2_ros2.py](examples/unitree_go2_ros2.py) - Unitree Go2 robot
+
+## Development
+
+```bash
+git clone https://github.com/TATP-233/MuJoCo-LiDAR.git
+cd MuJoCo-LiDAR
+uv sync --extra dev
+
+make test      # Run tests
+make lint      # Check code quality
+make benchmark # Run performance tests
+```
+
+See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines.
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.
+
+## Citation
+
+If you use this project in your research, please cite:
+
+```bibtex
+@software{mujoco_lidar,
+  title = {MuJoCo-LiDAR: High-Performance LiDAR Simulation},
+  author = {Yufei Jia},
+  year = {2024},
+  url = {https://github.com/TATP-233/MuJoCo-LiDAR}
+}
+```

mujoco_lidar-0.3.0/pyproject.toml

@@ -0,0 +1,78 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mujoco_lidar"
+dynamic = ["version"]
+authors = [
+    {name = "Yufei Jia", email = "jyf23@mails.tsinghua.edu.cn"}
+]
+description = "A LiDAR sensor designed for MuJoCo"
+readme = "README.md"
+license = {text = "MIT"}
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "numpy >= 1.20.0",
+    "mujoco >= 3.2.0",
+    "scipy",
+    "pynput",
+    "matplotlib",
+    "zhplot",
+]
+
+[project.optional-dependencies]
+# Taichi backend with all dependencies (taichi, tibvh, and utilities)
+taichi = [
+    "taichi >= 1.6.0",
+    "tibvh >= 0.1.2",
+]
+# JAX backend
+jax = [
+    "jax[cuda12]",
+]
+dev = [
+    "pytest>=7.0",
+    "pytest-cov>=4.0",
+    "ruff>=0.3.0",
+]
+examples = [
+    "onnxruntime>=1.19.2",
+    "etils[epath]>=1.13.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/TATP-233/MuJoCo-LiDAR.git"
+Repository = "https://github.com/TATP-233/MuJoCo-LiDAR.git"
+
+[tool.setuptools]
+packages = ["mujoco_lidar"]
+package-dir = {"" = "src"}
+
+[tool.setuptools.dynamic]
+version = {attr = "mujoco_lidar.__version__"}
+
+[tool.setuptools.package-data]
+"mujoco_lidar" = ["scan_mode/*.npy"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+exclude = ["examples/"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM"]
+ignore = ["E501"]
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = "test_*.py"
+addopts = "-v"

mujoco_lidar-0.3.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mujoco_lidar-0.3.0/src/mujoco_lidar/__init__.py

@@ -0,0 +1,43 @@
+# Lazy import to avoid loading dependencies when not needed
+# Import the wrapper class directly
+from typing import Any
+
+from mujoco_lidar.lidar_wrapper import MjLidarWrapper
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "MjLidarWrapper",
+    # Scan generation functions (imported lazily via __getattr__)
+    "LivoxGeneratorTi",
+    "LivoxGenerator",  # From scan_gen_livox (requires taichi)
+    "generate_grid_scan_pattern",
+    "create_lidar_single_line",
+    "generate_HDL64",  # From scan_gen (no taichi needed)
+    "generate_vlp32",
+    "generate_os128",
+    "generate_airy96",
+]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy import for scan generation functions."""
+    # LivoxGeneratorTi requires taichi - import from scan_gen_livox_ti
+    if name == "LivoxGeneratorTi":
+        from mujoco_lidar.scan_gen_livox_ti import LivoxGeneratorTi
+
+        return LivoxGeneratorTi
+    # Other scan functions don't require taichi - import from scan_gen
+    elif name in [
+        "LivoxGenerator",
+        "generate_grid_scan_pattern",
+        "create_lidar_single_line",
+        "generate_HDL64",
+        "generate_vlp32",
+        "generate_os128",
+        "generate_airy96",
+    ]:
+        from mujoco_lidar import scan_gen
+
+        return getattr(scan_gen, name)
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

mujoco_lidar-0.3.0/src/mujoco_lidar/lidar_wrapper.py

@@ -0,0 +1,205 @@
+from typing import Any
+
+import mujoco
+import numpy as np
+
+
+class MjLidarWrapper:
+    """
+    MuJoCo LiDAR wrapper supporting CPU, Taichi, and JAX backends.
+
+    Args:
+        mj_model: MuJoCo model object
+        site_name: Name of the LiDAR site in the model
+        backend: 'cpu', 'taichi', or 'jax' (default: 'taichi')
+        cutoff_dist: Maximum ray distance in meters (default: 100.0)
+        args: Backend-specific arguments (see docs/en/API.md)
+    """
+
+    def __init__(
+        self,
+        mj_model: mujoco.MjModel,
+        site_name: str,
+        backend: str = "taichi",
+        cutoff_dist: float = 100.0,
+        args: dict[str, Any] = None,
+    ):
+        if args is None:
+            args = {}
+        if backend == "gpu":
+            backend = "taichi"
+        self.backend = backend
+        self.mj_model = mj_model
+        self.cutoff_dist = cutoff_dist
+        self.args = args
+
+        if backend == "taichi":
+            self._init_taichi_backend()
+        elif backend == "jax":
+            self._init_jax_backend()
+        elif backend == "cpu":
+            self._init_cpu_backend()
+        else:
+            raise ValueError(
+                f"Unsupported backend: {backend}, choose from 'cpu', 'taichi', or 'jax'"
+            )
+
+        self.site_name = site_name
+        self._sensor_pose = np.eye(4, dtype=np.float32)
+        self._local_rays: np.ndarray | None = None
+        self._distances: np.ndarray | None = None
+
+    def _init_taichi_backend(self) -> None:
+        """Initialize Taichi backend"""
+        try:
+            # Lazy import: only import when Taichi backend is selected
+            import taichi as ti
+
+            from mujoco_lidar.core_ti.mjlidar_ti import MjLidarTi
+
+            # Initialize Taichi if not already done
+            if not hasattr(ti, "_is_initialized") or not ti._is_initialized:
+                ti.init(arch=ti.gpu, **self.args.get("ti_init_args", {}))
+
+            # Create Taichi backend instance
+            geomgroup = self.args.get("geomgroup", None)
+            bodyexclude = self.args.get("bodyexclude", -1)
+            max_candidates = self.args.get("max_candidates", 64)
+            self._backend_instance = MjLidarTi(
+                self.mj_model,
+                cutoff_dist=self.cutoff_dist,
+                geomgroup=geomgroup,
+                bodyexclude=bodyexclude,
+                max_candidates=max_candidates,
+            )
+
+        except ImportError as e:
+            raise ImportError(
+                f"Failed to import Taichi backend dependencies. "
+                f'Please install taichi: uv add "mujoco-lidar[taichi]"\n'
+                f"Error: {e}"
+            ) from e
+
+    def _init_jax_backend(self) -> None:
+        """Initialize JAX backend"""
+        try:
+            from mujoco_lidar.core_jax.mjlidar_jax import MjLidarJax
+
+            geomgroup = self.args.get("geomgroup", None)
+            bodyexclude = self.args.get("bodyexclude", -1)
+
+            # Pass mj_model directly. MjLidarJax will extract what it needs.
+            self._backend_instance = MjLidarJax(
+                self.mj_model,
+                geom_ids=self.args.get("geom_ids"),
+                geomgroup=geomgroup,
+                bodyexclude=bodyexclude,
+            )
+
+        except ImportError as e:
+            raise ImportError(f"Failed to import JAX backend dependencies.\nError: {e}") from e
+
+    def _init_cpu_backend(self) -> None:
+        """Initialize CPU backend"""
+        try:
+            from mujoco_lidar.core_cpu.mjlidar_cpu import MjLidarCPU
+
+            geomgroup = self.args.get("geomgroup", None)
+            bodyexclude = self.args.get("bodyexclude", -1)
+            self._backend_instance = MjLidarCPU(
+                self.mj_model,
+                cutoff_dist=self.cutoff_dist,
+                geomgroup=geomgroup,
+                bodyexclude=bodyexclude,
+            )
+
+        except ImportError as e:
+            raise ImportError(f"Failed to import CPU backend dependencies.\nError: {e}") from e
+
+    @property
+    def sensor_position(self) -> np.ndarray:
+        return self._sensor_pose[:3, 3].copy()
+
+    @property
+    def sensor_rotation(self) -> np.ndarray:
+        return self._sensor_pose[:3, :3].copy()
+
+    def update_sensor_pose(self, mj_data: mujoco.MjData, site_name: str) -> None:
+        # For CPU/Taichi/JAX backend, mj_data is mujoco.MjData
+        if self.backend in ["cpu", "taichi", "jax"]:
+            self._sensor_pose[:3, :3] = mj_data.site(site_name).xmat.reshape(3, 3)
+            self._sensor_pose[:3, 3] = mj_data.site(site_name).xpos
+
+    def trace_rays(
+        self,
+        mj_data: mujoco.MjData,
+        ray_theta: np.ndarray,
+        ray_phi: np.ndarray,
+        site_name: str | None = None,
+    ) -> np.ndarray:
+        """
+        Trace rays.
+        For JAX backend, mj_data can be mujoco.MjData.
+        """
+        target_site = self.site_name if site_name is None else site_name
+
+        if self.backend == "jax":
+            # Update sensor pose for consistency
+            self.update_sensor_pose(mj_data, target_site)
+
+            # Use JITed trace_rays from backend instance
+            # This handles ray generation, transformation and rendering in one JIT call
+            self._distances, self._local_rays = self._backend_instance.trace_rays(
+                mj_data.geom_xpos,
+                mj_data.geom_xmat,
+                mj_data.site(target_site).xpos,
+                mj_data.site(target_site).xmat.reshape(3, 3),
+                ray_theta,
+                ray_phi,
+            )
+
+            return self._distances
+
+        elif self.backend == "taichi":
+            # Taichi Backend
+            self.update_sensor_pose(mj_data, target_site)
+            self._backend_instance.update(mj_data)
+
+            import taichi as ti
+
+            # Convert numpy arrays to Taichi ndarrays if necessary
+            if isinstance(ray_theta, np.ndarray):
+                theta_ti = ti.ndarray(dtype=ti.f32, shape=ray_theta.shape[0])
+                theta_ti.from_numpy(ray_theta.astype(np.float32))
+            else:
+                theta_ti = ray_theta
+
+            if isinstance(ray_phi, np.ndarray):
+                phi_ti = ti.ndarray(dtype=ti.f32, shape=ray_phi.shape[0])
+                phi_ti.from_numpy(ray_phi.astype(np.float32))
+            else:
+                phi_ti = ray_phi
+
+            self._backend_instance.trace_rays(self._sensor_pose, theta_ti, phi_ti)
+            return self._backend_instance.get_distances()
+
+        else:
+            # CPU Backend
+            self.update_sensor_pose(mj_data, target_site)
+            self._backend_instance.update(mj_data)
+            self._backend_instance.trace_rays(self._sensor_pose, ray_theta, ray_phi)
+            return self._backend_instance.get_distances()
+
+    def get_hit_points(self) -> np.ndarray:
+        if self.backend == "jax":
+            if self._distances is None or self._local_rays is None:
+                return np.zeros((0, 3), dtype=np.float32)
+            return np.asarray(self._distances[:, np.newaxis] * self._local_rays)
+        return self._backend_instance.get_hit_points()
+
+    def get_distances(self) -> np.ndarray:
+        if self.backend == "jax":
+            if self._distances is None:
+                return np.zeros(0, dtype=np.float32)
+            return np.asarray(self._distances)
+        return self._backend_instance.get_distances()