Merge pull request #23652 from flferretti:feature/scalar_first_scipy_…

…spatial

PiperOrigin-RevId: 675286624

jax/_src/scipy/spatial/transform.py

@@ -167,12 +167,12 @@ def as_rotvec(self, degrees: bool = False) -> jax.Array:
     """Represent as rotation vectors."""
     return _as_rotvec(self.quat, degrees)
 
-  def as_quat(self, canonical: bool=False) -> jax.Array:
+  def as_quat(self, canonical: bool=False, scalar_first: bool=False) -> jax.Array:
     """Represent as quaternions."""
-    if canonical:
-      return _make_canonical(self.quat)
-    else:
-      return self.quat
+    quat = _make_canonical(self.quat) if canonical else self.quat
+    if scalar_first:
+        return jnp.roll(quat, shift=1, axis=-1)
+    return quat
 
   def inv(self):
     """Invert this rotation."""

tests/scipy_spatial_test.py

@@ -132,6 +132,20 @@ def testRotationAsQuatCanonical(self, shape, dtype):
     self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
     self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
 
+  @jtu.sample_product(
+    dtype=float_dtypes,
+    shape=[(4,), (num_samples, 4)],
+  )
+  def testRotationAsQuatScalarFirst(self, shape, dtype):
+    if scipy_version < (1, 14, 0):
+      self.skipTest("Scipy 1.14.0 added the `scalar_first` arg.")
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: (rng(shape, dtype),)
+    jnp_fn = lambda q: jsp_Rotation.from_quat(q).as_quat(scalar_first=True)
+    np_fn = lambda q: osp_Rotation.from_quat(q).as_quat(scalar_first=True).astype(dtype)
+    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)
+    self._CompileAndCheck(jnp_fn, args_maker, atol=1e-4)
+
   @jtu.sample_product(
     dtype=float_dtypes,
     shape=[(num_samples, 4)],