code: def testSoftmaxGradResiduals(self):
      if not config.softmax_custom_jvp.value:
          raise unittest.SkipTest(
              "only applies when upgrade flag enabled"
          )
      x = jnp.array([5.5, 1.3, -4.2, 0.9])
      res = ad_checkpoint.saved_residuals(nn.softmax, x)
      self.assertLen(res, 1)
  
imports: import collections
  from functools import partial
  import itertools
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import scipy.stats
  from jax._src import config
  from jax._src import core
  from jax._src import test_util
  from jax._src import ad_checkpoint
  from jax.test_util import check_grads
  from jax import nn
  from jax import random
  import jax
  import jax.numpy
label_desc: Use image node_id 26 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testSoftmaxGradResiduals() without return types