code: def test_segment_mean(self, nan_data):
      data = jnp.arange(9, dtype=jnp.float32)
      expected_out = jnp.array([4, 14 / 3.0, 2, 0, 4, 0])
      segment_ids = jnp.array([0, 1, 2, 0, 4, 0, 1, 1, 0])
      if nan_data:
          data = data.at[0].set(jnp.nan)
          expected_out = expected_out.at[segment_ids[0]].set(jnp.nan)
      result = utils.segment_mean(data, segment_ids, 6)
      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 18 for calling the GraphTest obj's underlying member method code with example usage: obj.test_segment_mean(nan_data) without return types