Feat/Split Operator #2490
Feat/Split Operator #2490
Conversation
…izes()`, add some clarity vs `chunk()`
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2490 +/- ##
==========================================
- Coverage 82.93% 82.86% -0.08%
==========================================
Files 815 817 +2
Lines 105344 105853 +509
==========================================
+ Hits 87371 87714 +343
- Misses 17973 18139 +166 ☔ View full report in Codecov by Sentry. 🚨 Try these New Features:
|
| @@ -4,6 +4,7 @@ mod tests { | |||
| use alloc::vec::Vec; | |||
| use burn_tensor::{Int, Shape, Tensor, TensorData}; | |||
|
|
|||
| #[test] | |||
There was a problem hiding this comment.
I think this was accidentally left out in #998
| @@ -107,6 +107,7 @@ macro_rules! testgen_quantization { | |||
| burn_tensor::testgen_q_sin!(); | |||
| burn_tensor::testgen_q_slice!(); | |||
| burn_tensor::testgen_q_sort_argsort!(); | |||
| // burn_tensor::testgen_q_split!(); | |||
There was a problem hiding this comment.
Adding quantized tests is not well documented at the moment, so you're not expected to do it.
@laggui will work on making it more straightforward and he'll be able to adapt your non-quantized tests.
| pub fn split<B: Backend, K: TensorKind<B> + BasicOps<B>>( | ||
| tensor: K::Primitive, | ||
| split_size: usize, | ||
| dim: usize, | ||
| ) -> Vec<K::Primitive> { | ||
| let size = K::shape(&tensor).dims[dim]; | ||
| let mut tensors = Vec::new(); | ||
|
|
||
| let mut start = 0; | ||
| while start < size { | ||
| let length = usize::min(split_size, size - start); | ||
| tensors.push(narrow::<B, K>(tensor.clone(), dim, start, length)); | ||
| start += length; | ||
| } | ||
|
|
||
| tensors | ||
| } |
There was a problem hiding this comment.
I think this function will likely by the fastest in most cases, since there is not kernel to execute, only metadata to update.
| /// | ||
| /// To split a tensor, users should prefer the [Tensor::split](Tensor::split) function, | ||
| /// which is more high-level and designed for public use. | ||
| fn split(tensor: Self::Primitive, split_size: usize, dim: usize) -> Vec<Self::Primitive>; |
There was a problem hiding this comment.
What's the difference between this version of split and chunk? Is it only that split takes the size of the new chunks while chunk takes the number of chunks?
There was a problem hiding this comment.
Correct- and then split_with_sizes goes a step further and lets you specify the size of each chunk.
| @@ -107,6 +107,7 @@ macro_rules! testgen_quantization { | |||
| burn_tensor::testgen_q_sin!(); | |||
| burn_tensor::testgen_q_slice!(); | |||
| burn_tensor::testgen_q_sort_argsort!(); | |||
| // burn_tensor::testgen_q_split!(); | |||
There was a problem hiding this comment.
Adding quantized tests is not well documented at the moment, so you're not expected to do it.
@laggui will work on making it more straightforward and he'll be able to adapt your non-quantized tests.
|
Thanks a lot @agelas for your implementation! And thank you, @louisfd and @nathanielsimard for your reviews! |
Pull Request Template
Checklist
run-checks allscript has been executed.Related Issues/PRs
#2440
Changes
Adds support for the split operation. The ONNX -> Burn conversion will be part of a separate PR.
Testing
Added tests under
burn-tensor.