How does reshape moves entries?

Hello, I have a neural network with two input tensors:

a shape (batch_size=128, 3)
b shape (60, batch_size=128, 1)

I need to augment a to have the same shape as b:

a.repeat(60, 1, 1) -> a shape (60, 128, 3)

Then I want to put a and b in the same neural network that accepts two dimensions. I need to reshape so that the 60 can be considered the same as the batch size:

result = neural_network(a.reshape(-1, 3), b.reshape(-1, 1))

This result now has the same ((60 x 128), 1)

Now I want to do the reshape so that the result has the shape (60, 128, 1) and it has the same order than a and b had initially. If I do:

result_reshaped = result.reshape(60, 128, 1) I will get that, however, how can I be sure that the order is correct?


One way to double check the reshaping operations is to initialize the tensor via torch.arange, reshape it, and print the tensor to see where the values landed. To do so I would reduce the shapes a bit, as the currently used number of elements is quite large.