Convert reduce_dim_naive kernel to use the #[cube] derive macro

[cBournhonesque]

Pull Request Template

Following a discussion with @nathanielsimard on discord, I wanted to try to convert the reduce kernel to use the cube derive macro.
(originally because I wanted to know how to set the CubeCount for a kernel).

The new macro is pretty ergonomic, it's actually crazy that it also works for impl Trait.
I have a couple remaining issues though:

Checklist

• Implement CubeType for a tuple of CubeTypes: Implement CubeType for a tuple of CubeTypes cubecl#48
• How do I get the max value/min value of a Numeric? Do I need to add a trait for it?
• Need to add the computation of the CubeCount
• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Testing

Unit tests

[nathanielsimard]

Some changes to be done, but that's crazy how flexible cube is, I never though we would be able to migrate the reduce trait to cube with that amount of ease.

[cBournhonesque]

Some changes to be done, but that's crazy how flexible cube is, I never though we would be able to migrate the reduce trait to cube with that amount of ease.

Yes, it's crazy that the code just looks like normal rust but is still handled correctly! Very nice that it works across traits as well

[codecov]

Codecov Report

Attention: Patch coverage is 21.42857% with 77 lines in your changes missing coverage. Please review.

Project coverage is 86.08%. Comparing base (ff8d030) to head (0e5d280).

Files	Patch %	Lines
crates/burn-jit/src/kernel/reduce/naive/shader.rs	45.71%	19 Missing ⚠️
crates/burn-jit/src/kernel/reduce/naive/argmax.rs	5.26%	18 Missing ⚠️
crates/burn-jit/src/kernel/reduce/naive/argmin.rs	5.88%	16 Missing ⚠️
...rates/burn-jit/src/kernel/reduce/naive/prod_dim.rs	10.00%	9 Missing ⚠️
crates/burn-jit/src/kernel/reduce/naive/sum_dim.rs	10.00%	9 Missing ⚠️
...rates/burn-jit/src/kernel/reduce/naive/mean_dim.rs	14.28%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2117      +/-   ##
==========================================
- Coverage   86.21%   86.08%   -0.13%     
==========================================
  Files         694      694              
  Lines       88854    88747     -107     
==========================================
- Hits        76606    76399     -207     
- Misses      12248    12348     +100     

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