Tensor expand operator

[antimora]

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

fixes #1412

Changes

1. All backends
2. Fix issue with WGPU when stride == 0
3. Add new tests
4. Update docs

Testing

New unit tests

[codecov]

Codecov Report

Attention: Patch coverage is 89.31699% with 61 lines in your changes are missing coverage. Please review.

Project coverage is 85.91%. Comparing base (dc45cf1) to head (0c95cf1).

❗ Current head 0c95cf1 differs from pull request most recent head 5f32187. Consider uploading reports for the commit 5f32187 to get more accurate results

Files	Patch %	Lines
crates/burn-jit/src/ops/base.rs	80.48%	8 Missing ⚠️
crates/burn-tensor/src/tensor/api/base.rs	85.96%	8 Missing ⚠️
crates/burn-tch/src/ops/base.rs	0.00%	7 Missing ⚠️
crates/burn-autodiff/src/ops/bool_tensor.rs	0.00%	6 Missing ⚠️
crates/burn-autodiff/src/ops/int_tensor.rs	0.00%	6 Missing ⚠️
crates/burn-tch/src/ops/bool_tensor.rs	0.00%	6 Missing ⚠️
crates/burn-tch/src/ops/int_tensor.rs	0.00%	6 Missing ⚠️
crates/burn-tch/src/ops/tensor.rs	0.00%	6 Missing ⚠️
crates/burn-autodiff/src/ops/tensor.rs	93.05%	5 Missing ⚠️
crates/burn-fusion/src/stream/operation.rs	50.00%	2 Missing ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1508      +/-   ##
==========================================
+ Coverage   85.88%   85.91%   +0.02%     
==========================================
  Files         661      663       +2     
  Lines       74876    75446     +570     
==========================================
+ Hits        64307    64816     +509     
- Misses      10569    10630      +61     

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

[louisfd]

Small things related to autodiff, otherwise looks good!

[laggui]

Not much to add to the others' comments. I was also confused by the constant switch between expand <-> broadcast but it seems that it will be addressed.

To be fair, in pytorch tensor.expand(shape) is equivalent to torch.broadcast_to(tensor, shape).

I noticed the docstrings still used "broadcast" but it makes sense since they're only meant to explain the function.