Chore/update cubecl (#2410)

Cargo.toml

@@ -152,8 +152,8 @@ ahash = { version = "0.8.11", default-features = false }
 portable-atomic-util = { version = "0.2.2", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "44be0489c04c1c88bef61bde4c386e6772693b94" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "44be0489c04c1c88bef61bde4c386e6772693b94" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

crates/burn-import/onnx-tests/tests/test_onnx.rs

@@ -1526,7 +1526,7 @@ mod tests {
         // because the output tensor is too large to compare with the expected tensor.
         let output_sum = output.sum().into_scalar();
 
-        let expected_sum = 33.810329; // example result running the corresponding PyTorch model (conv_transpose1d.py)
+        let expected_sum = 33.810_33; // example result running the corresponding PyTorch model (conv_transpose1d.py)
 
         assert!(expected_sum.approx_eq(output_sum, (1.0e-4, 2)));
     }

crates/burn-jit/src/kernel/index/slice.rs

@@ -110,21 +110,23 @@ pub(crate) fn slice<R: JitRuntime, E: JitElement>(
     indices: &[Range<usize>],
 ) -> JitTensor<R, E> {
     let mut dims = tensor.shape.dims.clone();
-    let mut offset_start = 0;
-    let mut offset_end = 0;
+    let mut offset_start = 0u64;
+    let mut offset_end = 0u64;
 
     for i in 0..indices.len() {
-        offset_start += tensor.strides[i] * indices[i].start;
-        offset_end += tensor.strides[i] * (dims[i] - indices[i].end);
+        offset_start += (tensor.strides[i] * indices[i].start) as u64;
+        offset_end += (tensor.strides[i] * (dims[i] - indices[i].end)) as u64;
         dims[i] = indices[i].end - indices[i].start;
     }
 
-    let offset_start = offset_start * E::cube_elem().size();
-    let offset_end = offset_end * E::cube_elem().size();
+    let offset_start = offset_start * E::cube_elem().size() as u64;
+    let offset_end = offset_end * E::cube_elem().size() as u64;
 
     let memory_offset_alignment = tensor.client.properties().memory_properties().alignment;
 
-    if offset_start % memory_offset_alignment == 0 && offset_end % memory_offset_alignment == 0 {
+    if offset_start % memory_offset_alignment == 0u64
+        && offset_end % memory_offset_alignment == 0u64
+    {
         JitTensor::new(
             tensor.client,
             tensor

crates/burn-jit/src/kernel/matmul/tune/base.rs

@@ -92,7 +92,7 @@ pub fn matmul_autotune<R: JitRuntime, E: FloatElement + Element>(
 
 macro_rules! matmul_tune_ops {
     ($name:ident, $func:expr) => {
-        #[derive(new)]
+        #[derive(new, Debug)]
         pub(crate) struct $name<R: JitRuntime, E: FloatElement> {
             lhs: JitTensor<R, E>,
             rhs: JitTensor<R, E>,

crates/burn-jit/src/kernel/reduce/base.rs

@@ -9,7 +9,7 @@ use super::{
 
 #[allow(dead_code)]
 pub(crate) trait ReduceDimAlgorithm<EI: JitElement>:
-    ReduceDimNaive<EI> + ReduceDimShared<EI>
+    core::fmt::Debug + ReduceDimNaive<EI> + ReduceDimShared<EI>
 {
 }
 
@@ -64,7 +64,9 @@ impl Default for ReduceStrategy {
 
 macro_rules! reduce_operation {
     ($name:ident, $ops:ident) => {
+        #[derive(Debug)]
         pub(crate) struct $ops;
+
         impl<EI: JitElement> ReduceDimAlgorithm<EI> for $ops {}
 
         /// Executes the reduce operation with the given strategy.

crates/burn-jit/src/kernel/reduce/tune/base.rs

@@ -133,7 +133,7 @@ pub(crate) fn reduce_dim_autotune<
     output
 }
 
-#[derive(new)]
+#[derive(new, Debug)]
 // Probably better on balanced tensor shapes
 pub(crate) struct ReduceDimNaiveAutotune<
     RD: ReduceDimAlgorithm<EI>,
@@ -169,7 +169,7 @@ where
     }
 }
 
-#[derive(new)]
+#[derive(new, Debug)]
 // Probably better on tensors large along reduce dim
 pub(crate) struct ReduceDimSharedAutotune<
     RD: ReduceDimAlgorithm<EI>,