Keep quaternion backend dispatch scriptable (#1604)

Summary:
Pull Request resolved: #1604

PyMomentum lazily imports its Triton quaternion backend to avoid a circular dependency. TorchScript attempts to compile that import when a scripted caller uses rotation conversion with `backend="auto"`, causing inference workflows to fail before dispatch can fall back to the Torch backend.

Return through the pure Torch implementation immediately when scripting, before resolving the lazy Triton backend. Eager automatic and Triton dispatch remain unchanged, while scripted graphs contain no Python operation and remain serializable. Cover all three rotation conversion entry points, including `torch.jit.save`.

___

Differential Revision: D114808178

fbshipit-source-id: 56928d20f44221765d78f02e72596a3ff4ee47a5

Author: kunhe001

pymomentum/quaternion.py

@@ -206,9 +206,11 @@ def to_rotation_matrix_assume_normalized(
     Convert quaternions to 3x3 rotation matrices.
 
     :parameter q: (nBatch x k x 4) tensor with the quaternions in ((x, y, z), w) format.
-    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch.
+    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch. Scripted calls always use the Torch backend.
     :return: (nBatch x k x 3 x 3) tensor with 3x3 rotation matrices.
     """
+    if torch.jit.is_scripting():
+        return torch_quaternion.to_rotation_matrix_assume_normalized(q)
     backend = resolve_backend(
         backend,
         q,
@@ -228,10 +230,12 @@ def to_rotation_matrix(
     Convert quaternions to 3x3 rotation matrices.
 
     :parameter q: (nBatch x k x 4) tensor with the quaternions in ((x, y, z), w) format.
-    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch.
+    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch. Scripted calls always use the Torch backend.
     :parameter eps: Minimum quaternion norm used during normalization.
     :return: (nBatch x k x 3 x 3) tensor with 3x3 rotation matrices.
     """
+    if torch.jit.is_scripting():
+        return torch_quaternion.to_rotation_matrix(q, eps=eps)
     backend = resolve_backend(
         backend,
         q,
@@ -329,9 +333,11 @@ def from_rotation_matrix(
 
     :parameter matrices: A tensor of shape (..., 3, 3) representing the rotation matrices.
     :parameter eta: Numerical precision threshold (unused, kept for compatibility).
-    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch.
+    :parameter backend: Backend to use. "auto" opts into the experimental Triton backend on CUDA float32, else torch. Scripted calls always use the Torch backend.
     :return: A tensor of shape (..., 4) representing the quaternions in ((x, y, z), w) format.
     """
+    if torch.jit.is_scripting():
+        return torch_quaternion.from_rotation_matrix(matrices, eta)
     backend = resolve_backend(
         backend,
         matrices,

pymomentum/test/test_quaternion.py

@@ -5,6 +5,7 @@
 
 # pyre-strict
 
+import io
 import math
 import unittest
 
@@ -26,6 +27,43 @@ def generateRandomQuats(sz: int) -> torch.Tensor:
 
 
 class TestQuaternion(unittest.TestCase):
+    def test_rotation_conversions_are_scriptable(self) -> None:
+        quaternions = quaternion.normalize(torch.randn(4, 4))
+        matrices = quaternion.to_rotation_matrix(quaternions)
+
+        scripted_to_matrix_assume_normalized = torch.jit.script(
+            quaternion.to_rotation_matrix_assume_normalized
+        )
+        scripted_to_matrix = torch.jit.script(quaternion.to_rotation_matrix)
+        scripted_from_matrix = torch.jit.script(quaternion.from_rotation_matrix)
+
+        self.assertTrue(
+            torch.allclose(
+                scripted_to_matrix_assume_normalized(quaternions, "auto"),
+                quaternion.to_rotation_matrix_assume_normalized(
+                    quaternions, backend="auto"
+                ),
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                scripted_to_matrix(quaternions, "auto"),
+                quaternion.to_rotation_matrix(quaternions, backend="auto"),
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                scripted_from_matrix(matrices, 1e-6, "auto"),
+                quaternion.from_rotation_matrix(matrices, backend="auto"),
+            )
+        )
+        for scripted_conversion in (
+            scripted_to_matrix_assume_normalized,
+            scripted_to_matrix,
+            scripted_from_matrix,
+        ):
+            torch.jit.save(scripted_conversion, io.BytesIO())
+
     def test_euler_conversion(self) -> None:
         torch.manual_seed(0)  # ensure repeatability
         nBatch = 6