how is an arbitrary dimension "reshaper" turned into a function given a linear index?
If the index of $(i,j,k)$ is $p = i + xj + xyk$ then you can compute $i,j,k$ given $p$ as follows:

i = p mod x
j = (p / x) mod y
k = (p / x) / y


Using this you can solve your problem.