Add support for Any, All operations to Tensor

[ashdtu]

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

#1341

Changes

• Add support for any() (logical_or), all() (logical_and) operations on Tensor struct for all types (Float, Int, Bool).
• For numeric tensor type(Float, Int), the input tensors are casted to Bool by checking against non-zero values.
• The implementation supports working with these operations along a dim using the any_dim(), all_dim() methods.
• Burn book updated with these Tensor ops.

Testing

Test cases written for the implemented methods (any(), any_dim(), all(), all_dim()) with all tensor types(float, int, bool) in burn-tensor/src/tests/ops/any.rs or all.rs.

[ashdtu]

@louisfd The bug in WGPU about the boolean casting to float/int is no longer a blocker in the updated code. I noticed that I can get rid of that by just changing the logic a bit and adding another bool_not op which converts the weird large values of 1065353216 for (1.0 f / 1 ) to exact zeros. There's no overflow now since this op happens before the final sum, so it's all ok.

TL;DR There's no blocker for this PR and it works for all backends.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.93%. Comparing base (4427768) to head (438a7c0).
Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1342      +/-   ##
==========================================
+ Coverage   78.77%   78.93%   +0.15%     
==========================================
  Files         551      563      +12     
  Lines       61836    62981    +1145     
==========================================
+ Hits        48712    49714    +1002     
- Misses      13124    13267     +143     

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

[louisfd]

Awesome! Thank you.
The logic is good everywhere. I do have some comments though, most of them minor but one is important, you must not do tensor.clone() if you're not obligated, this will augment reference count and may prevent operations to happen in-place. To make sure no occurrence of the clone is forgotten i commented them all 😄

[louisfd]

@ashdtu While we're here, can you change this LSTM test line in lstm.rs (which was the original reason I needed the any op):

 // Asserts the gradients exist and are non zero
 assert!(*some_gradient.abs().sum().into_data().value.first().unwrap() > 0.);

[louisfd]

LGTM
Thanks