Implement chunk for different backends #1032

Kelvinyu1117 · 2023-12-01T18:00:52Z

Pull Request Template

Checklist

Confirmed that run-checks all script has been executed.
Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Changes

This PR is in progress, mainly utilizing the backend implementation for chunk() instead of the default, currently only bool_chunk was implemented and encountered issues for burn-fusion.

Testing

As it is still in progress, it has not yet been tested.

Kelvinyu1117 · 2023-12-01T18:03:50Z

Currently, this PR has an error of not implementing the traits:

error[E0046]: not all trait items implemented, missing: `bool_chunk`
  --> burn-fusion/src/ops/boolean.rs:17:1
   |
17 | impl<B: FusionBackend> BoolTensorOps<Self> for Fusion<B> {
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ missing `bool_chunk` in implementation
   |
   = help: implement the missing item: `fn bool_chunk(_: <backend::Fusion<B> as burn_tensor::backend::Backend>::BoolTensorPrimitive<D>, _: usize, _: usize) -> std::vec::Vec<<backend::Fusion<B> as burn_tensor::backend::Backend>::BoolTensorPrimitive<D>> { todo!() }`

I think once we come up with the solution for burn-fusion, the implementation of other data types (Int/Float) will be trivial.

Kelvinyu1117 · 2023-12-02T16:52:33Z

The current solution is to have a default implementation on all the ops, such that the backend can fall back to the default.

nathanielsimard

I think you can remove the default function implementation of chunk and narrow in the Kind API. I think you should create files in burn_tensor for narrow and chunk where implementations will be located, which can reduce code duplication.

nathanielsimard · 2023-12-04T16:17:15Z

burn-candle/src/ops/base.rs

+    let tensor = tensor.tensor.narrow(dim, start, length);
+    match tensor {
+        Ok(tensor) => CandleTensor::new(tensor),
+        Err(e) => panic!("error chunk from Candle"),


error chunk => error narrow

nathanielsimard · 2023-12-04T16:19:29Z

burn-tensor/src/tensor/api/base.rs

+                if i == dim {
+                    start..(start + length)
+                } else {
+                    0..Self::shape(&tensor).dims[i]


I would extract the shape before the loop, just a bit cleaner imo.

nathanielsimard · 2023-12-04T16:23:58Z

burn-tensor/src/tensor/api/base.rs

+    fn chunk<const D: usize>(
+        tensor: Self::Primitive<D>,
+        chunks: usize,
+        dim: usize,
+    ) -> Vec<Self::Primitive<D>> {
+        let size = Self::shape(&tensor).dims[dim];


The default implementation should be in the Backend trait, not in the API Kind trait system.

So the change would be =>

Extract those functions into their own modules burn-tensor/src/ops/chunk.rs and burn-tensor/src/ops/narrow.rs. In the TensorOps, IntTensorOps and BoolTensorOps trait, we should call those functions by default, so backends are not forced to implement the function.

Just to confirm the interfaces, I presume they should be the following right?:

pub(crate) fn chunk<B: Backend, const D: usize, K: TensorKind>( tensor: Tensor<B, D, K>, chunks: usize, dim: usize, ) -> Vec<Tensor<B, D, K>> { } pub(crate) fn narrow<B: Backend, const D: usize, K: TensorKind>( tensor: Tensor<B, D, K>, dim: usize, start: usize, length: usize, ) -> Tensor<B, D, K> { }

@Kelvinyu1117 Yes sorry for the delay, I wasn't available last week 😅

The code is refactored. Please have a look, thanks.

louisfd

Thanks, it's good that we can leverage tch and candle implementations!
It appears we have worked a bit in parallel, see my other comment.

louisfd · 2023-12-15T18:32:09Z

burn-tensor/src/tensor/ops/tensor.rs

+    ///
+    /// A vectors of tensors
+    ///
+    fn chunk<const D: usize>(


It seems you have started your branch before #1006 was merged.
In #1006 we modified the default chunk algorithm --don't mind the name of the PR which is about fixing the docs, at first I thought the doc was wrong and the code was right, but in the end the doc was right and the code was wrong.
Could you please adjust the algorithm so it passes the test in #1006 (or in main as #1006 was merged) ?

Sure, let me work on that.

codecov · 2023-12-18T08:44:51Z

Codecov Report

Attention: 213 lines in your changes are missing coverage. Please review.

Comparison is base (4fb5b74) 85.63% compared to head (e1e9eee) 85.33%.
Report is 1 commits behind head on main.

Files	Patch %	Lines
burn-candle/src/ops/base.rs	0.00%	24 Missing ⚠️
burn-tch/src/ops/base.rs	0.00%	24 Missing ⚠️
burn-autodiff/src/ops/bool_tensor.rs	0.00%	15 Missing ⚠️
burn-autodiff/src/ops/int_tensor.rs	0.00%	15 Missing ⚠️
burn-candle/src/ops/bool_tensor.rs	0.00%	15 Missing ⚠️
burn-candle/src/ops/int_tensor.rs	0.00%	15 Missing ⚠️
burn-candle/src/ops/tensor.rs	0.00%	15 Missing ⚠️
burn-tch/src/ops/bool_tensor.rs	0.00%	15 Missing ⚠️
burn-tch/src/ops/int_tensor.rs	0.00%	15 Missing ⚠️
burn-tch/src/ops/tensor.rs	0.00%	15 Missing ⚠️
... and 3 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1032      +/-   ##
==========================================
- Coverage   85.63%   85.33%   -0.30%     
==========================================
  Files         507      509       +2     
  Lines       54074    54324     +250     
==========================================
+ Hits        46305    46357      +52     
- Misses       7769     7967     +198

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

louisfd · 2023-12-18T12:58:01Z

Thanks for the change, looks good to me now

nathanielsimard

Some smalls changes, but it's getting ready to be merged.

nathanielsimard · 2023-12-19T14:42:48Z

burn-tensor/src/tensor/api/base.rs

+    // fn chunk<const D: usize>(
+    //     tensor: Self::Primitive<D>,
+    //     chunks: usize,
+    //     dim: usize,
+    // ) -> Vec<Self::Primitive<D>> {
+    //     B::bool_chunk(tensor, chunks, dim)
+    // }


nathanielsimard · 2023-12-19T14:44:31Z

burn-tensor/src/tensor/api/chunk.rs

+/// Split the tensor along the given dimension into chunks.
+///
+/// # Arguments
+///
+/// * `tensor` - The tensor.
+/// * `chunks` - The number of chunks to be produced
+/// * `times` - The dimension along which the tensor will be split.
+///
+/// # Returns
+///
+/// A vectors of tensors
+///
+/// # Remarks
+///
+/// This is a fallback solution that used only when the backend doesn't have the corresponding implementation.
+/// Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved
+/// by static dispatch. It is not designed for direct usage by users, and not recommended to import
+/// or use this function directly.
+///


Small details, but I think spacing between the comments and the imports would be nice. Also, the last empty line of the comment isn't necessary.

… IntTensor

This reverts commit 7c6f017.

Kelvinyu1117 marked this pull request as draft December 2, 2023 06:02

Kelvinyu1117 marked this pull request as ready for review December 2, 2023 16:08

nathanielsimard requested changes Dec 4, 2023

View reviewed changes

louisfd requested changes Dec 15, 2023

View reviewed changes

louisfd approved these changes Dec 18, 2023

View reviewed changes

Kelvinyu1117 requested a review from nathanielsimard December 18, 2023 14:56

nathanielsimard requested changes Dec 19, 2023

View reviewed changes

Kelvinyu1117 requested a review from nathanielsimard December 20, 2023 16:14

Kelvinyu1117 added 5 commits December 21, 2023 00:17

Move narrow(), chunk() from base to BasicOps

2ae4a67

Implemented chunk for burn-candle, burn-ndarray, burn-tch

c81f23c

Implement narrow() on every backend, remove chunk from ndarray

cb36487

utilize backend implement of narrow() and chunk() for FloatTensor and…

a75b92c

… IntTensor

refactor chunk and narrow

e1e9eee

nathanielsimard approved these changes Dec 20, 2023

View reviewed changes

nathanielsimard merged commit 7c6f017 into tracel-ai:main Dec 20, 2023
13 of 14 checks passed

Kelvinyu1117 deleted the kelvin/feature/chunk-optimize branch December 21, 2023 09:43

syl20bnr added a commit to syl20bnr/burn that referenced this pull request Jan 4, 2024

Revert "Implement chunk for different backends (tracel-ai#1032)"

83278ad

This reverts commit 7c6f017.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement chunk for different backends #1032

Implement chunk for different backends #1032

Kelvinyu1117 commented Dec 1, 2023 •

edited

Loading

Kelvinyu1117 commented Dec 1, 2023 •

edited

Loading

Kelvinyu1117 commented Dec 2, 2023

nathanielsimard left a comment

nathanielsimard Dec 4, 2023

nathanielsimard Dec 4, 2023

nathanielsimard Dec 4, 2023

Kelvinyu1117 Dec 6, 2023

nathanielsimard Dec 12, 2023

Kelvinyu1117 Dec 18, 2023

louisfd left a comment

louisfd Dec 15, 2023

Kelvinyu1117 Dec 17, 2023

codecov bot commented Dec 18, 2023 •

edited

Loading

louisfd commented Dec 18, 2023

nathanielsimard left a comment

nathanielsimard Dec 19, 2023

Kelvinyu1117 Dec 20, 2023

nathanielsimard Dec 19, 2023

Kelvinyu1117 Dec 20, 2023

Implement chunk for different backends #1032

Implement chunk for different backends #1032

Conversation

Kelvinyu1117 commented Dec 1, 2023 • edited Loading

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

Kelvinyu1117 commented Dec 1, 2023 • edited Loading

Kelvinyu1117 commented Dec 2, 2023

nathanielsimard left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

louisfd left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Dec 18, 2023 • edited Loading

Codecov Report

louisfd commented Dec 18, 2023

nathanielsimard left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Kelvinyu1117 commented Dec 1, 2023 •

edited

Loading

Kelvinyu1117 commented Dec 1, 2023 •

edited

Loading

codecov bot commented Dec 18, 2023 •

edited

Loading