I want to perform operations in last to dimensions.
e.g. on each (5, 5) subtensor.

Let’s say…
Mean of all 5 rows, mean of all 5 colums => 10 values
Max. of each row and column = > 10 values
Min. of each row and column = > 10 values
So, I have now 10+10+10= 30 values, which can be viewed as (5,6) subtensor.

So, I want to replace original (5, 5) subtensors by these (5, 6) subtensors such that the shape of that 4D tensor now becomes (100, 10, 5, 6) or maybe (100, 10, 30).

I want to avoid for loops because backward method is very slow in that case.

I was thinking to view it as a 3D tensor maybe and iterate through each (5, 5) sub tensors with for loop.
But that way would be extremely slow as I mentioned.

I has a similar use case where a tensor lets say A, A.shape = (100,10,5,5)
I wanted to add the last two dimensions such that new Ashape = (100,10,5) and following did the trick