code: def test_gradients(self, expectile, predictions, expected_grads):
      predictions = jnp.array(predictions)
      expected_grads = jnp.array(expected_grads)
      targets = jnp.zeros_like(predictions)
      expectile_loss_fn = self.variant(
          functools.partial(losses.expectile_loss, expectile=expectile)
      )
      batch_grad_fn = jax.vmap(jax.grad(expectile_loss_fn), (0, 0))
      predicted_grads = batch_grad_fn(predictions, targets)
      np.testing.assert_allclose(
          predicted_grads, expected_grads, atol=1e-4
      )
  
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 2 for calling the ExpectileLossTest obj's underlying member method code with example usage: obj.test_gradients(expectile, predictions, expected_grads) without return types