scipy.spatial.transform: Differentiable Framework-Agnostic 3D Transformations in Python
This provides a production-grade, framework-agnostic solution for 3D spatial math in differentiable systems and ML, addressing a bottleneck in robotics, vision, and simulation workflows, though it is incremental as it builds on existing SciPy functionality.
The authors tackled the problem of implementing robust and differentiable 3D transformations for machine learning by overhauling SciPy's spatial.transform module to be compatible with various array libraries like JAX and PyTorch, enabling GPU execution and autodiff. This resulted in a framework-agnostic tool that supports scalable 3D transforms and applications such as drone simulations, with the changes merged into SciPy for upcoming release.
Three-dimensional rigid-body transforms, i.e. rotations and translations, are central to modern differentiable machine learning pipelines in robotics, vision, and simulation. However, numerically robust and mathematically correct implementations, particularly on SO(3), are error-prone due to issues such as axis conventions, normalizations, composition consistency and subtle errors that only appear in edge cases. SciPy's spatial$.$transform module is a rigorously tested Python implementation. However, it historically only supported NumPy, limiting adoption in GPU-accelerated and autodiff-based workflows. We present a complete overhaul of SciPy's spatial$.$transform functionality that makes it compatible with any array library implementing the Python array API, including JAX, PyTorch, and CuPy. The revised implementation preserves the established SciPy interface while enabling GPU/TPU execution, JIT compilation, vectorized batching, and differentiation via native autodiff of the chosen backend. We demonstrate how this foundation supports differentiable scientific computing through two case studies: (i) scalability of 3D transforms and rotations and (ii) a JAX drone simulation that leverages SciPy's Rotation for accurate integration of rotational dynamics. Our contributions have been merged into SciPy main and will ship in the next release, providing a framework-agnostic, production-grade basis for 3D spatial math in differentiable systems and ML.