Enhance gradient api #2965

VirtualNonsense · 2025-03-27T17:24:30Z

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

I adjusted the TensorContainer according the discussion in #2924 and adjusted the usage accordingly.
Question: Should the return type be propagated?

Testing

laggui

Question: Should the return type be propagated?

We have two options:

Propagate the result to the GradientsParams methods (get(...) and remove(...))
.unwrap() the result, which will panic with the error type (or .expect(...) with a descriptive message)

but I don't think we should not simply ignore the error with .ok(). This will be even more difficult to understand for user error 😅

codecov · 2025-03-27T20:34:38Z

Codecov Report

Attention: Patch coverage is 64.42953% with 53 lines in your changes missing coverage. Please review.

Project coverage is 81.06%. Comparing base (70e2bc7) to head (3ea5e58).
Report is 187 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-core/src/optim/grad_accum.rs	12.12%	29 Missing ⚠️
crates/burn-core/src/optim/visitor.rs	0.00%	12 Missing ⚠️
crates/burn-core/src/optim/simple/adaptor.rs	74.19%	8 Missing ⚠️
crates/burn-autodiff/src/backend.rs	76.92%	3 Missing ⚠️
crates/burn-tensor/src/tensor/container.rs	91.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2965      +/-   ##
==========================================
- Coverage   82.17%   81.06%   -1.11%     
==========================================
  Files         871      872       +1     
  Lines      120526   121221     +695     
==========================================
- Hits        99041    98268     -773     
- Misses      21485    22953    +1468

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

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

VirtualNonsense · 2025-03-29T14:16:43Z

Hey!
Please excuse my rather cautious approach. I'm a little bit burned by my corporate environment where even minor changes to lead to people scream at you 😅

If you don't mind I'll choose the first option after getting the check to run at my machine.

laggui · 2025-03-31T13:09:51Z

Ha, no need to apologize! 😄

Your feedback is important, especially since you were one of the first users to report this discrepancy. So if you feel that one option is more natural / helpful for an end-user, that's great!

…ptions and avoiding unwraps

…non disruptive way

VirtualNonsense · 2025-04-07T15:05:03Z

Please excuse the delay.

Here is my proposed solution for #2924.
I finally did not change the return type of the backend trait methods and chose to panic instead since it might break a lot of code and I don't feel like this change is important enough to justify such a change.
Let me know if I can enhance this PR in any way!

laggui

Overall, the addition of the result makes sense.

Just have some comments on the usage. See below 🙂

crates/burn-autodiff/src/backend.rs

crates/burn-core/src/optim/grads.rs

crates/burn-tensor/src/tests/mod.rs

VirtualNonsense · 2025-04-09T16:14:23Z

Thank you for taking the time for reviewing my PR and sorry for that sloppy oversight!
I'll try to account for those issues as soon as possible.

laggui · 2025-04-09T17:09:48Z

Ha no worries! That's what reviews are for 🙂

…ing to handle the result value in methods downstream in a way that is consequent and helpful.

…ncast error

…rContainerError gets another variant for whatever reason and to increase the ease of debugging.

laggui · 2025-04-14T12:29:12Z

@VirtualNonsense also let me know when this is ready for another round of review 🙂

VirtualNonsense · 2025-04-14T12:54:43Z

@laggui
I tried to address all issues/comments and the pipeline seems to be happy as well so it should be fine for a review.
Thanks in advance!

nathanielsimard · 2025-04-14T13:29:07Z

crates/burn-tensor/src/tensor/container.rs

+    /// Gradients are not stored on the autodiff backend.
+    DowncastError,
+}
+
 /// Contains tensor of arbitrary dimension.
 #[derive(Debug)]
 pub struct TensorContainer<ID> {


The main reason to have a Box<dyn Any> here was because the TensorPrimitive used to be generic over the rank. Now that's not the case anymore, we could actually remove the Box<dyn Any> and see if we can use the right backend as generic.

Thank you for taking the time to review my PR!

If I unterstand you correctly you are suggesting to change the TensorContainer to something like this:

/// Contains tensor of arbitrary dimension. #[derive(Debug)] pub struct TensorContainer<ID, B: Backend> { tensors: HashMap<ID, TensorPrimitive<B>>, }

and thereby sidestepping the issue entirely?
I have to admit I'm not that good with rust yet and do not know burn well enough to properly foresee the consequences of that change but I can try if you like :D
Should I create another PR for that version?

github-actions · 2025-05-15T12:11:50Z

This PR has been marked as stale because it has not been updated for over a month

github-actions · 2025-06-16T12:12:48Z

This PR has been marked as stale because it has not been updated for over a month

laggui reviewed Mar 27, 2025

View reviewed changes

Andreas Nachtmann added 7 commits March 31, 2025 21:31

Improved the api for tensor containers by return results instead of o…

887eba5

…ptions and avoiding unwraps

error type should be debug

a783a10

fixed build issues that were caused by the container api change in a …

b871f74

…non disruptive way

adding PartialEq for ease of testing.

ee2f2a6

added tests for container

b738f7c

propagating errors

e0842e5

Downcast error should panic in backend with a clear message.

7741896

VirtualNonsense force-pushed the enhance_gradient_api branch from d9f4ee2 to 7741896 Compare March 31, 2025 19:32

formatting

4cb2c79

VirtualNonsense marked this pull request as ready for review April 7, 2025 14:59

laggui requested changes Apr 7, 2025

View reviewed changes

crates/burn-autodiff/src/backend.rs Outdated Show resolved Hide resolved

crates/burn-core/src/optim/grads.rs Outdated Show resolved Hide resolved

crates/burn-tensor/src/tests/mod.rs Outdated Show resolved Hide resolved

Andreas Nachtmann and others added 7 commits April 10, 2025 13:09

Changed return type of GradientsParams::remove and get to Result; Try…

e0dc6c9

…ing to handle the result value in methods downstream in a way that is consequent and helpful.

fixed typo

5e0a164

Summarized duplicates into macro

d0675a8

Moved existing tests into grads of burn-core; added test case for dow…

f931314

…ncast error

using alloc instead of std implementation of format! and String

22a884f

Explicit match arms to catch errors down stream quicker in case Tenso…

269d5b0

…rContainerError gets another variant for whatever reason and to increase the ease of debugging.

Formatting.

3ea5e58

nathanielsimard reviewed Apr 14, 2025

View reviewed changes

github-actions bot added the stale The issue or pr has been open for too long label May 15, 2025

github-actions bot removed the stale The issue or pr has been open for too long label May 16, 2025

github-actions bot added stale The issue or pr has been open for too long and removed stale The issue or pr has been open for too long labels Jun 16, 2025

VirtualNonsense closed this Jun 18, 2025

Enhance gradient api #2965

Enhance gradient api #2965

Uh oh!

Conversation

VirtualNonsense commented Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

Uh oh!

laggui left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

VirtualNonsense commented Mar 29, 2025

Uh oh!

laggui commented Mar 31, 2025

Uh oh!

VirtualNonsense commented Apr 7, 2025

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

VirtualNonsense commented Apr 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

laggui commented Apr 9, 2025

Uh oh!

laggui commented Apr 14, 2025

Uh oh!

VirtualNonsense commented Apr 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nathanielsimard Apr 14, 2025

Choose a reason for hiding this comment

Uh oh!

VirtualNonsense Apr 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented May 15, 2025

Uh oh!

github-actions bot commented Jun 16, 2025

Uh oh!

Uh oh!

VirtualNonsense commented Mar 27, 2025 •

edited

Loading

laggui left a comment •

edited

Loading

codecov bot commented Mar 27, 2025 •

edited

Loading

VirtualNonsense commented Apr 9, 2025 •

edited

Loading

VirtualNonsense commented Apr 14, 2025 •

edited

Loading

VirtualNonsense Apr 15, 2025 •

edited

Loading