jaxlie
Matrix Lie groups in JAX
Decision gist · record as of 2026-08-14
Yes, if you are building JAX-based computer vision or robotics code that requires differentiable rigid body transformations. The library is well-designed for manifold optimization and integrates cleanly with JAX's ecosystem. The aging maintenance status (477 days since last release) is not a blocker—the package is stable and the repo remains active—but check whether recent JAX API changes affect your use case. No known vulnerabilities.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires JAX installed; jaxlie requires Python >=3.8.
- Low friction: pure Python wheel with five runtime dependencies (jax, numpy, jax_dataclasses, typing_extensions, tyro).
- Maintenance status is aging—last commit 477 days ago—but the repo remains active and the package has been stable since its 2021 release.
License · maintenance · safety
MIT (permissive) — MIT license (permissive) places no restrictions on use, modification, or distribution in proprietary or open-source contexts.
last release 2025-04-24 (477 days) · last repo commit 2025-04-24 · 338 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 96,173 downloads/mo, #13,225 on PyPI
Alternatives
Verify before relying
pip install jaxlie
import jaxlie
# Create a 3D rotation from quaternion (wxyz)
rotation = jaxlie.SO3(jnp.array([1.0, 0.0, 0.0, 0.0]))
# Apply rotation to a 3D point
point = jnp.array([1.0, 0.0, 0.0])
rotated = rotation.apply(point)- Whether the package's AD support covers all use cases or has known limitations in reverse-mode differentiation.
- Performance characteristics and scalability when used with large batches or vmap over many group elements.
- Stability and numerical accuracy of Taylor approximations near singularities in practice.
What it is and what it does
jaxlie is a JAX library for working with matrix Lie groups commonly used in rigid body kinematics and transformations. It provides high-level dataclass implementations of SO2 (2D rotations), SE2 (2D rigid transforms), SO3 (3D rotations), and SE3 (3D rigid transforms), each parameterized in a way suitable for optimization and differentiation. Every group supports forward and reverse-mode autodiff-friendly operations including exp, log, matrix conversion, composition, and inversion.
The library is designed for computer vision and robotics workflows where you need to optimize over transformation manifolds or compose transformations in differentiable code. It integrates with JAX's function transformations (vmap, jit, grad) and supports broadcasting, pytree flattening, and serialization via flax. Taylor approximations handle numerical stability near singularities, and utilities like uniform random sampling and Euler angle conversion are included for SO3.
Use it for
- Optimize camera poses or object transforms in 3D vision pipelines using manifold-aware gradient descent.
- Compose and apply rigid body transformations in robotics simulation or control code with automatic differentiation.
- Batch-process rotations and transforms across multiple frames or trajectories using vmap.
- Convert between rotation representations (quaternions, matrices, Euler angles) while maintaining differentiability.
- Implement factor graph optimization for SLAM or pose estimation with Lie group constraints.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
Yes, if you are building JAX-based computer vision or robotics code that requires differentiable rigid body transformations.
The library is well-designed for manifold optimization and integrates cleanly with JAX's ecosystem. The aging maintenance status (477 days since last release) is not a blocker—the package is stable and the repo remains active—but check whether recent JAX API changes affect your use case. No known vulnerabilities.
Install
jaxlie on PyPI
Before you install
Low friction: pure Python wheel with five runtime dependencies (jax, numpy, jax_dataclasses, typing_extensions, tyro). Maintenance status is aging—last commit 477 days ago—but the repo remains active and the package has been stable since its 2021 release.
Requires JAX installed; jaxlie requires Python >=3.8.
License in practice
MIT license (permissive) places no restrictions on use, modification, or distribution in proprietary or open-source contexts.
Quickstart
pip install jaxlie
import jaxlie
# Create a 3D rotation from quaternion (wxyz)
rotation = jaxlie.SO3(jnp.array([1.0, 0.0, 0.0, 0.0]))
# Apply rotation to a 3D point
point = jnp.array([1.0, 0.0, 0.0])
rotated = rotation.apply(point)Verify before relying
- Whether the package's AD support covers all use cases or has known limitations in reverse-mode differentiation.
- Performance characteristics and scalability when used with large batches or vmap over many group elements.
- Stability and numerical accuracy of Taylor approximations near singularities in practice.
Package facts
| License | MIT permissive |
| Python support | Supports the current Python release >=3.8 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 5 packagesjaxjax_dataclassesnumpytyping_extensionstyro |
| Maintenance | Aging 477 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 96,173 / month, #13,225 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | License :: OSI Approved :: MIT LicenseOperating System :: OS IndependentProgramming Language :: Python :: 3.10Programming Language :: Python :: 3.7Programming Language :: Python :: 3.8Programming Language :: Python :: 3.9 |
Evidence: jaxlie-1.5.0-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “lie groups JAX”
- jaxliejaxlie implements Lie groups (SO2, SE2, SO3, SE3) for rigid body…
- passagemath-gapProvides Python interfaces to the GAP computational algebra system…
- equinoxEquinox provides neural network and model building on top of JAX with…
Give your agent the search over MCP, or paste the wish link into any chat.
More Mathematics packages
NetworkX provides data structures and algorithms for creating, analyzing, and manipulating graphs and networks, supporting everything from simple undirected graphs to complex directed and weighted networks.
kiwisolver is a Python binding to a fast C++ implementation of the Cassowary constraint solver, enabling you to solve systems of linear constraints and inequalities.
Install it if you need to solve constraint systems; skip it if you only need simple linear algebra.
SymPy is a Python library for symbolic mathematics, performing algebraic manipulation, calculus, equation solving, and mathematical expression simplification without numerical approximation.
ContourPy calculates contours of 2D quadrilateral grids using C++11 algorithms wrapped in Python, offering serial and multithreaded implementations without requiring Matplotlib as a dependency.
PyTorch provides GPU-accelerated tensor computation and automatic differentiation for building and training deep neural networks in Python.
onnxruntime loads and executes Open Neural Network Exchange (ONNX) models with a focus on inference performance across CPUs and accelerators.
Install it if you have ONNX models to run in production or development.
See also transformations · transforms3d · roma · jaxellip · json-e · libpinocchio · e3nn-jax · drjax · pin · pyquaternion