Hi I found one workaround based on @J_Johnson’s suggestion!
I am showing the demo numpy implementation for the same.
x = np.random.randint(low=10,size=(10,3))
# array([[6, 0, 3],
       [0, 7, 0],
       [4, 6, 8],
       [1, 6, 7],
       [2, 6, 1],
       [8, 5, 5],
       [9, 8, 3],
       [4, 9, 1],
       [5, 8, 6],
       [6, 6, 3]])

splits = np.array([0,2,3,6,7,10])
splits_ = np.subtract(splits[1:],splits[:-1])
# splits_ = [2,1,3,1,3]
l = len(splits_) # l =5
m = max(splits_) # m = 3
neighbourRagged = np.zeros([l,m,3]) 
# this will be my output array where I can copy
# the input with zero padding

I set output's second axis length to the maximum split value( in this case 3). This is the maximum value of the second axis after zero-padding
Now for each block(among l blocks) each of size (m,3), I fill up the first i rows (obtained from each split value)   with the i==j-prev rows from input array and rest rows stay at zero.
prev = 0
idxs = range(l)
for idx,i,j in zip(idxs,splits_, splits[1:]):
...     neighbourRagged[idx, :i,:]=x[prev:j,:]
...     prev=j

#neighbourRagged
array([[[6., 0., 3.],
        [0., 7., 0.],
        [0., 0., 0.]],

       [[4., 6., 8.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[1., 6., 7.],
        [2., 6., 1.],
        [8., 5., 5.]],

       [[9., 8., 3.],
        [0., 0., 0.],
        [0., 0., 0.]],

       [[4., 9., 1.],
        [5., 8., 6.],
        [6., 6., 3.]]])

Now to obtain the max output, I can take the maximum along axis=1
neighbourFeatures = neighbourRagged.max(axis=1)
#neighbourFeatures
array([[6., 7., 3.],
       [4., 6., 8.],
       [8., 6., 7.],
       [9., 8., 3.],
       [6., 9., 6.]])

This is the output I want.
I was wondering if we can implement this in PyTorch using some efficient function.(without using loops etc.)
@ptrblck