tf_agents.bandits.agents.exp3_agent.selective_sum

View source on GitHub

Sums entries in values, partitioned using partitions.

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.