How to convert a (M, L) tensor to (N, L) based on counts vector of size (N) where M is sum of counts, using aggregation by adding

user_10854970 0
Viewed 30

So I have a 2D tensor A of shape (M, L) and I want to convert it into B, a (N, L) tensor.
I also have a counts tensor C (N) which has the counts of how many rows belong to which group such that sum(C) = M.

For example :

# shape = (6, 3)
A = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [1, 1, 1],
                  [5, 3, 1],
                  [5, 7, 1],
                  [2, 1, 3]])

# counts
C = torch.tensor([2, 1, 3])
# torch.sum(C) == A.shape[0]    

# final tensor shape = (3, 3)
B = torch.tensor([[5, 7, 9],
                  [1, 1, 1], 
                  [12, 11, 5]])

The aggregation of rows is done by element wise addition.

I tried to create a simple function for this as follows :

def convert_to_batch_results(self, results, counts):
  knowledge_cumsum = torch.cumsum(results, dim=0) # [M, L]
  inds = torch.cumsum(counts, dim=0) - 1 # [N]
  knowledge = knowledge_cumsum[inds] # [N, L]
  diff = torch.zeros(knowledge.shape) # [N, L]
  diff[1:] = knowledge[:-1]
  user_knowledge = torch.sub(knowledge, diff) # [N, L]
  return user_knowledge

This works when counts vector has all elements non-zero.
But if counts have some 0 elements the summation becomes wrong.
I case counts have 0s, I want the corresponding output rows in B to have value 0s.

This is what I changed :


def convert_to_batch_results(results, counts):
  knowledge_cumsum = torch.cumsum(results, dim=0) # [M, L]
  inds = torch.cumsum(counts, dim=0) - 1 # [N] # print(inds)
  knowledge = knowledge_cumsum[inds] # [N, L]
  knowledge = torch.mul(knowledge, torch.sign(counts.unsqueeze(1)))
  diff = torch.zeros(knowledge.shape).to('cpu') # [N, L]
  diff[1:] = knowledge[:-1]
  user_knowledge = torch.sub(knowledge, diff) # [N, L]
  user_knowledge = torch.mul(user_knowledge, torch.sign(counts.unsqueeze(1)))
  return user_knowledge

And ran it like so :

res = torch.ones(5, 3)
counts = torch.tensor([2, 0, 3])
print(convert_to_batch_results(res, counts))

# Output :
# tensor([[2., 2., 2.],
#        [-0., -0., -0.],
#        [5., 5., 5.]])

# Expected Output :
# tensor([[2., 2., 2.],
#        [-0., -0., -0.],
#        [3., 3., 3.]])

I have tried some other things but am unable to get the correct results in case of zeros present in counts. Please help me with the correct way to achieve the desired results.

Just to reiterate, counts can have zeros anywhere multiple times like [0, 3, 4], [0, 3, 0, 0, 1], [4, 0, 6, 0 1], etc.

Edit :
Adding the solution here by modifying my first approach and Shai's answer below.

The problem was with leading zeros in count which were becoming -1 and used as index into cumsum giving us the last row. Fixed this by using a mask.

def convert_to_batch_results(self, A, C):
  I = C.cumsum(dim=0)-1
  I = torch.mul(A.cumsum(dim=0)[I, :], (I != -1).int().unsqueeze(1)) # here
  I = torch.cat((torch.zeros_like(A[:1,:]), I), dim=0)
  return I.diff(dim=0)
pytorch python
1 Answers

You should use torch.cumsum to solve your problem:
Your output is simply the cumulative sum of rows up to the row induces of C. Taking the diff of the cumulative sum from the beginning of the tensor will give you the sum over the intervals you want:

B = torch.cat((torch.zeros_like(A[:1,:]), A.cumsum(dim=0)[C.cumsum(dim=0)-1, :]), dim=0).diff(dim=0)

Works with zeros in C "out-of-the-box".

We can break this one-liner to better understand what is going on:

A.cumsum(dim=0)  # cumulative sum of each column of A from the very beginning
C.cumsum(dim=0)  # (2, 3, 6) - At what rows we want to look at the cumulative sum: the second row will give us the first row of B. The third one would be the second row of B _but_ including the first one, etc.
diff(dim=0)  # subtrat the cumulative rows to get the intervals we want of A into B.

This trick is a simplified 1D version of integral images.

user_17144100
Related