How to average tensor elements via an index vector

1

Hi,

I have an index matrix that is of the form

idx = [[-1,2,3,4],                   # Agent 1
          [0, -1, 2,3],                 # Agent 2
          [1, 1, -1, 2]                 # Agent 3
          [1, 0, 0, -1]                 # Agent 4
]```

Now, I have a bunch of tensors (shape is (#Agents, hidden_layer_size)).  I need a function that averages the agents based on the index for every agent.

So, for example, for agent #2, I need to average hidden layer values from 1 and 2.  -1 represents itself and I won't pool its layer and agent #4 will be as is in #2.  I have the following pooled output tensor of size (#Agents, #Sectors, hidden_layer_size). I have used loops and a dictionary to do them but how would I go about doing this?

Thank you.
2

Did you figure out a way to do this?