search

kokkos / mdspan

Reference implementation of mdspan targeting C++23
Other
413 stars 69 forks source link

Reshaping mdspans #354

Closed headmyshoulder closed 3 months ago

headmyshoulder commented 3 months ago

Is there a way of reshaping a mdspan?

I'm looking for the equivalent of the following numpy code:

arr = np.zeros([10, 5, 3])
arr2 = arr[2:8, :, :2]
arr3 = arr2.reshape(6 * 5, 2)

Does the proposal support that? Its one of the basic functionalities of numpy/pytorch and it would be great if that could be supported in mdspan too.

dalg24 commented 3 months ago

No there is not. mdspan semantics to not lend themselves well to this functionality. It is not always possible to change the shape without copying the data and numpy.reshape may return something that is not a view into the original md array. mdspan is non-owning so it would not be able to do that.

headmyshoulder commented 3 months ago

@dalg24 Thank you for the details.

In the above example arr3 was still pointing into arr, but I understand that does not need to be the case in general.

Is it possible to have a generic reshape layout which wraps the input layout and applies the reshape mapping? It might not be optimal in terms of performance but might just work for most cases.

nmm0 commented 3 months ago

You might be able to add a custom layout mapping that applies the reshape operations. One thing is that you would pay the cost of your reshaping calculations each access as mdspan uses the mapping to calculate the offset into the underlying buffer.

One intriguing possibility is that if we have something like the mdspan copy proposal I'm working on in the standard, you could combine the two to create a reshape function that copies. It wouldn't be semantically the same as numpy reshape, since mdspans don't own memory, but could end up with something like reshape(src_mdspan, dst_mdspan)

headmyshoulder commented 3 months ago

Yes, I will give that a try. The custom layout mapping might solve my use case.

mhoemmen commented 3 months ago

  1. mdspan objects may have an arbitrary mix of compile-time and run-time extents. It's not possible to change compile-time extents at run time.

  2. mdspan objects are generally cheap to create, so the mdspan idiom for "reshape" would just create a new mdspan with different extents.

It sounds like the original issue is resolved, so I'm closing this issue. Please feel free to reopen or to create another issue if you have further concerns. Thanks!