Add `chunk` Tensor Method

[dcvz]

Feature description

Add a chunk method on tensor to faciliate splitting a tensor into a given number of chunks

Feature motivation

Its a fairly common operation supported by other ML frameworks.

Actions

• [x] Default method implementation (@dcvz)
• [ ] Backend implementation (@Kelvinyu1117)

[Kelvinyu1117]

I would like to work on this issue, can you assign it to me?

[dcvz]

I would like to work on this issue, can you assign it to me?

Nice! Let me know if you need any guidance. You can also find me in the Discord as @dcvz.

[antimora]

I would like to work on this issue, can you assign it to me?

Done

[Kelvinyu1117]

I just had a brief look at the source code, the file I should modify is src/tensor/api/base.rs right? And the interface is as follows:

pub fn chunk<const D2: usize>(
    Tensor<B, D2, K>& tensor, 
    chunks: usize, 
    dim: usize = 0) -> Vec<Tensor<B, D2, K>>
{
}

[dcvz]

I just had a brief look at the source code, the file I should modify is src/tensor/api/base.rs right? And the interface is as follows:

Correct! I think interface can likely be the following, if you do the chunking in base space:

pub fn chunk(self, chunks: usize, dim: usize) -> Vec<Self>

Although given that some of the backends already have their own implementations, it might be better to do this backend specific instead of general (depends on what optimizations they have). What do you say @nathanielsimard ?

[Kelvinyu1117]

I think leveraging the implementation of the backend is a good idea!

[antimora]

@dcvz can this issue be closed as the main PR is merged?

[Kelvinyu1117]

I'm working on the backend implementation, we can use this issue or open another issue as a tracker.