code: def testSoftmaxWhereMask(self, fn):
      x = jnp.array([5.5, 1.3, -4.2, 0.9])
      m = jnp.array([True, False, True, True])
  
      out = fn(x, where=m, initial=-jnp.inf)
      self.assertAllClose(out[m], fn(x[m]))
  
      probs = out if fn is nn.softmax else jnp.exp(out)
      self.assertAllClose(probs.sum(), 1.0)
  
      # TODO(mattjj): include log_softmax in these extra tests if/when we add a
      # custom_jvp rule for it (since otherwise it doesn't pass the numerical
      # checks below).
      if fn is nn.softmax and config.softmax_custom_jvp.value:
          g_fun = lambda x: jnp.take(
              fn(x, where=m, initial=-jnp.inf), jnp.array([0, 2, 3])
          )
          jtu.check_grads(g_fun, (x,), order=2)
  
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 24 for calling the NNFunctionsTest obj's underlying member method code with example usage: obj.testSoftmaxWhereMask(fn) without return types