code: def testSoftmaxGradFlag(self):
      x = jnp.array([5.5, 1.3, -4.2, 0.9])
  
      with jax.softmax_custom_jvp(False):
          res = ad_checkpoint.saved_residuals(nn.softmax, x)
      self.assertLen(res, 3)
      self.assertEqual(sum(a.size for a, _ in res), 6)
  
      with jax.softmax_custom_jvp(True):
          res = ad_checkpoint.saved_residuals(nn.softmax, x)
      self.assertLen(res, 1)
      self.assertEqual(sum(a.size for a, _ in res), 4)
  
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 27 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testSoftmaxGradFlag() without return types