code: def test_segment_softmax(self, nan_data):
      data = jnp.arange(9, dtype=jnp.float32)
      segment_ids = jnp.array([0, 1, 2, 0, 4, 0, 1, 1, 0])
      num_segments = 6
      expected_out = jnp.array(
          [
              3.1741429e-04,
              1.8088353e-03,
              1.0000000e00,
              6.3754367e-03,
              1.0000000e00,
              4.7108460e-02,
              2.6845494e-01,
              7.2973621e-01,
              9.4619870e-01,
          ]
      )
      if nan_data:
          data = data.at[0].set(jnp.nan)
          expected_out = expected_out.at[np.array([0, 3, 5, 8])].set(
              jnp.nan
          )
      with self.subTest("nojit"):
          result = utils.segment_softmax(
              data, segment_ids, num_segments
          )
          np.testing.assert_allclose(result, expected_out)
          result = utils.segment_softmax(data, segment_ids)
          np.testing.assert_allclose(result, expected_out)
      with self.subTest("jit"):
          result = jax.jit(utils.segment_softmax, static_argnums=2)(
              data, segment_ids, num_segments
          )
          np.testing.assert_allclose(result, expected_out)
  
imports: import functools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax.lib import xla_bridge
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 27 for calling the GraphTest obj's underlying member method code with example usage: obj.test_segment_softmax(nan_data) without return types