code: def test_gradients(self):
      l2_loss = self.variant(losses.l2_loss)
      # Compute gradient in batch
      batch_grad_func = jax.vmap(jax.grad(l2_loss), (0))
      actual = batch_grad_func(self.xs)
      np.testing.assert_allclose(actual, self.dys)
  
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import losses
label_desc: Use image node_id 5 for calling the L2LossTest obj's underlying member method code with example usage: obj.test_gradients() without return types