code: def selective_sum(
      values: types.Tensor, partitions: types.Int, num_partitions: int
  ) -> types.Tensor:
      """Sums entries in `values`, partitioned using `partitions`.
  
      For example,
  
      ```python
         # Returns `[0 + 4 + 5, 2 + 3 + 4]` i.e. `[9, 6]`.
         selective_sum(values=[0, 1, 2, 3, 4, 5],
                       partitions=[0, 1, 1, 1, 0, 0]),
                       num_partitions=2)
      ```
  
      Args:
        values: a `Tensor` with numerical type.
        partitions: an integer `Tensor` with the same shape as `values`. Entry
          `partitions[i]` indicates the partition to which `values[i]` belongs.
        num_partitions: the number of partitions. All values in `partitions` must
          lie in `[0, num_partitions)`.
  
      Returns:
        A vector of size `num_partitions` with the same dtype as `values`. Entry `i`
        is the sum of all entries in `values` belonging to partition `i`.
      """
      partitioned_values = tf.dynamic_partition(
          values, partitions, num_partitions
      )
      return tf.stack(
          [tf.reduce_sum(partition) for partition in partitioned_values]
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.policies import categorical_policy
  from tf_agents.policies import utils
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 1 for calling a global function with example usage: selective_sum(values, partitions, num_partitions) and returns: tf