code: def testSoftmaxGrad(self):
      x = jnp.array([5.5, 1.3, -4.2, 0.9])
      jtu.check_grads(nn.softmax, (x,), order=2, atol=5e-3)
  
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 25 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testSoftmaxGrad() without return types