Reshape in-place using Relay

I’m building some functions in Relay, which have reshape stages for some tensors. However, there are special cases where from a memory layout perspective the reshape operation is the identity.

E.g. I might have a volume A of shape [a, b, c], but I have an operation which reshapes it into B*, of shape [a, b, c//2, c//2].

In 1D memory A might look like: [1, 2, 3, 4, 5, ....].

For most cases of the reshape, it would return something different, e.g. B = [1, 2, 5, 6, 3, 4, ...]. However, in other cases B will still be identical: [1, 2, 3, 4, 5, ....].

I haven’t figured out a way to ensure that in this case the reshape is done in-place, i.e. there is no copy of A made. I could always skip the operation, however I need the indexing of B later in the algorithm.

Is there a way I can ensure this happens? Would it be a schedule problem? Or is it something I can specify elsewhere? I still see the reshape in the TVM IR.

2 Likes