code: def test_segment_normalize(self, nan_data):
      norm = lambda x: (x - jnp.mean(x)) * jax.lax.rsqrt(jnp.var(x))
      data = jnp.arange(8, dtype=jnp.float32)
      segment_ids = jnp.array([0, 0, 0, 1, 1, 2, 2, 2])
      expected_out = jnp.concatenate(
          [
              norm(jnp.arange(3, dtype=jnp.float32)),
              norm(jnp.arange(3, 5, dtype=jnp.float32)),
              norm(jnp.arange(5, 8, dtype=jnp.float32)),
          ]
      )
      if nan_data:
          data = data.at[0].set(jnp.nan)
          expected_out = expected_out.at[:3].set(jnp.nan)
      result = utils.segment_normalize(data, segment_ids, 3)
      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 20 for calling the GraphTest obj's underlying member method code with example usage: obj.test_segment_normalize(nan_data) without return types