code: def testCustomJVP(self):
      square = jax.custom_derivatives.custom_jvp(lambda x: x**2)
      square.defjvp(lambda p, t: (p[0] ** 2, 2 * t[0] * p[0]))
      x = BCOO.fromdense(jnp.arange(5.0))
  
      # Test calling the function itself.
      result = self.sparsify(square)(x)
      expected = self.sparsify(lambda x: x**2)(x)
      self.assertArraysEqual(result.indices, expected.indices)
      self.assertArraysAllClose(result.data, expected.data)
  
      # Test evaluating the custom gradient.
      grad_square_sum = jax.grad(lambda x: square(x).sum())
      result = self.sparsify(grad_square_sum)(x)
      expected = self.sparsify(jax.grad(lambda x: jnp.sum(x**2)))(x)
      self.assertArraysEqual(result.indices, expected.indices)
      self.assertArraysAllClose(result.data, expected.data)
  
imports: from functools import partial
  import math
  import operator
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import config, jit, lax
  import jax.numpy
  import jax._src.test_util
  from jax.experimental.sparse import BCOO, BCSR, sparsify, todense, SparseTracer
  from jax.experimental.sparse.transform import arrays_to_spvalues, spvalues_to_arrays, sparsify_raw, SparsifyValue, SparsifyEnv
  from jax.experimental.sparse.util import CuSparseEfficiencyWarning
  from jax.experimental.sparse import test_util
label_desc: Use image node_id 35 for calling the SparsifyTest obj's underlying member method code with example usage: obj.testCustomJVP() without return types