code: def test_segment_max_or_constant_2d(
      self, indices_are_sorted, unique_indices
  ):
      if unique_indices:
          data = jnp.stack(
              [jnp.arange(6), jnp.arange(6, 0, -1)], axis=1
          )
          if indices_are_sorted:
              segment_ids = jnp.array([0, 1, 2, 3, 4, 5])
              expected_out = jnp.array(
                  [[0, 6], [1, 5], [2, 4], [3, 3], [4, 2], [5, 1]]
              )
              num_segments = 6
          else:
              segment_ids = jnp.array([1, 0, 2, 4, 3, -5])
              expected_out = jnp.array(
                  [[1, 5], [0, 6], [2, 4], [4, 2], [3, 3]]
              )
              num_segments = 5
      else:
          data = jnp.stack(
              [jnp.arange(9), jnp.arange(9, 0, -1)], axis=1
          )
          if indices_are_sorted:
              segment_ids = jnp.array([0, 0, 0, 1, 1, 1, 2, 3, 4])
              expected_out = jnp.array(
                  [[2, 9], [5, 6], [6, 3], [7, 2], [8, 1], [0, 0]]
              )
          else:
              segment_ids = jnp.array([0, 1, 2, 0, 4, 0, 1, 1, -6])
              expected_out = jnp.array(
                  [[5, 9], [7, 8], [2, 7], [0, 0], [4, 5], [0, 0]]
              )
          num_segments = 6
  
      with self.subTest("nojit"):
          result = utils.segment_max_or_constant(
              data,
              segment_ids,
              num_segments,
              indices_are_sorted,
              unique_indices,
          )
          np.testing.assert_allclose(result, expected_out)
          result = utils.segment_max_or_constant(
              data,
              segment_ids,
              indices_are_sorted=indices_are_sorted,
              unique_indices=unique_indices,
          )
          num_unique_segments = jnp.max(segment_ids) + 1
          np.testing.assert_allclose(
              result, expected_out[:num_unique_segments]
          )
      with self.subTest("jit"):
          result = jax.jit(
              utils.segment_max_or_constant, static_argnums=(2, 3, 4)
          )(
              data,
              segment_ids,
              num_segments,
              indices_are_sorted,
              unique_indices,
          )
          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 23 for calling the GraphTest obj's underlying member method code with example usage: obj.test_segment_max_or_constant_2d(indices_are_sorted, unique_indices) without return types