Hi I am trying to move a TF code to PyTorch. But I am unable to replicate one functionality into Torch. I am pasting the two lines from tensorflow here
neighbourRagged = tf.RaggedTensor.from_row_splits(x, row_splits=rowSplits)
neighbourFeatures = tf.math.reduce_max(neighbourRagged, axis=1)

where x is a huge float array of size (257997,64) and row_splits is an int32 array of size (32769,).
I am trying to implement these two lines from tf to torch. I am aware that Torch does not support RaggedTensor  yet. But is there any work-around that to “specifically” implement these two lines?
I have uploaded the exact .npy files(zipped) on drive which can be accessed from here 1
The x and rowSplits array can be accessed from “files” folder after extracting zip file.
x = np.load("x.npy")
rowSplits = np.load("numNeighboursArray.npy")

If you are able to replicate the functionality please do let me know!
Thank you in advance!!
Tagging @ptrblck