code: def test_partition_softmax(self):
      data = jnp.arange(9)
      partitions = jnp.array([3, 2, 4])
      sum_partitions = 9
      expected_out = np.array(
          [
              0.090031,
              0.244728,
              0.665241,
              0.268941,
              0.731059,
              0.032059,
              0.087144,
              0.236883,
              0.643914,
          ]
      )
      with self.subTest("nojit"):
          result = utils.partition_softmax(
              data, partitions, sum_partitions
          )
          jax.tree_util.tree_map(
              functools.partial(
                  np.testing.assert_allclose, atol=1e-5, rtol=1e-5
              ),
              result,
              expected_out,
          )
          result = utils.partition_softmax(data, partitions)
          jax.tree_util.tree_map(
              functools.partial(
                  np.testing.assert_allclose, atol=1e-5, rtol=1e-5
              ),
              result,
              expected_out,
          )
      with self.subTest("jit"):
          result = jax.jit(utils.partition_softmax, static_argnums=2)(
              data, partitions, sum_partitions
          )
          jax.tree_util.tree_map(
              functools.partial(
                  np.testing.assert_allclose, atol=1e-5, rtol=1e-5
              ),
              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 28 for calling the GraphTest obj's underlying member method code with example usage: obj.test_partition_softmax() without return types