code: def partition_softmax(
      logits: ArrayTree,
      partitions: jnp.ndarray,
      sum_partitions: Optional[int] = None,
  ):
      """Compute a softmax within partitions of an array.
  
        For example::
          logits = jnp.ndarray([1.0, 2.0, 3.0, 1.0, 2.0])
          partitions = jnp.ndarray([3, 2])
          partition_softmax(node_logits, n_node)
          >> DeviceArray(
          >> [0.09003057, 0.24472848, 0.66524094, 0.26894142, 0.7310586],
          >> dtype=float32)
  
      Args:
        logits: the logits for the softmax.
        partitions: the number of nodes per graph. It is a vector of integers with
          shape ``[n_graphs]``, such that ``graph.n_node[i]`` is the number of nodes
          in the i-th graph.
        sum_partitions: the sum of n_node. If not passed, the result of this method
          is data dependent and so not ``jit``-able.
  
      Returns:
        The softmax over partitions.
      """
      n_partitions = len(partitions)
      segment_ids = jnp.repeat(
          jnp.arange(n_partitions),
          partitions,
          axis=0,
          total_repeat_length=sum_partitions,
      )
      return segment_softmax(
          logits, segment_ids, n_partitions, indices_are_sorted=True
      )
  
imports: import functools
  from typing import Any, Callable, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Union
  import jax
  from jax import lax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  import numpy
label_desc: Use image node_id 11 for calling a global function with example usage: partition_softmax(logits, partitions, sum_partitions) and returns: segment_softmax