[QST] What is the easiest way to partition/slice a tensor in CUTE to get a subtensor?

[srikantvv]

For example, if I have a tensorA (_256,_64):(_64,_1). I want to get a subtensor with the first two columns to end up with a shape tensorB (_2, _64):(_64,_1). Something with logical_divide seems to be the way, but it is haven't been able to figure out the easiest way.

[ccecka]

If you just want that "first" subtensor

Tensor my_2x64 = composition(my_256x64, Shape<_2,_64>{});

and if you want the slice off any of those subtensors

Tensor tmp = flat_divide(my_256x64, Shape<_2,_64>{});    // (_2,_64,_128,_1)
Tensor my_2x64 = tmp(_,_,my_idx0,my_idx1);

Similar to these examples: https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/03_tensor.md#examples

[irasin]

Hi, @ccecka It looks like that composition does not work for tensor

here is my test code

    {
        auto my_256x64 = make_tensor<float>(Shape<_256, _64>{}, LayoutRight{});
        Tensor my_2x64 = composition(my_256x64, Shape<_2,_64>{});
    }

and I got function not found error

[ccecka]

Ah, right. composition(Tensor, Layout) is intentionally not defined because it implicitly slices off the "first" subtensor. Use a divide to tile a tensor and make the slice explicit instead.

[irasin]

Thanks. But I was wondering why not remove the composition(Tensor, Layout) API from the documentation? This can be very confusing.

[ccecka]

I don't believe it is in the documentation? It should only appear erroneously in this thread.

[irasin]

While, it does exist. In https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/03_tensor.md#tiling-a-tensor

And in https://github.com/NVIDIA/cutlass/blob/main/media/docs/cute/03_tensor.md#thread-value-partitioning

Maybe we should not use the terminology composition in the tensor doc.

[ccecka]

I've very surprised that composition(Tensor, Layout) doesn't exist, even if it's just for testing and documentation/explanation.

I suppose it should be added (even though it is of little use in real applications).

[github-actions[bot]]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.