Add 0-dim tensor checks for creation ops and validate TensorData shape w/ num values

[laggui]

Checklist

• [x] Confirmed that run-checks all script has been executed.

Related Issues/PRs

Fixes #2002

Changes

To prevent users from creating 0-dim tensors (invalid), I've added checks for all creation ops, and while I was at it some checks when creating TensorData struct to make sure the number of values matches the provided shape.

Users trying to create a 0-dim tensor will now be prompted with the following error:

=== Tensor Operation Error ===
  Operation: 'Zeros'
  Reason:
    1. Tried to create a 0-dim tensor, which is invalid. Tensor rank: '0', given dimensions: '[]'.

note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

[antimora]

CI is failing

[laggui]

Yep looks like the checks are also failing locally for a single broadcast test with wgpu. I opened the PR before they completed, will check it out.

/edit:

Uh, looks like the behavior for expand() + into_data() is different on wgpu compared to other backends.

let t = Tensor::<B, 1>::from_floats([1.0], &device);
let out = t.expand([2, 3]);

println!("{}", out);
println!("{}", out.into_data());

For other backends (e.g., ndarray or tch), we have

Tensor {
  data:
[[1.0, 1.0, 1.0],
 [1.0, 1.0, 1.0]],
  shape:  [2, 3],
  device:  Cpu,
  backend:  "tch",
  kind:  "Float",
  dtype:  "f32",
}
[1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

But for wgpu, the TensorData contains a single value.

Tensor {
  data:
[[1.0, 1.0, 1.0],
 [1.0, 1.0, 1.0]],
  shape:  [2, 3],
  device:  BestAvailable,
  backend:  "fusion<jit<wgpu>>",
  kind:  "Float",
  dtype:  "f32",
}
[1.0]

This is why the test is failing

---- tests::jit::expand::tests::broadcast_single stdout ----
thread 'tests::jit::expand::tests::broadcast_single' panicked at /home/laggui/workspace/burn/crates/burn-tensor/src/tensor/data.rs:70:9:
assertion `left == right` failed: Shape [2, 3] is invalid for input of size 1
  left: 6
 right: 1

---- tests::jit_fusion::expand::tests::broadcast_single stdout ----
thread 'tests::jit_fusion::expand::tests::broadcast_single' panicked at /home/laggui/workspace/burn/crates/burn-tensor/src/tensor/data.rs:70:9:
assertion `left == right` failed: Shape [2, 3] is invalid for input of size 1
  left: 6
 right: 1


failures:
    tests::jit::expand::tests::broadcast_single
    tests::jit_fusion::expand::tests::broadcast_single

test result: FAILED. 1550 passed; 2 failed; 12 ignored; 0 measured; 0 filtered out; finished in 21.48s

[laggui]

Once the linked cubecl PR is merged I'll update the dependency and this PR should have tests passing.

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 78.12500% with 7 lines in your changes missing coverage. Please review.

Project coverage is 86.11%. Comparing base (16239db) to head (c6362fc). Report is 1 commits behind head on main.

Files	Patch %	Lines
crates/burn-tensor/src/tensor/api/check.rs	33.33%	6 Missing :warning:
crates/burn-tensor/src/tensor/data.rs	85.71%	1 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2137      +/-   ##
==========================================
- Coverage   86.11%   86.11%   -0.01%     
==========================================
  Files         694      694              
  Lines       88822    88845      +23     
==========================================
+ Hits        76490    76507      +17     
- Misses      12332    12338       +6     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.