code: def test_segment_variance(self, nan_data):
      data = jnp.arange(8, dtype=jnp.float32)
      expected_out = jnp.stack(
          [
              jnp.var(jnp.arange(3)),
              jnp.var(jnp.arange(3, 5)),
              jnp.var(jnp.arange(5, 8)),
          ]
      )
      segment_ids = jnp.array([0, 0, 0, 1, 1, 2, 2, 2])
      if nan_data:
          data = data.at[0].set(jnp.nan)
          expected_out = expected_out.at[segment_ids[0]].set(jnp.nan)
      result = utils.segment_variance(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 19 for calling the GraphTest obj's underlying member method code with example usage: obj.test_segment_variance(nan_data) without return types