I have a one-dimensional tensor X whose components are either 0 or 1. Let’s say that this tensor has n elements.
On the other hand, i have a 2D list/numpy array N whose rows dont have the same lenght.
For example, N = [[1,4,8], [2,5],[0]…]. In particular, this array represents the structure of a given graph G, so that row i corresponds to the neigbhors nodes of vertex i.
What i want is a way of simultaneously (on my GPU) index the 1D tensor X with each row of N, so i can do operations like computing the number of neigbhors whose X value is 1, for each node i.
Long story short, is there a way of doing: torch.sum(X[N[i]]) for each row i of N, and obtain a tensor like this: B = [torch.sum(X[N[0]), torch.sum(X[N[1]), torch.sum(X[N[2]]), …] of lenght n?