Fix VonMises tests.

PiperOrigin-RevId: 485835879

Author: franrruiz

distrax/_src/distributions/von_mises_test.py

@@ -368,6 +368,7 @@ def f(c):
     # should not be zero.
     self.assertNotEqual(jax.grad(f)(0.), 0)
 
+  @chex.all_variants(with_pmap=False, with_device=False)
   @parameterized.named_parameters(
       ('small concentration', 1.),
       ('medium concentration', 10.),
@@ -382,29 +383,26 @@ def f(seed, l, c):
       vm = self.distrax_cls(l, c)
       x = vm.sample(seed=seed)  # pylint: disable=cell-var-from-loop
       return x
-    jax_sample_and_grad = jax.value_and_grad(f, argnums=2)
+    jax_sample_and_grad = self.variant(jax.value_and_grad(f, argnums=2))
 
     def samples_grad(s, concentration):
       broadcast_concentration = concentration
       _, dcdf_dconcentration = tfp.math.value_and_gradient(
           lambda conc: tfp.distributions.von_mises.von_mises_cdf(s, conc),
           broadcast_concentration)
       inv_prob = np.exp(-concentration * (np.cos(s) - 1.)) * (
-          (2. * np.pi) * scipy.special.i0e(concentration)
-      )
+          (2. * np.pi) * scipy.special.i0e(concentration))
       # Computes the implicit derivative,
       # dz = dconc * -(dF(z; conc) / dconc) / p(z; conc)
       dsamples = -dcdf_dconcentration * inv_prob
       return dsamples
 
     for seed in range(10):
       sample, sample_grad = jax_sample_and_grad(
-          seed, jnp.array(locs), jnp.array(concentration)
-      )
+          seed, jnp.array(locs), jnp.array(concentration))
       comparison = samples_grad(sample, concentration)
       np.testing.assert_allclose(
-          comparison, sample_grad, rtol=1e-06, atol=1e-06
-      )
+          comparison, sample_grad, rtol=1e-06, atol=1e-06)
 
   def test_von_mises_sample_moments(self):
     locs_v = np.array([-1., 0.3, 2.3])