[FIXED] Add padding based on partial sum

Issue

I have four given variables:

• group size
• total of groups
• partial sum
• 1-D tensor and I want to add zeros when the sum within a group reached the partial sum. For example:

groupsize = 4
totalgroups = 3
partialsum = 15
d1tensor = torch.tensor([ 3, 12,  5,  5,  5,  4, 11])

The expected result is: [ 3, 12, 0, 0, 5, 5, 5, 0, 4, 11, 0, 0]

I have no clue how can I achieve that in pure pytorch. In python it would be something like this:

target = [0]*(groupsize*totalgroups)

cursor = 0
current_count = 0
d1tensor = [ 3, 12,  5,  5,  5,  4, 11]

for idx, ele in enumerate(target):
  subgroup_start = (idx//groupsize) *groupsize
  subgroup_end = subgroup_start + groupsize 
  if sum(target[subgroup_start:subgroup_end]) < partialsum: 
     target[idx] = d1tensor[cursor]
     cursor +=1

Can anyone help me with that? I have already googled it but couldn't find anything.

Solution

Some logic, Numpy and list comprehensions are sufficient here. I will break it down step by step, you can make it slimmer and prettier afterwards:

import numpy as np

my_val = 15
block_size = 4
total_groups = 3
d1 = [3, 12,  5,  5,  5,  4, 11]
d2 = np.cumsum(d1)
d3 = d2 % my_val == 0 #find where sum of elements is 15 or multiple
split_points= [i+1 for i, x in enumerate(d3) if x] # find index where cumsum == my_val
#### Option 1
split_array = np.split(d1, split_points, axis=0)
padded_arrays = [np.pad(array, (0, block_size - len(array)), mode='constant') for array in split_array] #pad arrays
padded_d1 = np.concatenate(padded_arrays[:total_groups]) #put them together, discard extra group if present
#### Option 2
split_points = [el for el in split_points if el <len(d1)] #make sure we are not splitting on the last element of d1
split_array = np.split(d1, split_points, axis=0)
padded_arrays = [np.pad(array, (0, block_size - len(array)), mode='constant') for array in split_array] #pad arrays
padded_d1 = np.concatenate(padded_arrays)

Answered By - CAPSLOCK