tf_agents.bandits.agents.exp3_agent.selective_sum

Sums entries in values, partitioned using partitions.

tf_agents.bandits.agents.exp3_agent.selective_sum(
    values: tf_agents.typing.types.Tensor,
    partitions: tf_agents.typing.types.Int,
    num_partitions: int
) -> tf_agents.typing.types.Tensor

For example,

   # 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)

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).

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.