scipy.spatial.transform: Differentiable Framework-Agnostic 3D Transformations in Python

TL;DR

This paper introduces a framework-agnostic, differentiable implementation of 3D transformations in Python, compatible with various array libraries, enabling robust, GPU-accelerated, and autodiff-enabled workflows for robotics, vision, and simulation.

Contribution

It overhauls SciPy's spatial.transform to support any array library with autodiff, GPU acceleration, and JIT, making 3D spatial math more robust and widely usable in differentiable systems.

Findings

Supports GPU/TPU execution and JIT compilation.

Enables differentiation through native autodiff.

Improves robustness and correctness of 3D transformations.

Abstract

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.