nanomanifold 0.1.1__tar.gz

nanomanifold-0.1.1/.github/workflows/release.yml

@@ -0,0 +1,65 @@
+name: Release
+
+on:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: write
+  id-token: write
+
+jobs:
+  release:
+    runs-on: ubuntu-latest
+    environment: pypi
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: astral-sh/setup-uv@v4
+        with:
+          python-version: "3.11"
+      - name: Detect version change
+        id: version
+        run: |
+          new_version=$(uv version --short)
+          echo "version=$new_version" >> "$GITHUB_OUTPUT"
+          tmp_dir=$(mktemp -d)
+          git show HEAD^:pyproject.toml 2>/dev/null > "$tmp_dir/pyproject.toml" || true
+          if [ -s "$tmp_dir/pyproject.toml" ]; then
+            old_version=$(uv version --short --project "$tmp_dir")
+          else
+            old_version=""
+          fi
+          if [ "$new_version" != "$old_version" ]; then
+            echo "changed=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "changed=false" >> "$GITHUB_OUTPUT"
+          fi
+      - name: Push tag
+        if: steps.version.outputs.changed == 'true'
+        run: |
+          git fetch --tags
+          if git rev-parse "v${{ steps.version.outputs.version }}" >/dev/null 2>&1; then
+            echo "Tag v${{ steps.version.outputs.version }} already exists"
+            exit 0
+          fi
+          git config user.name "${GITHUB_ACTOR}"
+          git config user.email "${GITHUB_ACTOR}@users.noreply.github.com"
+          git tag "v${{ steps.version.outputs.version }}"
+          git push origin "v${{ steps.version.outputs.version }}"
+      - name: Build
+        if: steps.version.outputs.changed == 'true'
+        run: uv build
+      - name: Publish to PyPI
+        if: steps.version.outputs.changed == 'true'
+        run: uv publish --check-url https://pypi.org/simple
+      - name: Create GitHub release
+        if: steps.version.outputs.changed == 'true'
+        uses: softprops/action-gh-release@v1
+        with:
+          tag_name: v${{ steps.version.outputs.version }}
+          name: v${{ steps.version.outputs.version }}
+          generate_release_notes: true
+          files: dist/*

nanomanifold-0.1.1/.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

nanomanifold-0.1.1/.python-version

@@ -0,0 +1 @@
+3.11

nanomanifold-0.1.1/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: nanomanifold
+Version: 0.1.1
+Summary: SO3/SE3 operations on any backend
+Author-email: Andrea Boscolo Camiletto <abcamiletto@gmail.com>
+Requires-Python: >=3.11
+Requires-Dist: array-api-compat>=1.12.0
+Requires-Dist: jaxtyping>=0.3.2
+Requires-Dist: numpy>=2.3.2
+Description-Content-Type: text/markdown
+
+# nanomanifold
+
+Fast, batched and differentiable SO3/SE3 transforms for any backend (NumPy, PyTorch, JAX, ...)
+Works directly on arrays, defined as:
+
+- **SO3**: unit quaternions `[w, x, y, z]` for 3D rotations, shape `(..., 4)`
+- **SE3**: concatenated `[quat, translation]`, shape `(..., 7)`
+
+```python
+import numpy as np
+from nanomanifold import SO3, SE3
+
+# Rotations stored as quaternion arrays [w,x,y,z]
+q = SO3.from_axis_angle(np.array([0, 0, 1]), np.pi/4)  # 45° around Z
+points = np.array([[1, 0, 0], [0, 1, 0]])
+rotated = SO3.rotate_points(q, points)
+
+# Rigid transforms stored as 7D arrays [quat, translation]
+T = SE3.from_rt(q, np.array([1, 0, 0]))  # rotation + translation
+transformed = SE3.transform_points(T, points)
+```
+
+## Features
+
+**Array-based API** — all functions operate directly on arrays from any backend  
+**Backend agnostic** — works with NumPy, PyTorch, JAX, CuPy, Dask, and more  
+**Batched operations** — process thousands of transforms at once  
+**Differentiable** — automatic gradients where supported (PyTorch/JAX)  
+**Memory efficient** — quaternions instead of matrices  
+**Numerically stable** — handles edge cases and singularities
+
+## Quick Start
+
+### Rotations (SO3)
+
+```python
+from nanomanifold import SO3
+
+# Create rotations
+q1 = SO3.from_axis_angle([1, 0, 0], np.pi/2)    # 90° around X
+q2 = SO3.from_euler([0, 0, np.pi/4])            # 45° around Z
+q3 = SO3.from_matrix(rotation_matrix)
+
+# Compose and interpolate
+q_combined = SO3.multiply(q1, q2)
+q_halfway = SO3.slerp(q1, q2, t=0.5)
+
+# Apply to points
+points = np.array([[1, 0, 0], [0, 1, 0]])
+rotated = SO3.rotate_points(q_combined, points)
+```
+
+### Rigid Transforms (SE3)
+
+```python
+from nanomanifold import SE3
+
+# Create transforms
+T1 = SE3.from_rt(q1, [1, 2, 3])               # rotation + translation
+T2 = SE3.from_matrix(transformation_matrix)
+
+# Compose transforms
+T_combined = SE3.multiply(T1, T2)
+T_inverse = SE3.inverse(T_combined)
+
+# Apply to points
+transformed = SE3.transform_points(T_combined, points)
+```
+
+## API Reference
+
+### SO3 (3D Rotations)
+
+| Function                       | Input → Output                      | Description                 |
+| ------------------------------ | ----------------------------------- | --------------------------- |
+| `from_axis_angle(axis, angle)` | `(...,3), (...) → (...,4)`          | Create from axis-angle      |
+| `from_euler(angles)`           | `(...,3) → (...,4)`                 | Create from Euler angles    |
+| `from_matrix(R)`               | `(...,3,3) → (...,4)`               | Create from rotation matrix |
+| `multiply(q1, q2)`             | `(...,4), (...,4) → (...,4)`        | Compose rotations           |
+| `slerp(q1, q2, t)`             | `(...,4), (...,4), (...) → (...,4)` | Spherical interpolation     |
+| `rotate_points(q, points)`     | `(...,4), (...,N,3) → (...,N,3)`    | Rotate 3D points            |
+
+### SE3 (Rigid Transforms)
+
+| Function                      | Input → Output                   | Description                        |
+| ----------------------------- | -------------------------------- | ---------------------------------- |
+| `from_rt(q, t)`               | `(...,4), (...,3) → (...,7)`     | Create from rotation + translation |
+| `from_matrix(T)`              | `(...,4,4) → (...,7)`            | Create from 4×4 matrix             |
+| `multiply(T1, T2)`            | `(...,7), (...,7) → (...,7)`     | Compose transforms                 |
+| `inverse(T)`                  | `(...,7) → (...,7)`              | Invert transform                   |
+| `transform_points(T, points)` | `(...,7), (...,N,3) → (...,N,3)` | Transform 3D points                |
+
+## Installation
+
+```bash
+pip install nanomanifold
+```
+
+## License
+
+MIT

nanomanifold-0.1.1/README.md

@@ -0,0 +1,101 @@
+# nanomanifold
+
+Fast, batched and differentiable SO3/SE3 transforms for any backend (NumPy, PyTorch, JAX, ...)
+Works directly on arrays, defined as:
+
+- **SO3**: unit quaternions `[w, x, y, z]` for 3D rotations, shape `(..., 4)`
+- **SE3**: concatenated `[quat, translation]`, shape `(..., 7)`
+
+```python
+import numpy as np
+from nanomanifold import SO3, SE3
+
+# Rotations stored as quaternion arrays [w,x,y,z]
+q = SO3.from_axis_angle(np.array([0, 0, 1]), np.pi/4)  # 45° around Z
+points = np.array([[1, 0, 0], [0, 1, 0]])
+rotated = SO3.rotate_points(q, points)
+
+# Rigid transforms stored as 7D arrays [quat, translation]
+T = SE3.from_rt(q, np.array([1, 0, 0]))  # rotation + translation
+transformed = SE3.transform_points(T, points)
+```
+
+## Features
+
+**Array-based API** — all functions operate directly on arrays from any backend  
+**Backend agnostic** — works with NumPy, PyTorch, JAX, CuPy, Dask, and more  
+**Batched operations** — process thousands of transforms at once  
+**Differentiable** — automatic gradients where supported (PyTorch/JAX)  
+**Memory efficient** — quaternions instead of matrices  
+**Numerically stable** — handles edge cases and singularities
+
+## Quick Start
+
+### Rotations (SO3)
+
+```python
+from nanomanifold import SO3
+
+# Create rotations
+q1 = SO3.from_axis_angle([1, 0, 0], np.pi/2)    # 90° around X
+q2 = SO3.from_euler([0, 0, np.pi/4])            # 45° around Z
+q3 = SO3.from_matrix(rotation_matrix)
+
+# Compose and interpolate
+q_combined = SO3.multiply(q1, q2)
+q_halfway = SO3.slerp(q1, q2, t=0.5)
+
+# Apply to points
+points = np.array([[1, 0, 0], [0, 1, 0]])
+rotated = SO3.rotate_points(q_combined, points)
+```
+
+### Rigid Transforms (SE3)
+
+```python
+from nanomanifold import SE3
+
+# Create transforms
+T1 = SE3.from_rt(q1, [1, 2, 3])               # rotation + translation
+T2 = SE3.from_matrix(transformation_matrix)
+
+# Compose transforms
+T_combined = SE3.multiply(T1, T2)
+T_inverse = SE3.inverse(T_combined)
+
+# Apply to points
+transformed = SE3.transform_points(T_combined, points)
+```
+
+## API Reference
+
+### SO3 (3D Rotations)
+
+| Function                       | Input → Output                      | Description                 |
+| ------------------------------ | ----------------------------------- | --------------------------- |
+| `from_axis_angle(axis, angle)` | `(...,3), (...) → (...,4)`          | Create from axis-angle      |
+| `from_euler(angles)`           | `(...,3) → (...,4)`                 | Create from Euler angles    |
+| `from_matrix(R)`               | `(...,3,3) → (...,4)`               | Create from rotation matrix |
+| `multiply(q1, q2)`             | `(...,4), (...,4) → (...,4)`        | Compose rotations           |
+| `slerp(q1, q2, t)`             | `(...,4), (...,4), (...) → (...,4)` | Spherical interpolation     |
+| `rotate_points(q, points)`     | `(...,4), (...,N,3) → (...,N,3)`    | Rotate 3D points            |
+
+### SE3 (Rigid Transforms)
+
+| Function                      | Input → Output                   | Description                        |
+| ----------------------------- | -------------------------------- | ---------------------------------- |
+| `from_rt(q, t)`               | `(...,4), (...,3) → (...,7)`     | Create from rotation + translation |
+| `from_matrix(T)`              | `(...,4,4) → (...,7)`            | Create from 4×4 matrix             |
+| `multiply(T1, T2)`            | `(...,7), (...,7) → (...,7)`     | Compose transforms                 |
+| `inverse(T)`                  | `(...,7) → (...,7)`              | Invert transform                   |
+| `transform_points(T, points)` | `(...,7), (...,N,3) → (...,N,3)` | Transform 3D points                |
+
+## Installation
+
+```bash
+pip install nanomanifold
+```
+
+## License
+
+MIT

nanomanifold-0.1.1/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "nanomanifold"
+version = "0.1.1"
+description = "SO3/SE3 operations on any backend"
+readme = "README.md"
+authors = [
+    { name = "Andrea Boscolo Camiletto", email = "abcamiletto@gmail.com" },
+]
+requires-python = ">=3.11"
+dependencies = ["array-api-compat>=1.12.0", "jaxtyping>=0.3.2", "numpy>=2.3.2"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = ["jax>=0.7.1", "pytest>=8.4.1", "scipy>=1.16.1", "torch>=2.8.0"]
+
+[tool.ruff]
+line-length = 140
+lint.ignore = ["E741", "E743", "F722"] # ignore ambiguous variable names
+lint.extend-select = ["I"]

nanomanifold-0.1.1/src/nanomanifold/SE3/__init__.py

@@ -0,0 +1,21 @@
+from .canonicalize import canonicalize
+from .conversions.matrix import from_matrix, to_matrix
+from .conversions.rt import from_rt, to_rt
+from .exp import exp
+from .inverse import inverse
+from .log import log
+from .multiply import multiply
+from .transform_points import transform_points
+
+__all__ = [
+    "from_matrix",
+    "to_matrix",
+    "from_rt",
+    "to_rt",
+    "canonicalize",
+    "multiply",
+    "inverse",
+    "transform_points",
+    "log",
+    "exp",
+]

nanomanifold-0.1.1/src/nanomanifold/SE3/canonicalize.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3.canonicalize import canonicalize as canonicalize_quat
+
+
+def canonicalize(se3: Float[Any, "... 7"]) -> Float[Any, "... 7"]:
+    """Canonicalize SE(3) representation by canonicalizing the quaternion part.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        Canonicalized SE(3) representation with quaternion w >= 0
+    """
+    xp = get_namespace(se3)
+
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+
+    quat_canonical = canonicalize_quat(quat)
+
+    return xp.concatenate([quat_canonical, translation], axis=-1)

nanomanifold-0.1.1/src/nanomanifold/SE3/conversions/__init__.py

@@ -0,0 +1,9 @@
+from .matrix import from_matrix, to_matrix
+from .rt import from_rt, to_rt
+
+__all__ = [
+    "from_matrix",
+    "to_matrix",
+    "from_rt",
+    "to_rt",
+]

nanomanifold-0.1.1/src/nanomanifold/SE3/conversions/matrix.py

@@ -0,0 +1,57 @@
+"""Matrix conversions for SE(3) transformations."""
+
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold import SO3
+from nanomanifold.common import get_namespace
+
+from ..canonicalize import canonicalize
+
+
+def to_matrix(se3: Float[Any, "... 7"]) -> Float[Any, "... 4 4"]:
+    """Convert SE(3) representation to 4x4 transformation matrix.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        4x4 transformation matrix (..., 4, 4)
+    """
+    xp = get_namespace(se3)
+
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+
+    R = SO3.to_matrix(quat)
+
+    translation_column = translation[..., None]
+    top_block = xp.concatenate([R, translation_column], axis=-1)
+
+    zeros = xp.zeros(top_block.shape[:-2] + (1, 3), dtype=se3.dtype)
+    ones = xp.ones(top_block.shape[:-2] + (1, 1), dtype=se3.dtype)
+    bottom_row = xp.concatenate([zeros, ones], axis=-1)
+
+    return xp.concatenate([top_block, bottom_row], axis=-2)
+
+
+def from_matrix(matrix: Float[Any, "... 4 4"]) -> Float[Any, "... 7"]:
+    """Convert 4x4 transformation matrix to SE(3) representation.
+
+    Args:
+        matrix: 4x4 transformation matrix (..., 4, 4)
+
+    Returns:
+        SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+    """
+    xp = get_namespace(matrix)
+
+    R = matrix[..., :3, :3]
+
+    quat = SO3.from_matrix(R)
+
+    translation = matrix[..., :3, 3]
+
+    se3 = xp.concatenate([quat, translation], axis=-1)
+    return canonicalize(se3)

nanomanifold-0.1.1/src/nanomanifold/SE3/conversions/rt.py

@@ -0,0 +1,30 @@
+from nanomanifold.common import get_namespace
+
+
+def from_rt(quat, translation):
+    """Create SE(3) representation from rotation quaternion and translation.
+
+    Args:
+        quat: Rotation quaternion (..., 4) as [w, x, y, z]
+        translation: Translation vector (..., 3)
+
+    Returns:
+        SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+    """
+    xp = get_namespace(quat)
+    return xp.concatenate([quat, translation], axis=-1)
+
+
+def to_rt(se3):
+    """Extract rotation quaternion and translation from SE(3) representation.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        quat: Rotation quaternion (..., 4) as [w, x, y, z]
+        translation: Translation vector (..., 3)
+    """
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+    return quat, translation

nanomanifold-0.1.1/src/nanomanifold/SE3/exp.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import exp as so3_exp
+from nanomanifold.SO3 import hat
+
+
+def exp(tangent_vector: Float[Any, "... 6"]) -> Float[Any, "... 7"]:
+    """Compute the exponential map from se(3) tangent space to SE(3) manifold.
+
+    The exponential map takes a tangent vector in the Lie algebra se(3)
+    and returns the corresponding SE(3) transformation. This is the inverse
+    operation of log().
+
+    The se(3) exponential map takes a 6-vector [ω, ρ] where:
+    - ω ∈ ℝ³ is the angular velocity (rotation part)
+    - ρ ∈ ℝ³ is the translational velocity
+
+    The formula involves:
+    - R = exp_SO3(ω) for the rotation quaternion
+    - t = V * ρ where V is the left Jacobian matrix
+
+    Args:
+        tangent_vector: Tangent vector in se(3) as [ω, ρ] of shape (..., 6)
+
+    Returns:
+        SE(3) transformation in [w, x, y, z, tx, ty, tz] format of shape (..., 7)
+    """
+    xp = get_namespace(tangent_vector)
+
+    omega = tangent_vector[..., :3]
+    rho = tangent_vector[..., 3:6]
+
+    q = so3_exp(omega)
+    omega_norm = xp.linalg.norm(omega, axis=-1, keepdims=True)
+
+    eps = xp.finfo(omega.dtype).eps
+    small_angle_threshold = xp.asarray(max(1e-6, float(eps)), dtype=omega.dtype)
+    small_angle_mask = omega_norm < small_angle_threshold
+
+    omega_cross = hat(omega)
+    omega_cross_sq = xp.matmul(omega_cross, omega_cross)
+
+    identity = xp.eye(3, dtype=omega.dtype)
+    identity = xp.broadcast_to(identity, omega.shape[:-1] + (3, 3))
+
+    V_small = identity + 0.5 * omega_cross + (1.0 / 12.0) * omega_cross_sq
+
+    cos_norm = xp.cos(omega_norm)
+    sin_norm = xp.sin(omega_norm)
+
+    safe_norm = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm)
+    safe_norm_sq = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm**2)
+    safe_norm_cub = safe_norm_sq * safe_norm
+
+    A = (1.0 - cos_norm) / safe_norm_sq
+    B = (safe_norm - sin_norm) / safe_norm_cub
+    A = xp.reshape(A, A.shape[:-1])[..., None, None]
+    B = xp.reshape(B, B.shape[:-1])[..., None, None]
+
+    V_large = identity + A * omega_cross + B * omega_cross_sq
+
+    mask = xp.reshape(small_angle_mask, small_angle_mask.shape[:-1])[..., None, None]
+    V = xp.where(mask, V_small, V_large)
+    V = xp.reshape(V, omega.shape[:-1] + (3, 3))
+
+    t = xp.matmul(V, rho[..., None])[..., 0]
+
+    se3 = xp.concatenate([q, t], axis=-1)
+
+    return se3

nanomanifold-0.1.1/src/nanomanifold/SE3/inverse.py

@@ -0,0 +1,37 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import inverse as so3_inverse
+from nanomanifold.SO3 import rotate_points
+
+from .canonicalize import canonicalize
+
+
+def inverse(se3: Float[Any, "... 7"]) -> Float[Any, "... 7"]:
+    """Compute the inverse of SE(3) transformations.
+
+    For an SE(3) transformation T = [R, t] represented as [q, t],
+    the inverse is T^(-1) = [R^T, -R^T * t] represented as [q^(-1), -q^(-1) * t].
+
+    Args:
+        se3: SE(3) transformation in [w, x, y, z, tx, ty, tz] format
+
+    Returns:
+        Inverse SE(3) transformation
+    """
+    xp = get_namespace(se3)
+
+    se3 = canonicalize(se3)
+
+    q = se3[..., :4]
+    t = se3[..., 4:7]
+
+    q_inv = so3_inverse(q)
+
+    t_inv = -rotate_points(q_inv, t[..., None, :]).squeeze(-2)
+
+    result = xp.concatenate([q_inv, t_inv], axis=-1)
+
+    return canonicalize(result)

nanomanifold-0.1.1/src/nanomanifold/SE3/log.py

@@ -0,0 +1,75 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import hat
+from nanomanifold.SO3 import log as so3_log
+
+from .canonicalize import canonicalize
+
+
+def log(se3: Float[Any, "... 7"]) -> Float[Any, "... 6"]:
+    """Compute the logarithmic map of SE(3) to its Lie algebra se(3).
+
+    The logarithmic map takes an SE(3) transformation and returns the corresponding
+    tangent vector in the Lie algebra se(3). This is the inverse operation of exp().
+
+    The SE(3) logarithmic map computes a 6-vector [ω, ρ] where:
+    - ω ∈ ℝ³ is the angular velocity (rotation part, same as SO(3) log)
+    - ρ ∈ ℝ³ is the transformed translation part
+
+    The formula involves:
+    - ω = log_SO3(R) where R is the rotation quaternion
+    - ρ = V^(-1) * t where V is the left Jacobian inverse and t is the translation
+
+    Args:
+        se3: SE(3) transformation in [w, x, y, z, tx, ty, tz] format of shape (..., 7)
+
+    Returns:
+        Tangent vector in se(3) as [ω, ρ] of shape (..., 6)
+    """
+    xp = get_namespace(se3)
+    se3 = canonicalize(se3)
+
+    q = se3[..., :4]
+    t = se3[..., 4:7]
+
+    omega = so3_log(q)
+    omega_norm = xp.linalg.norm(omega, axis=-1, keepdims=True)
+
+    eps = xp.finfo(omega.dtype).eps
+    small_angle_threshold = xp.asarray(max(1e-6, float(eps)), dtype=omega.dtype)
+    small_angle_mask = omega_norm < small_angle_threshold
+
+    omega_cross = hat(omega)
+    omega_cross_sq = xp.matmul(omega_cross, omega_cross)
+
+    identity = xp.eye(3, dtype=omega.dtype)
+    identity = xp.broadcast_to(identity, omega.shape[:-1] + (3, 3))
+
+    V_inv_small = identity - 0.5 * omega_cross + (1.0 / 12.0) * omega_cross_sq
+
+    half_norm = omega_norm / 2.0
+    cos_half = xp.cos(half_norm)
+    sin_half = xp.sin(half_norm)
+
+    safe_sin_half = xp.where(small_angle_mask, xp.ones_like(sin_half), sin_half)
+    cot_half = cos_half / safe_sin_half
+
+    safe_norm = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm)
+    safe_norm_sq = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm**2)
+    B = (1.0 - 0.5 * safe_norm * cot_half) / safe_norm_sq
+    B = xp.reshape(B, B.shape[:-1])[..., None, None]
+
+    V_inv_large = identity - 0.5 * omega_cross + B * omega_cross_sq
+
+    mask = xp.reshape(small_angle_mask, small_angle_mask.shape[:-1])[..., None, None]
+    V_inv = xp.where(mask, V_inv_small, V_inv_large)
+    V_inv = xp.reshape(V_inv, omega.shape[:-1] + (3, 3))
+
+    rho = xp.matmul(V_inv, t[..., None])[..., 0]
+
+    tangent = xp.concatenate([omega, rho], axis=-1)
+
+    return tangent