class_docstr: Tests for LAX-backed scipy.spatial implementations
class_docstr_tok: ["Tests","for","LAX-backed","scipy.spatial","implementations"]
code: def testRotationMean(self, shape, dtype, rng_weights):
      rng = jtu.rand_default(self.rng())
      args_maker = lambda: (
          rng(shape, dtype),
          jnp.abs(rng(shape[0], dtype)) if rng_weights else None,
      )
      jnp_fn = (
          lambda q, w: jsp_Rotation.from_quat(q).mean(w).as_rotvec()
      )
      np_fn = (
          lambda q, w: osp_Rotation.from_quat(q)
          .mean(w)
          .as_rotvec()
          .astype(dtype)
      )  # HACK
      tol = 5e-3  # 1e-4 too tight for TF32
      self._CheckAgainstNumpy(
          np_fn, jnp_fn, args_maker, check_dtypes=True, tol=tol
      )
      self._CompileAndCheck(jnp_fn, args_maker, tol=tol)
  
imports: from absl.testing import absltest
  import jax
  import scipy.version
  from jax._src import test_util
  from jax.scipy.spatial.transform import Rotation
  from scipy.spatial.transform import Rotation
  from jax.scipy.spatial.transform import Slerp
  from scipy.spatial.transform import Slerp
  import jax.numpy
  import numpy
  from jax import config
label_desc: Use image node_id 15 for calling the LaxBackedScipySpatialTransformTests obj's underlying member method code with example usage: obj.testRotationMean(shape, dtype, rng_weights) without return types