feat: update workspace paths and enhance gitignore

- Updated stablediffusion crate path from "../stable-diffusion-burn" to "./crates/stable-diffusion-burn" for proper workspace resolution
- Enhanced .gitignore to include generated model files (.mpk, .pt, .bin, .safetensors, .ckpt) and user_data directory
- Added Cargo.lock to gitignore with appropriate comment
- Reorganized IDE files section in gitignore for better clarity
- Added newline at end of file for proper formatting

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/avg_pool2d.rs

@@ -0,0 +1,96 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{
+    Distribution, Tensor, TensorPrimitive, backend::Backend, module, ops::ModuleOps,
+};
+
+#[test]
+fn avg_pool2d_should_match_reference_backend() {
+    let tensor = Tensor::<TestBackend, 4>::random(
+        [32, 32, 32, 32],
+        Distribution::Default,
+        &Default::default(),
+    );
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &Default::default());
+    let kernel_size = [3, 4];
+    let stride = [1, 2];
+    let padding = [1, 2];
+    let count_include_pad = true;
+
+    let pooled = module::avg_pool2d(
+        tensor,
+        kernel_size,
+        stride,
+        padding,
+        count_include_pad,
+        false,
+    );
+    let pooled_ref = module::avg_pool2d(
+        tensor_ref,
+        kernel_size,
+        stride,
+        padding,
+        count_include_pad,
+        false,
+    );
+
+    pooled
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&pooled_ref.into_data(), Tolerance::default());
+}
+
+#[test]
+fn avg_pool2d_backward_should_match_reference_backend() {
+    let device = Default::default();
+
+    TestBackend::seed(&device, 0);
+    ReferenceBackend::seed(&Default::default(), 0);
+
+    let tensor = Tensor::<TestBackend, 4>::random([32, 32, 32, 32], Distribution::Default, &device);
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &Default::default());
+    let kernel_size = [3, 3];
+    let stride = [1, 1];
+    let padding = [1, 1];
+    let count_include_pad = true;
+
+    let shape_out = module::avg_pool2d(
+        tensor.clone(),
+        kernel_size,
+        stride,
+        padding,
+        count_include_pad,
+        false,
+    )
+    .shape();
+    let grad_output =
+        Tensor::<TestBackend, 4>::random(shape_out, Distribution::Default, &Default::default());
+    let grad_output_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(grad_output.to_data(), &Default::default());
+
+    let grad: Tensor<TestBackend, 4> =
+        Tensor::from_primitive(TensorPrimitive::Float(TestBackend::avg_pool2d_backward(
+            tensor.into_primitive().tensor(),
+            grad_output.into_primitive().tensor(),
+            kernel_size,
+            stride,
+            padding,
+            count_include_pad,
+            false,
+        )));
+    let grad_ref: Tensor<ReferenceBackend, 4> = Tensor::from_primitive(TensorPrimitive::Float(
+        ReferenceBackend::avg_pool2d_backward(
+            tensor_ref.into_primitive().tensor(),
+            grad_output_ref.into_primitive().tensor(),
+            kernel_size,
+            stride,
+            padding,
+            count_include_pad,
+            false,
+        ),
+    ));
+
+    grad.into_data()
+        .assert_approx_eq::<FloatElem>(&grad_ref.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/bernoulli.rs

@@ -0,0 +1,48 @@
+use super::*;
+
+use serial_test::serial;
+
+use core::f32;
+
+use burn_tensor::{Distribution, Shape, Tensor, backend::Backend};
+
+use cubek::random::{assert_number_of_1_proportional_to_prob, assert_wald_wolfowitz_runs_test};
+
+#[test]
+#[serial]
+fn number_of_1_proportional_to_prob() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let shape: Shape = [40, 40].into();
+    let prob = 0.7;
+
+    let tensor =
+        Tensor::<TestBackend, 2>::random(shape.clone(), Distribution::Bernoulli(prob), &device)
+            .into_data();
+
+    let numbers = tensor
+        .as_slice::<<TestBackend as Backend>::FloatElem>()
+        .unwrap();
+
+    assert_number_of_1_proportional_to_prob(numbers, prob as f32);
+}
+
+#[test]
+#[serial]
+fn wald_wolfowitz_runs_test() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let shape = Shape::new([512, 512]);
+    let device = Default::default();
+    let tensor = Tensor::<TestBackend, 2>::random(shape, Distribution::Bernoulli(0.5), &device);
+
+    let data = tensor.into_data();
+    let numbers = data
+        .as_slice::<<TestBackend as Backend>::FloatElem>()
+        .unwrap();
+
+    // High bound slightly over 1 so 1.0 is included in second bin
+    assert_wald_wolfowitz_runs_test(numbers, 0., 1.1);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/cast.rs

@@ -0,0 +1,45 @@
+use super::*;
+use burn_tensor::{Int, Tensor, TensorData};
+
+#[test]
+fn should_cast_int_to_float() {
+    const START: usize = 0;
+    const END: usize = 100;
+
+    let device = Default::default();
+    let tensor = Tensor::<TestBackend, 1, Int>::arange(START as i64..END as i64, &device);
+
+    let data_int = tensor.to_data();
+    let data_int = data_int.as_slice::<i32>().unwrap();
+    let data_float = tensor.float().into_data();
+    let data_float = data_float.as_slice::<f32>().unwrap();
+
+    for i in START..END {
+        assert_eq!(data_int[i], i as i32);
+        assert_eq!(data_float[i], i as f32);
+    }
+}
+
+#[test]
+fn should_cast_bool_to_int() {
+    let device = Default::default();
+
+    let tensor_1 = Tensor::<TestBackend, 2>::from_floats([[1., 0., 3.], [0., 0., 900.]], &device);
+    let tensor_2: Tensor<TestBackend, 2, Int> = tensor_1.clone().greater_elem(0.0).int();
+
+    tensor_2
+        .to_data()
+        .assert_eq(&TensorData::from([[1, 0, 1], [0, 0, 1]]), false);
+}
+
+#[test]
+fn should_cast_bool_to_float() {
+    let device = Default::default();
+
+    let tensor_1 = Tensor::<TestBackend, 2>::from_floats([[1., 0., 3.], [0., 0., 900.]], &device);
+    let tensor_2: Tensor<TestBackend, 2> = tensor_1.clone().greater_elem(0.0).float();
+
+    tensor_2
+        .to_data()
+        .assert_eq(&TensorData::from([[1., 0., 1.], [0., 0., 1.]]), false);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/cat.rs

@@ -0,0 +1,42 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, backend::Backend};
+
+#[test]
+fn cat_should_match_reference_backend_dim0() {
+    test_same_as_reference([6, 256], 2, 0);
+}
+
+#[test]
+fn cat_should_match_reference_backend_dim1() {
+    test_same_as_reference([6, 256], 2, 1);
+}
+
+#[test]
+fn cat_should_support_uneven_launch() {
+    test_same_as_reference([1, 137], 2, 0);
+}
+
+fn test_same_as_reference(shape: [usize; 2], num_tensors: usize, dim: usize) {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let tensors = (0..num_tensors)
+        .map(|_| {
+            Tensor::<TestBackend, 2>::random(shape, Distribution::Default, &Default::default())
+        })
+        .collect::<Vec<_>>();
+    let tensors_ref = tensors
+        .iter()
+        .map(|tensor| {
+            Tensor::<ReferenceBackend, 2>::from_data(tensor.to_data(), &Default::default())
+        })
+        .collect::<Vec<_>>();
+
+    let tensor = Tensor::<TestBackend, 2>::cat(tensors, dim);
+    let tensor_ref = Tensor::<ReferenceBackend, 2>::cat(tensors_ref, dim);
+
+    tensor
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&tensor_ref.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/clamp.rs

@@ -0,0 +1,20 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor};
+
+#[test]
+fn clamp_should_match_reference() {
+    let input = Tensor::<TestBackend, 4>::random(
+        [1, 5, 32, 32],
+        Distribution::Default,
+        &Default::default(),
+    );
+    let input_ref = Tensor::<ReferenceBackend, 4>::from_data(input.to_data(), &Default::default());
+
+    let output = input.clamp(0.3, 0.7);
+
+    output.into_data().assert_approx_eq::<FloatElem>(
+        &input_ref.clamp(0.3, 0.7).into_data(),
+        Tolerance::default(),
+    );
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/contiguous.rs

@@ -0,0 +1,40 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Int, Tensor};
+
+#[test]
+pub fn into_contiguous_match_reference_backend_1() {
+    for shape in [
+        [4, 4, 4, 4],
+        [32, 42, 24, 48],
+        [8, 3, 7, 4],
+        [1, 4, 1, 1],
+        [1, 32, 256, 128],
+    ] {
+        let num_elems = shape.iter().product::<usize>() as i64;
+        let tensor: Tensor<TestBackend, 4> =
+            Tensor::<TestBackend, 1, Int>::arange(0..num_elems, &Default::default())
+                .reshape(shape)
+                .float();
+        let tensor_ref =
+            Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &Default::default());
+
+        for (i, j) in get_combinations(shape.len()) {
+            let view = tensor.clone().swap_dims(i, j);
+            let view_ref = tensor_ref.clone().swap_dims(i, j);
+            let data = view.into_data();
+            let data_ref = view_ref.into_data();
+
+            data_ref.assert_approx_eq::<FloatElem>(&data, Tolerance::default());
+        }
+    }
+}
+
+fn get_combinations(n: usize) -> impl Iterator<Item = (usize, usize)> {
+    // Iterate from 0 up to n
+    (0..n).flat_map(move |i| {
+        // For each i, iterate from i + 1 up to n
+        // This ensures no repeats (i == j) and no duplicates (j, i)
+        (i + 1..n).map(move |j| (i, j))
+    })
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/conv2d.rs

@@ -0,0 +1,78 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, module};
+
+#[test]
+fn conv2d_should_match_reference_backend() {
+    let test_device = Default::default();
+    let input =
+        Tensor::<TestBackend, 4>::random([6, 16, 32, 32], Distribution::Default, &test_device);
+    let weight =
+        Tensor::<TestBackend, 4>::random([12, 8, 3, 3], Distribution::Default, &test_device);
+    let bias = Tensor::<TestBackend, 1>::random([12], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+
+    let input_ref = Tensor::<ReferenceBackend, 4>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 4>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let options = burn_tensor::ops::ConvOptions::new([2, 3], [2, 3], [2, 3], 2);
+
+    let output = module::conv2d(input, weight, Some(bias), options.clone());
+    let output_ref = module::conv2d(input_ref, weight_ref, Some(bias_ref), options);
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), Tolerance::default());
+}
+
+#[test]
+fn conv2d_should_match_reference_backend_implicit() {
+    let test_device = Default::default();
+    let input =
+        Tensor::<TestBackend, 4>::random([4, 16, 6, 6], Distribution::Default, &test_device);
+    let weight =
+        Tensor::<TestBackend, 4>::random([16, 16, 3, 3], Distribution::Default, &test_device);
+    let bias = Tensor::<TestBackend, 1>::random([16], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+
+    let input_ref = Tensor::<ReferenceBackend, 4>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 4>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let options = burn_tensor::ops::ConvOptions::new([1, 1], [2, 2], [1, 1], 1);
+
+    let output = module::conv2d(input, weight, Some(bias), options.clone());
+    let output_ref = module::conv2d(input_ref, weight_ref, Some(bias_ref), options);
+
+    let tolerance = Tolerance::default();
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), tolerance);
+}
+
+/// Regression test for bias loader in new implicit GEMM
+#[test]
+fn conv2d_should_match_reference_backend_bias_regression() {
+    let test_device = Default::default();
+    let input = Tensor::<TestBackend, 4>::random([1, 1, 1, 1], Distribution::Default, &test_device);
+    let weight =
+        Tensor::<TestBackend, 4>::random([32, 1, 3, 3], Distribution::Default, &test_device);
+    let bias = Tensor::<TestBackend, 1>::random([32], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+
+    let input_ref = Tensor::<ReferenceBackend, 4>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 4>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let options = burn_tensor::ops::ConvOptions::new([1, 1], [1, 1], [1, 1], 1);
+
+    let output = module::conv2d(input, weight, Some(bias), options.clone()).permute([0, 2, 3, 1]);
+    let output_ref =
+        module::conv2d(input_ref, weight_ref, Some(bias_ref), options).permute([0, 2, 3, 1]);
+
+    let tolerance = Tolerance::default();
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), tolerance);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/conv3d.rs

@@ -0,0 +1,27 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, module};
+
+#[test]
+fn conv3d_should_match_reference_backend() {
+    let test_device = Default::default();
+    let input =
+        Tensor::<TestBackend, 5>::random([6, 16, 32, 32, 32], Distribution::Default, &test_device);
+    let weight =
+        Tensor::<TestBackend, 5>::random([12, 8, 3, 3, 3], Distribution::Default, &test_device);
+    let bias = Tensor::<TestBackend, 1>::random([12], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+
+    let input_ref = Tensor::<ReferenceBackend, 5>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 5>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let options = burn_tensor::ops::ConvOptions::new([2, 3, 4], [2, 3, 4], [2, 3, 4], 2);
+
+    let output = module::conv3d(input, weight, Some(bias), options.clone());
+    let output_ref = module::conv3d(input_ref, weight_ref, Some(bias_ref), options);
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/conv_transpose2d.rs

@@ -0,0 +1,48 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, backend::Backend, module};
+
+#[test]
+fn conv_transpose2d_should_match_reference_backend() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let height = 8;
+    let width = 8;
+    let in_channels = 8;
+    let out_channels = 8;
+    let batch_size = 32;
+    let kernel_size_0 = 3;
+    let kernel_size_1 = 3;
+    let options = burn_tensor::ops::ConvTransposeOptions::new([1, 1], [1, 1], [0, 0], [1, 1], 1);
+
+    let test_device = Default::default();
+    let input = Tensor::<TestBackend, 4>::random(
+        [batch_size, in_channels, height, width],
+        Distribution::Default,
+        &test_device,
+    );
+    let weight = Tensor::<TestBackend, 4>::random(
+        [
+            in_channels,
+            out_channels / options.groups,
+            kernel_size_0,
+            kernel_size_1,
+        ],
+        Distribution::Default,
+        &test_device,
+    );
+    let bias =
+        Tensor::<TestBackend, 1>::random([out_channels], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+    let input_ref = Tensor::<ReferenceBackend, 4>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 4>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let output = module::conv_transpose2d(input, weight, Some(bias), options.clone());
+    let output_ref = module::conv_transpose2d(input_ref, weight_ref, Some(bias_ref), options);
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), Tolerance::rel_abs(0.01, 0.02));
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/conv_transpose3d.rs

@@ -0,0 +1,51 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, backend::Backend, module};
+
+#[test]
+fn conv_transpose3d_should_match_reference_backend() {
+    let test_device = Default::default();
+    TestBackend::seed(&test_device, 0);
+
+    let depth = 8;
+    let height = 8;
+    let width = 8;
+    let in_channels = 8;
+    let out_channels = 8;
+    let batch_size = 32;
+    let kernel_size_0 = 3;
+    let kernel_size_1 = 3;
+    let kernel_size_2 = 3;
+    let options =
+        burn_tensor::ops::ConvTransposeOptions::new([1, 1, 1], [1, 1, 1], [0, 0, 0], [1, 1, 1], 1);
+
+    let input = Tensor::<TestBackend, 5>::random(
+        [batch_size, in_channels, depth, height, width],
+        Distribution::Default,
+        &test_device,
+    );
+    let weight = Tensor::<TestBackend, 5>::random(
+        [
+            in_channels,
+            out_channels / options.groups,
+            kernel_size_0,
+            kernel_size_1,
+            kernel_size_2,
+        ],
+        Distribution::Default,
+        &test_device,
+    );
+    let bias =
+        Tensor::<TestBackend, 1>::random([out_channels], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+    let input_ref = Tensor::<ReferenceBackend, 5>::from_data(input.to_data(), &ref_device);
+    let weight_ref = Tensor::<ReferenceBackend, 5>::from_data(weight.to_data(), &ref_device);
+    let bias_ref = Tensor::<ReferenceBackend, 1>::from_data(bias.to_data(), &ref_device);
+
+    let output = module::conv_transpose3d(input, weight, Some(bias), options.clone());
+    let output_ref = module::conv_transpose3d(input_ref, weight_ref, Some(bias_ref), options);
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/cross.rs

@@ -0,0 +1,159 @@
+use super::*;
+use burn_tensor::Tensor;
+use burn_tensor::Tolerance;
+
+#[test]
+fn test_cross_product() {
+    let device = Default::default();
+    // Test with well-known orthogonal vectors for clearer validation
+    let a = Tensor::<TestBackend, 2>::from_data([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], &device);
+    let b = Tensor::<TestBackend, 2>::from_data([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], &device);
+
+    let result = a.cross(b, 1);
+    // For orthogonal unit vectors:
+    // i × j = k
+    // j × k = i
+    let expected = Tensor::<TestBackend, 2>::from_data([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], &device);
+
+    // Use Tolerance for floating-point comparisons
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}
+
+#[test]
+fn test_cross_product_zeros() {
+    let device = Default::default();
+    // Test cross product with zero vector - should always give zero vector
+    let a = Tensor::<TestBackend, 2>::from_data([[2.0, 3.0, 4.0]], &device);
+    let b = Tensor::<TestBackend, 2>::zeros([1, 3], &device);
+
+    let result = a.cross(b, 1);
+    let expected = Tensor::<TestBackend, 2>::zeros([1, 3], &device);
+
+    // For zeros, we can use exact equality or a very tight tolerance
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}
+
+#[test]
+fn test_cross_product_batch() {
+    let device = Default::default();
+    // Test typical cross product computations in batch
+    let a = Tensor::<TestBackend, 2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
+    let b = Tensor::<TestBackend, 2>::from_data([[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device);
+
+    let result = a.cross(b, 1);
+    // Cross products:
+    // [1,2,3] × [4,5,6] = [-3,6,-3]
+    // [4,5,6] × [7,8,9] = [-3,6,-3]
+    let expected =
+        Tensor::<TestBackend, 2>::from_data([[-3.0, 6.0, -3.0], [-3.0, 6.0, -3.0]], &device);
+
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}
+
+#[test]
+#[should_panic]
+fn test_cross_product_invalid_dimension() {
+    let device = Default::default();
+    let a = Tensor::<TestBackend, 2>::zeros([1, 4], &device);
+    let b = Tensor::<TestBackend, 2>::zeros([1, 4], &device);
+
+    let _ = a.cross(b, 1);
+}
+
+#[test]
+fn test_cross_product_parallel_vectors() {
+    let device = Default::default();
+    // Test cross product of parallel vectors (should be zero)
+    let a = Tensor::<TestBackend, 2>::from_data([[1.0, 2.0, 3.0]], &device);
+    let b = Tensor::<TestBackend, 2>::from_data([[2.0, 4.0, 6.0]], &device); // b = 2 * a
+
+    let result = a.cross(b, 1);
+    let expected = Tensor::<TestBackend, 2>::zeros([1, 3], &device);
+
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}
+
+#[test]
+fn test_cross_product_3d_tensor() {
+    let device = Default::default();
+    // Test with 3D tensor (batch of matrices)
+    let a = Tensor::<TestBackend, 3>::from_data(
+        [
+            [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
+            [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+        ],
+        &device,
+    );
+
+    let b = Tensor::<TestBackend, 3>::from_data(
+        [
+            [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+            [[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
+        ],
+        &device,
+    );
+
+    let result = a.cross(b, 2); // Cross on last dimension
+    let expected = Tensor::<TestBackend, 3>::from_data(
+        [
+            [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]],
+            [[-3.0, 6.0, -3.0], [-3.0, 6.0, -3.0]],
+        ],
+        &device,
+    );
+
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}
+
+// Test to verify that padding doesn't affect results
+#[test]
+fn test_cross_product_with_padding_awareness() {
+    let device = Default::default();
+    // Create tensors that would span multiple 4-element blocks
+    // This tests that the padding doesn't corrupt adjacent data
+    let a = Tensor::<TestBackend, 2>::from_data(
+        [
+            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // Two vectors: [1,2,3] and [4,5,6]
+        ],
+        &device,
+    );
+
+    let b = Tensor::<TestBackend, 2>::from_data(
+        [
+            [7.0, 8.0, 9.0, 10.0, 11.0, 12.0], // Two vectors: [7,8,9] and [10,11,12]
+        ],
+        &device,
+    );
+
+    // Reshape to have proper 3-element vectors in last dimension
+    let a_reshaped = a.reshape([2, 3]);
+    let b_reshaped = b.reshape([2, 3]);
+
+    let result = a_reshaped.cross(b_reshaped, 1);
+
+    // Expected cross products:
+    // [1,2,3] × [7,8,9] = [-6,12,-6]
+    // [4,5,6] × [10,11,12] = [-6,12,-6]
+    let expected =
+        Tensor::<TestBackend, 2>::from_data([[-6.0, 12.0, -6.0], [-6.0, 12.0, -6.0]], &device);
+
+    let tolerance = Tolerance::<FloatElem>::default();
+    result
+        .to_data()
+        .assert_approx_eq(&expected.to_data(), tolerance);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/gather.rs

@@ -0,0 +1,44 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Int, Shape, Tensor, backend::Backend};
+
+#[test]
+fn gather_should_work_with_multiple_workgroups_dim0() {
+    test_same_as_ref([6, 256], 0);
+}
+
+#[test]
+fn gather_should_work_with_multiple_workgroups_dim1() {
+    test_same_as_ref([6, 256], 1);
+}
+
+fn test_same_as_ref<const D: usize>(shape: [usize; D], dim: usize) {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let max = shape[dim];
+    let shape = Shape::new(shape);
+    let tensor =
+        Tensor::<TestBackend, D>::random(shape.clone(), Distribution::Default, &Default::default());
+    let indices = Tensor::<TestBackend, 1, Int>::from_data(
+        Tensor::<TestBackend, 1>::random(
+            [shape.num_elements()],
+            Distribution::Uniform(0., max as f64),
+            &Default::default(),
+        )
+        .into_data(),
+        &Default::default(),
+    )
+    .reshape(shape);
+    let tensor_ref =
+        Tensor::<ReferenceBackend, D>::from_data(tensor.to_data(), &Default::default());
+    let indices_ref =
+        Tensor::<ReferenceBackend, D, Int>::from_data(indices.to_data(), &Default::default());
+
+    let actual = tensor.gather(dim, indices);
+    let expected = tensor_ref.gather(dim, indices_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/mask_fill.rs

@@ -0,0 +1,64 @@
+use super::*;
+use burn_cubecl::kernel::{MaskFillStrategy, mask_fill};
+use burn_tensor::Tolerance;
+use burn_tensor::{Bool, Distribution, Element, Tensor, TensorPrimitive, backend::Backend};
+use cubecl::prelude::InputScalar;
+
+#[test]
+fn mask_fill_should_match_reference_backend() {
+    let (tensor, mask, tensor_ref, mask_ref) = inputs_mask_fill();
+    let dtype_bool = <<TestBackend as Backend>::BoolElem as Element>::dtype();
+    let dtype_ft = <FloatElem as Element>::dtype();
+
+    let actual = Tensor::<TestBackend, 3>::from_primitive(TensorPrimitive::Float(mask_fill(
+        tensor.into_primitive().tensor(),
+        mask.into_primitive(),
+        InputScalar::new(4.0, dtype_ft),
+        MaskFillStrategy::Readonly,
+        dtype_bool,
+    )));
+    let expected = tensor_ref.mask_fill(mask_ref, 4.0);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[test]
+fn mask_fill_inplace_should_match_reference_backend() {
+    let (tensor, mask, tensor_ref, mask_ref) = inputs_mask_fill();
+    let dtype_bool = <<TestBackend as Backend>::BoolElem as Element>::dtype();
+    let dtype_ft = <FloatElem as Element>::dtype();
+
+    let actual = Tensor::<TestBackend, 3>::from_primitive(TensorPrimitive::Float(mask_fill::<_>(
+        tensor.into_primitive().tensor(),
+        mask.into_primitive(),
+        InputScalar::new(4.0, dtype_ft),
+        MaskFillStrategy::Inplace,
+        dtype_bool,
+    )));
+    let expected = tensor_ref.mask_fill(mask_ref, 4.0);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[allow(clippy::type_complexity)]
+fn inputs_mask_fill() -> (
+    Tensor<TestBackend, 3>,
+    Tensor<TestBackend, 3, Bool>,
+    Tensor<ReferenceBackend, 3>,
+    Tensor<ReferenceBackend, 3, Bool>,
+) {
+    let test_device = Default::default();
+    let tensor = Tensor::<TestBackend, 3>::random([2, 6, 256], Distribution::Default, &test_device);
+    let mask =
+        Tensor::<TestBackend, 3>::random([2, 6, 256], Distribution::Uniform(0., 1.), &test_device)
+            .lower_equal_elem(0.5);
+    let ref_device = Default::default();
+    let tensor_ref = Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &ref_device);
+    let mask_ref = Tensor::<ReferenceBackend, 3, Bool>::from_data(mask.to_data(), &ref_device);
+
+    (tensor, mask, tensor_ref, mask_ref)
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/mask_where.rs

@@ -0,0 +1,80 @@
+use super::*;
+use burn_cubecl::kernel::{MaskWhereStrategy, mask_where};
+use burn_tensor::Tolerance;
+use burn_tensor::{Bool, Distribution, Element, Tensor, TensorPrimitive, backend::Backend};
+
+#[test]
+fn mask_where_should_match_reference_backend() {
+    let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where();
+
+    let actual = tensor.mask_where(mask, value);
+    let expected = tensor_ref.mask_where(mask_ref, value_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+#[test]
+fn mask_where_inplace_lhs_should_match_reference_backend() {
+    let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where();
+    let dtype_bool = <<TestBackend as Backend>::BoolElem as Element>::dtype();
+
+    let actual = Tensor::<TestBackend, 3>::from_primitive(TensorPrimitive::Float(mask_where::<_>(
+        tensor.into_primitive().tensor(),
+        mask.into_primitive(),
+        value.into_primitive().tensor(),
+        MaskWhereStrategy::InplaceLhs,
+        dtype_bool,
+    )));
+    let expected = tensor_ref.mask_where(mask_ref, value_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[test]
+fn mask_where_inplace_rhs_should_match_reference_backend() {
+    let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where();
+    let dtype_bool = <<TestBackend as Backend>::BoolElem as Element>::dtype();
+
+    let actual = Tensor::<TestBackend, 3>::from_primitive(TensorPrimitive::Float(mask_where::<_>(
+        tensor.into_primitive().tensor(),
+        mask.into_primitive(),
+        value.into_primitive().tensor(),
+        MaskWhereStrategy::InplaceRhs,
+        dtype_bool,
+    )));
+    let expected = tensor_ref.mask_where(mask_ref, value_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[allow(clippy::type_complexity)]
+fn inputs_mask_where() -> (
+    Tensor<TestBackend, 3>,
+    Tensor<TestBackend, 3>,
+    Tensor<TestBackend, 3, Bool>,
+    Tensor<ReferenceBackend, 3>,
+    Tensor<ReferenceBackend, 3>,
+    Tensor<ReferenceBackend, 3, Bool>,
+) {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let tensor = Tensor::<TestBackend, 3>::random([2, 6, 256], Distribution::Default, &device);
+    let value = Tensor::<TestBackend, 3>::random([2, 6, 256], Distribution::Default, &device);
+    let mask =
+        Tensor::<TestBackend, 3>::random([2, 6, 256], Distribution::Uniform(0., 1.), &device)
+            .lower_equal_elem(0.5);
+
+    let device_ref = Default::default();
+    let tensor_ref = Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &device_ref);
+    let value_ref = Tensor::<ReferenceBackend, 3>::from_data(value.to_data(), &device_ref);
+    let mask_ref = Tensor::<ReferenceBackend, 3, Bool>::from_data(mask.to_data(), &device_ref);
+    mask.to_data().assert_eq(&mask_ref.to_data(), false);
+
+    (tensor, value, mask, tensor_ref, value_ref, mask_ref)
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/max_pool2d.rs

@@ -0,0 +1,52 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, module};
+
+#[test]
+pub fn max_pool2d_should_match_reference_backends() {
+    let tensor = Tensor::<TestBackend, 4>::random(
+        [32, 32, 32, 32],
+        Distribution::Default,
+        &Default::default(),
+    );
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &Default::default());
+    let kernel_size = [3, 3];
+    let stride = [2, 2];
+    let padding = [1, 1];
+    let dilation = [1, 1];
+
+    let pooled = module::max_pool2d(tensor, kernel_size, stride, padding, dilation, false);
+    let pooled_ref = module::max_pool2d(tensor_ref, kernel_size, stride, padding, dilation, false);
+
+    pooled
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&pooled_ref.into_data(), Tolerance::default());
+}
+
+#[test]
+pub fn max_pool2d_with_indices_should_match_reference_backend() {
+    let tensor = Tensor::<TestBackend, 4>::random(
+        [32, 32, 32, 32],
+        Distribution::Default,
+        &Default::default(),
+    );
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &Default::default());
+    let kernel_size = [3, 3];
+    let stride = [2, 2];
+    let padding = [1, 1];
+    let dilation = [1, 1];
+
+    let (pooled, indices) =
+        module::max_pool2d_with_indices(tensor, kernel_size, stride, padding, dilation, false);
+    let (pooled_ref, indices_ref) =
+        module::max_pool2d_with_indices(tensor_ref, kernel_size, stride, padding, dilation, false);
+
+    pooled
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&pooled_ref.into_data(), Tolerance::default());
+    indices
+        .into_data()
+        .assert_eq(&indices_ref.into_data(), false);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/max_pool2d_backward.rs

@@ -0,0 +1,67 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor, TensorPrimitive, module, ops::ModuleOps};
+
+#[test]
+pub fn max_pool2d_with_indices_backward_should_match_reference_backend() {
+    let test_device = Default::default();
+    let tensor =
+        Tensor::<TestBackend, 4>::random([32, 32, 32, 32], Distribution::Default, &test_device);
+    let grad_output =
+        Tensor::<TestBackend, 4>::random([32, 32, 16, 16], Distribution::Default, &test_device);
+    let ref_device = Default::default();
+    let tensor_ref = Tensor::<ReferenceBackend, 4>::from_data(tensor.to_data(), &ref_device);
+    let grad_output_ref =
+        Tensor::<ReferenceBackend, 4>::from_data(grad_output.to_data(), &ref_device);
+    let kernel_size = [3, 3];
+    let stride = [2, 2];
+    let padding = [1, 1];
+    let dilation = [1, 1];
+
+    let (_, indices) = module::max_pool2d_with_indices(
+        tensor.clone(),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        false,
+    );
+    let (_, indices_ref) = module::max_pool2d_with_indices(
+        tensor_ref.clone(),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        false,
+    );
+    let grad = TestBackend::max_pool2d_with_indices_backward(
+        tensor.into_primitive().tensor(),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        false,
+        grad_output.into_primitive().tensor(),
+        indices.into_primitive(),
+    )
+    .x_grad;
+    let grad_ref = ReferenceBackend::max_pool2d_with_indices_backward(
+        tensor_ref.into_primitive().tensor(),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        false,
+        grad_output_ref.into_primitive().tensor(),
+        indices_ref.into_primitive(),
+    )
+    .x_grad;
+
+    Tensor::<TestBackend, 4>::from_primitive(TensorPrimitive::Float(grad))
+        .into_data()
+        .assert_approx_eq::<FloatElem>(
+            &Tensor::<ReferenceBackend, 4>::from_primitive(TensorPrimitive::Float(grad_ref))
+                .into_data(),
+            Tolerance::default(),
+        );
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/mod.rs

@@ -0,0 +1,30 @@
+// #[allow(unused_imports)] // required for re-included modules
+pub use super::*;
+
+mod avg_pool2d;
+mod bernoulli;
+mod cast;
+mod cat;
+mod clamp;
+mod contiguous;
+mod conv2d;
+mod conv3d;
+mod conv_transpose2d;
+mod conv_transpose3d;
+mod cross;
+mod gather;
+mod mask_fill;
+mod mask_where;
+mod max_pool2d;
+mod max_pool2d_backward;
+mod normal;
+mod quantization;
+mod reduce;
+mod repeat_dim;
+mod scatter;
+mod select;
+mod select_assign;
+mod slice;
+mod slice_assign;
+mod unary;
+mod uniform;

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/normal.rs

@@ -0,0 +1,36 @@
+use super::*;
+use burn_tensor::{Distribution, Shape, Tensor, backend::Backend};
+use cubek::random::{assert_mean_approx_equal, assert_normal_respects_68_95_99_rule};
+use serial_test::serial;
+
+#[test]
+#[serial]
+fn empirical_mean_close_to_expectation() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let shape = [100, 100];
+    let mean = 10.;
+    let tensor = Tensor::<TestBackend, 2>::random(shape, Distribution::Normal(mean, 2.), &device)
+        .into_data();
+    let numbers = tensor.as_slice::<FloatElem>().unwrap();
+
+    assert_mean_approx_equal(numbers, mean as f32);
+}
+
+#[test]
+#[serial]
+fn normal_respects_68_95_99_rule() {
+    // https://en.wikipedia.org/wiki/68%E2%80%9395%E2%80%9399.7_rule
+    let shape: Shape = [1000, 1000].into();
+    let device = Default::default();
+    let mu = 0.;
+    let s = 1.;
+    let tensor =
+        Tensor::<TestBackend, 2>::random(shape.clone(), Distribution::Normal(mu, s), &device)
+            .into_data();
+
+    let numbers = tensor.as_slice::<FloatElem>().unwrap();
+
+    assert_normal_respects_68_95_99_rule(numbers, mu as f32, s as f32);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/quantization.rs

@@ -0,0 +1,240 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{
+    Shape, Tensor,
+    backend::Backend,
+    quantization::{QuantLevel, QuantScheme, QuantStore, QuantValue},
+};
+
+fn should_quantize_dequantize_symmetric_arange<S: Into<Shape>>(
+    value: QuantValue,
+    store: QuantStore,
+    shape: S,
+) {
+    let shape = shape.into();
+    assert_eq!(shape.rank(), 2); // 2D tests
+
+    let scheme = QuantScheme::default().with_value(value).with_store(store);
+    let scheme_ref = scheme.clone().with_store(QuantStore::Native);
+
+    let input: Tensor<TestBackend, 2> =
+        Tensor::arange(0..shape.num_elements() as i64, &Default::default())
+            .float()
+            .reshape(shape);
+    let input_ref = Tensor::<ReferenceBackend, 2>::from_data(input.to_data(), &Default::default());
+
+    let output = input.quantize_dynamic(&scheme);
+    let output_ref = input_ref.quantize_dynamic(&scheme_ref);
+
+    output.to_data().assert_eq(&output_ref.to_data(), false);
+
+    let output = output.dequantize();
+    let output_ref = output_ref.dequantize();
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.to_data(), Tolerance::default());
+}
+
+fn should_quantize_dequantize_symmetric_per_block_arange<S: Into<Shape>>(
+    value: QuantValue,
+    block_size: usize,
+    store: QuantStore,
+    shape: S,
+) {
+    let scheme = QuantScheme::default()
+        .with_value(value)
+        .with_level(QuantLevel::block([block_size as u8]))
+        .with_store(store);
+    let scheme_ref = scheme.clone().with_store(QuantStore::Native);
+
+    let shape = shape.into();
+    let input: Tensor<TestBackend, 2> =
+        Tensor::arange(0..shape.num_elements() as i64, &Default::default())
+            .float()
+            .reshape(shape);
+    let input_ref = Tensor::<ReferenceBackend, 2>::from_data(input.to_data(), &Default::default());
+
+    let output = input.quantize_dynamic(&scheme);
+    let output_ref = input_ref.quantize_dynamic(&scheme_ref);
+
+    output.to_data().assert_eq(&output_ref.to_data(), false);
+
+    let output = output.dequantize();
+    let output_ref = output_ref.dequantize();
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.to_data(), Tolerance::default());
+}
+
+fn should_quantize_dequantize_symmetric_per_block(
+    value: QuantValue,
+    block_size: usize,
+    store: QuantStore,
+) {
+    let scheme = QuantScheme::default()
+        .with_value(value)
+        .with_level(QuantLevel::block([block_size as u8]))
+        .with_store(store);
+    let scheme_ref = scheme.clone().with_store(QuantStore::Native);
+
+    let input = Tensor::<TestBackend, 2>::from_floats(
+        [
+            [
+                -1.8, -1.0, 0.0, 0.5, -1.8, -1.0, 0.0, 0.5, 0.01, 0.025, 0.03, 0.04, 0.01, 0.025,
+                0.03, 0.04,
+            ],
+            [
+                1.8, 1.0, 0.0, -0.5, 1.8, 1.0, 0.0, -0.5, -0.01, -0.025, -0.03, -0.04, -0.01,
+                -0.025, -0.03, -0.04,
+            ],
+        ],
+        &Default::default(),
+    );
+    let input_ref = Tensor::<ReferenceBackend, 2>::from_data(input.to_data(), &Default::default());
+
+    let output = input.quantize_dynamic(&scheme);
+    let output_ref = input_ref.quantize_dynamic(&scheme_ref);
+
+    output.to_data().assert_eq(&output_ref.to_data(), false);
+
+    let output = output.dequantize();
+    let output_ref = output_ref.dequantize();
+
+    output
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&output_ref.to_data(), Tolerance::default());
+}
+
+fn supports_native() -> bool {
+    let name = <TestBackend as Backend>::name(&Default::default());
+    // TODO: Proper checks for i8 support.
+    name.contains("cuda")
+        || name.contains("rocm")
+        || name.contains("hip")
+        || name.contains("vulkan")
+        || name.contains("spirv")
+        || name.contains("metal")
+        || name.contains("msl")
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q8s_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q8S, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q8f_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q8F, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q4s_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q4S, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q4f_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q4F, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q2s_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q2S, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q2f_packed() {
+    should_quantize_dequantize_symmetric_arange(QuantValue::Q2F, QuantStore::PackedU32(0), [8, 16])
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_q8s_packed() {
+    should_quantize_dequantize_symmetric_per_block(QuantValue::Q8S, 8, QuantStore::PackedU32(0))
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_q4s_packed() {
+    should_quantize_dequantize_symmetric_per_block(QuantValue::Q4S, 8, QuantStore::PackedU32(0))
+}
+
+#[test]
+#[should_panic = "Block size must be divisible by 16"]
+fn should_panic_when_block_size_cannot_store_num_quants() {
+    // num_quants in u32 = 32 bits / 2 bits = 16
+    should_quantize_dequantize_symmetric_per_block(QuantValue::Q2S, 8, QuantStore::PackedU32(0))
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_q2s_packed() {
+    should_quantize_dequantize_symmetric_per_block(QuantValue::Q2S, 16, QuantStore::PackedU32(0))
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_q8s_native() {
+    if supports_native() {
+        should_quantize_dequantize_symmetric_arange(QuantValue::Q8S, QuantStore::Native, [32, 32])
+    }
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_q8s_native() {
+    if supports_native() {
+        should_quantize_dequantize_symmetric_per_block(QuantValue::Q8S, 8, QuantStore::Native)
+    }
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_arange_q8s_packed() {
+    should_quantize_dequantize_symmetric_per_block_arange(
+        QuantValue::Q8S,
+        32,
+        QuantStore::PackedU32(0),
+        [32, 32],
+    )
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_per_block_arange_q8s_native() {
+    if supports_native() {
+        should_quantize_dequantize_symmetric_per_block_arange(
+            QuantValue::Q8S,
+            32,
+            QuantStore::Native,
+            [32, 32],
+        )
+    }
+}
+
+#[test]
+fn should_quantize_dequantize_symmetric_arange_128x256_q8s_native() {
+    if supports_native() {
+        should_quantize_dequantize_symmetric_per_block_arange(
+            QuantValue::Q8S,
+            32,
+            QuantStore::Native,
+            [128, 256],
+        )
+    }
+}
+#[test]
+fn should_quantize_dequantize_symmetric_arange_128x256_q8s_packed() {
+    should_quantize_dequantize_symmetric_per_block_arange(
+        QuantValue::Q8S,
+        32,
+        QuantStore::PackedU32(0),
+        [128, 256],
+    )
+}
+
+#[test]
+#[should_panic = "Can't store in u32"]
+fn should_panic_when_shape_cannot_store_quants() {
+    let device = Default::default();
+    let scheme = QuantScheme::default();
+
+    let _tensor_1 =
+        Tensor::<TestBackend, 2>::from_floats([[1.0, 6.35], [2.0, 3.0], [1.0, 3.0]], &device)
+            .quantize_dynamic(&scheme);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/reduce.rs

@@ -0,0 +1,135 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor};
+
+const RANK: usize = 4;
+const SHAPE: [usize; RANK] = [2, 4, 8, 16];
+
+#[test]
+fn reduction_argmax_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    for dim in 0..RANK {
+        tensor
+            .clone()
+            .argmax(dim)
+            .into_data()
+            .assert_eq(&tensor_ref.clone().argmax(dim).into_data(), false);
+    }
+}
+
+#[test]
+fn reduction_argmin_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    for dim in 0..RANK {
+        tensor
+            .clone()
+            .argmin(dim)
+            .into_data()
+            .assert_eq(&tensor_ref.clone().argmin(dim).into_data(), false);
+    }
+}
+
+#[test]
+fn reduction_mean_dim_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    for dim in 0..RANK {
+        tensor
+            .clone()
+            .mean_dim(dim)
+            .into_data()
+            .assert_approx_eq::<FloatElem>(
+                &tensor_ref.clone().mean_dim(dim).into_data(),
+                Tolerance::default(),
+            );
+    }
+}
+
+#[test]
+fn reduction_mean_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    tensor
+        .clone()
+        .mean()
+        .into_data()
+        .assert_approx_eq::<FloatElem>(
+            &tensor_ref.clone().mean().into_data(),
+            Tolerance::default(),
+        );
+}
+
+#[test]
+fn reduction_prod_dim_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    for dim in 0..RANK {
+        tensor
+            .clone()
+            .prod_dim(dim)
+            .into_data()
+            .assert_approx_eq::<FloatElem>(
+                &tensor_ref.clone().prod_dim(dim).into_data(),
+                Tolerance::default(),
+            );
+    }
+}
+
+#[test]
+fn reduction_prod_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    tensor
+        .clone()
+        .prod()
+        .into_data()
+        .assert_approx_eq::<FloatElem>(
+            &tensor_ref.clone().prod().into_data(),
+            Tolerance::default(),
+        );
+}
+
+#[test]
+fn reduction_sum_dim_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    for dim in 0..RANK {
+        tensor
+            .clone()
+            .sum_dim(dim)
+            .into_data()
+            .assert_approx_eq::<FloatElem>(
+                &tensor_ref.clone().sum_dim(dim).into_data(),
+                Tolerance::default(),
+            );
+    }
+}
+
+#[test]
+fn reduction_sum_should_match_reference_backend() {
+    let tensor =
+        Tensor::<TestBackend, RANK>::random(SHAPE, Distribution::Default, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, RANK>::from_data(tensor.to_data(), &Default::default());
+    tensor
+        .clone()
+        .sum()
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&tensor_ref.clone().sum().into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/repeat_dim.rs

@@ -0,0 +1,71 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor};
+
+#[test]
+fn repeat_dim_0_few_times() {
+    let tensor =
+        Tensor::<TestBackend, 3>::random([1, 6, 6], Distribution::Default, &Default::default());
+    let dim = 0;
+    let times = 4;
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &Default::default());
+
+    let actual = tensor.repeat_dim(dim, times);
+    let expected = tensor_ref.repeat_dim(dim, times);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[test]
+fn repeat_dim_1_few_times() {
+    let tensor =
+        Tensor::<TestBackend, 3>::random([6, 1, 6], Distribution::Default, &Default::default());
+    let dim = 1;
+    let times = 4;
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &Default::default());
+
+    let actual = tensor.repeat_dim(dim, times);
+    let expected = tensor_ref.repeat_dim(dim, times);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[test]
+fn repeat_dim_2_few_times() {
+    let tensor =
+        Tensor::<TestBackend, 3>::random([6, 6, 1], Distribution::Default, &Default::default());
+    let dim = 2;
+    let times = 4;
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &Default::default());
+
+    let actual = tensor.repeat_dim(dim, times);
+    let expected = tensor_ref.repeat_dim(dim, times);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+#[test]
+fn repeat_dim_2_many_times() {
+    let tensor =
+        Tensor::<TestBackend, 3>::random([10, 10, 1], Distribution::Default, &Default::default());
+    let dim = 2;
+    let times = 200;
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 3>::from_data(tensor.to_data(), &Default::default());
+
+    let actual = tensor.repeat_dim(dim, times);
+    let expected = tensor_ref.repeat_dim(dim, times);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/scatter.rs

@@ -0,0 +1,66 @@
+use super::*;
+use burn_tensor::{Distribution, Int, Tensor, backend::Backend};
+use burn_tensor::{IndexingUpdateOp, Tolerance};
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_2d_dim0() {
+    same_as_reference_same_shape(0, [256, 32]);
+}
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_2d_dim1() {
+    same_as_reference_same_shape(1, [32, 256]);
+}
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_3d_dim0() {
+    same_as_reference_same_shape(0, [256, 6, 6]);
+}
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_3d_dim1() {
+    same_as_reference_same_shape(1, [6, 256, 6]);
+}
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_3d_dim2() {
+    same_as_reference_same_shape(2, [6, 6, 256]);
+}
+
+#[test]
+fn scatter_should_work_with_multiple_workgroups_diff_shapes() {
+    same_as_reference_diff_shape(1, [32, 128], [32, 1]);
+}
+
+fn same_as_reference_diff_shape<const D: usize>(
+    dim: usize,
+    shape1: [usize; D],
+    shape2: [usize; D],
+) {
+    let test_device = Default::default();
+    TestBackend::seed(&test_device, 0);
+
+    let tensor = Tensor::<TestBackend, D>::random(shape1, Distribution::Default, &test_device);
+    let value = Tensor::<TestBackend, D>::random(shape2, Distribution::Default, &test_device);
+    let indices = Tensor::<TestBackend, 1, Int>::random(
+        [shape2.iter().product::<usize>()],
+        Distribution::Uniform(0., shape2[dim] as f64),
+        &test_device,
+    )
+    .reshape(shape2);
+    let ref_device = Default::default();
+    let tensor_ref = Tensor::<ReferenceBackend, D>::from_data(tensor.to_data(), &ref_device);
+    let value_ref = Tensor::<ReferenceBackend, D>::from_data(value.to_data(), &ref_device);
+    let indices_ref = Tensor::<ReferenceBackend, D, Int>::from_data(indices.to_data(), &ref_device);
+
+    let actual = tensor.scatter(dim, indices, value, IndexingUpdateOp::Add);
+    let expected = tensor_ref.scatter(dim, indices_ref, value_ref, IndexingUpdateOp::Add);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}
+
+fn same_as_reference_same_shape<const D: usize>(dim: usize, shape: [usize; D]) {
+    same_as_reference_diff_shape(dim, shape, shape);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/select.rs

@@ -0,0 +1,21 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Int, Tensor};
+
+#[test]
+fn select_should_work_with_multiple_workgroups() {
+    let tensor =
+        Tensor::<TestBackend, 2>::random([6, 256], Distribution::Default, &Default::default());
+    let indices = Tensor::<TestBackend, 1, Int>::arange(0..100, &Default::default());
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 2>::from_data(tensor.to_data(), &Default::default());
+    let indices_ref =
+        Tensor::<ReferenceBackend, 1, Int>::from_data(indices.to_data(), &Default::default());
+
+    let actual = tensor.select(1, indices);
+    let expected = tensor_ref.select(1, indices_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/select_assign.rs

@@ -0,0 +1,54 @@
+use super::*;
+use burn_tensor::{Distribution, Int, Tensor, backend::Backend};
+use burn_tensor::{IndexingUpdateOp, Tolerance};
+
+#[test]
+fn select_add_should_work_with_multiple_workgroups_2d_dim0() {
+    select_add_same_as_ref(0, [256, 6]);
+}
+
+#[test]
+fn select_add_should_work_with_multiple_workgroups_2d_dim1() {
+    select_add_same_as_ref(1, [6, 256]);
+}
+
+#[test]
+fn select_add_should_work_with_multiple_workgroups_3d_dim0() {
+    select_add_same_as_ref(0, [256, 6, 6]);
+}
+
+#[test]
+fn select_add_should_work_with_multiple_workgroups_3d_dim1() {
+    select_add_same_as_ref(1, [6, 256, 6]);
+}
+
+#[test]
+fn select_add_should_work_with_multiple_workgroups_3d_dim2() {
+    select_add_same_as_ref(2, [6, 6, 256]);
+}
+
+fn select_add_same_as_ref<const D: usize>(dim: usize, shape: [usize; D]) {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+
+    let tensor =
+        Tensor::<TestBackend, D>::random(shape, Distribution::Default, &Default::default());
+    let value = Tensor::<TestBackend, D>::random(shape, Distribution::Default, &Default::default());
+    let indices = Tensor::<TestBackend, 1, Int>::random(
+        [shape[dim]],
+        Distribution::Uniform(0., shape[dim] as f64),
+        &Default::default(),
+    );
+    let tensor_ref =
+        Tensor::<ReferenceBackend, D>::from_data(tensor.to_data(), &Default::default());
+    let value_ref = Tensor::<ReferenceBackend, D>::from_data(value.to_data(), &Default::default());
+    let indices_ref =
+        Tensor::<ReferenceBackend, 1, Int>::from_data(indices.to_data(), &Default::default());
+
+    let actual = tensor.select_assign(dim, indices, value, IndexingUpdateOp::Add);
+    let expected = tensor_ref.select_assign(dim, indices_ref, value_ref, IndexingUpdateOp::Add);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/slice.rs

@@ -0,0 +1,19 @@
+use super::*;
+use burn_tensor::Tolerance;
+use burn_tensor::{Distribution, Tensor};
+
+#[test]
+fn slice_should_work_with_multiple_workgroups() {
+    let tensor =
+        Tensor::<TestBackend, 2>::random([6, 256], Distribution::Default, &Default::default());
+    let indices = [3..5, 45..256];
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 2>::from_data(tensor.to_data(), &Default::default());
+
+    let actual = tensor.slice(indices.clone());
+    let expected = tensor_ref.slice(indices);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/slice_assign.rs

@@ -0,0 +1,21 @@
+use super::*;
+use burn_tensor::{Distribution, Tensor, Tolerance};
+
+#[test]
+fn slice_assign_should_work_with_multiple_workgroups() {
+    let tensor =
+        Tensor::<TestBackend, 2>::random([6, 256], Distribution::Default, &Default::default());
+    let value =
+        Tensor::<TestBackend, 2>::random([2, 211], Distribution::Default, &Default::default());
+    let indices = [3..5, 45..256];
+    let tensor_ref =
+        Tensor::<ReferenceBackend, 2>::from_data(tensor.to_data(), &Default::default());
+    let value_ref = Tensor::<ReferenceBackend, 2>::from_data(value.to_data(), &Default::default());
+
+    let actual = tensor.slice_assign(indices.clone(), value);
+    let expected = tensor_ref.slice_assign(indices, value_ref);
+
+    expected
+        .into_data()
+        .assert_approx_eq::<FloatElem>(&actual.into_data(), Tolerance::default());
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/unary.rs

@@ -0,0 +1,24 @@
+use super::*;
+use burn_tensor::Tensor;
+
+#[test]
+fn tanh_should_not_have_numerical_bugs_on_macos() {
+    fn tanh_one_value(input: f32) -> f32 {
+        let tensor = Tensor::<TestBackend, 1>::ones([1], &Default::default()) * input;
+        let output = tensor.tanh().into_primitive();
+        Tensor::<TestBackend, 1>::from_primitive(output)
+            .into_data()
+            .as_slice()
+            .unwrap()[0]
+    }
+
+    let ok = tanh_one_value(43.0); // metal tanh gives 1.0 which is the right answer
+    let zero = tanh_one_value(44.0); // metal tanh gives zero when within 43.67..44.36
+    let nan = tanh_one_value(45.0); // metal tanh gives nan when over 44.36
+    let neg = tanh_one_value(-45.0); //  metal works correctly here
+
+    assert!(!ok.is_nan() && ok == 1.0);
+    assert!(!zero.is_nan() && zero == 1.0);
+    assert!(!nan.is_nan() && nan == 1.0);
+    assert!(!neg.is_nan() && neg == -1.0);
+}

crates/stable-diffusion-burn/burn-crates/burn-backend-tests/tests/cubecl/uniform.rs

@@ -0,0 +1,113 @@
+use super::*;
+use burn_tensor::{Distribution, Int, Shape, Tensor, backend::Backend};
+use burn_tensor::{ElementConversion, Tolerance};
+
+use serial_test::serial;
+
+use cubek::random::{assert_at_least_one_value_per_bin, assert_wald_wolfowitz_runs_test};
+
+#[test]
+#[serial]
+fn values_all_within_interval_default() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = [24, 24];
+
+    let tensor = Tensor::<TestBackend, 2>::random(shape, Distribution::Default, &device);
+    tensor
+        .to_data()
+        .assert_within_range::<FloatElem>(0.elem()..1.elem());
+}
+
+#[test]
+#[serial]
+fn values_all_within_interval_uniform() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = [24, 24];
+
+    let tensor = Tensor::<TestBackend, 2>::random(shape, Distribution::Uniform(5., 17.), &device);
+    tensor
+        .to_data()
+        .assert_within_range::<FloatElem>(5.elem()..17.elem());
+}
+
+#[test]
+#[serial]
+fn at_least_one_value_per_bin_uniform() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = [64, 64];
+
+    let tensor = Tensor::<TestBackend, 2>::random(shape, Distribution::Uniform(-5., 10.), &device)
+        .into_data();
+    let numbers = tensor.as_slice::<FloatElem>().unwrap();
+
+    assert_at_least_one_value_per_bin(numbers, 3, -5., 10.);
+}
+
+#[test]
+#[serial]
+fn runs_test() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = Shape::new([512, 512]);
+    let tensor =
+        Tensor::<TestBackend, 2>::random(shape, Distribution::Default, &device).into_data();
+
+    let numbers = tensor.as_slice::<FloatElem>().unwrap();
+
+    assert_wald_wolfowitz_runs_test(numbers, 0., 1.);
+}
+
+#[test]
+#[serial]
+fn int_values_all_within_interval_uniform() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = Shape::new([20, 20]);
+    let tensor: Tensor<TestBackend, 2, Int> = Tensor::random(shape, Distribution::Default, &device);
+
+    let data_float = tensor.float().into_data();
+
+    data_float.assert_within_range(0..255);
+}
+
+#[test]
+#[serial]
+fn at_least_one_value_per_bin_int_uniform() {
+    let device = Default::default();
+    TestBackend::seed(&device, 0);
+    let shape = Shape::new([64, 64]);
+
+    let tensor: Tensor<TestBackend, 2, Int> =
+        Tensor::random(shape, Distribution::Uniform(-10.0, 10.0), &device);
+
+    let data_float = tensor.float().into_data();
+
+    let numbers = data_float.as_slice::<FloatElem>().unwrap();
+
+    assert_at_least_one_value_per_bin(numbers, 10, -10., 10.);
+}
+
+#[test]
+fn should_not_fail_on_non_float_autotune() {
+    let device = Default::default();
+    let tensor_1 = Tensor::<TestBackend, 2>::from_floats([[1., 2., 3.], [3., 4., 5.]], &device);
+
+    // Autotune of all (reduce) on lower_equal_elem's output calls uniform distribution
+    tensor_1.lower_equal_elem(1.0).all();
+}
+
+#[test]
+#[serial]
+fn test_seed_reproducibility() {
+    let device = Default::default();
+    TestBackend::seed(&device, 42);
+    let t1 = TestTensor::<1>::random([5], Distribution::Default, &device);
+    TestBackend::seed(&device, 42);
+    let t2 = TestTensor::<1>::random([5], Distribution::Default, &device);
+
+    t1.into_data()
+        .assert_approx_eq::<FloatElem>(&t2.into_data(), Tolerance::default());
+}