mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: change random matrix generation for Rotation

Browse Source
This commit is contained in:
Jake VanderPlas 2025-01-31 09:21:46 -08:00
parent bf9671731c
commit 4d433d063e
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tests/scipy_spatial_test.py
View File

@ -196,7 +196,10 @@ class LaxBackedScipySpatialTransformTests(jtu.JaxTestCase):
)
def testRotationFromMatrix(self, shape, dtype):
rng = jtu.rand_default(self.rng())
args_maker = lambda: (rng(shape, dtype),)
def args_maker():
# Use QR to ensure valid positive-definite rotation matrix.
q, _ = onp.linalg.qr(rng(shape, dtype))
return [q]
jnp_fn = lambda m: jsp_Rotation.from_matrix(m).as_rotvec()
np_fn = lambda m: osp_Rotation.from_matrix(m).as_rotvec().astype(dtype)
self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=True, tol=1e-4)