From ea84756d63c22c25ea1693ad863cbc1c21bb4d95 Mon Sep 17 00:00:00 2001
From: nathaniel <nathaniel.simard.42@gmail.com>
Date: Thu, 11 Jul 2024 17:36:36 -0400
Subject: [PATCH] Cleanup

---
 crates/burn-cube/src/frontend/cmma.rs       | 18 +++++++++---------
 crates/burn-cube/src/frontend/subcube.rs    | 20 ++++++++++----------
 crates/burn-cube/src/runtime_tests/slice.rs |  6 +++---
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/crates/burn-cube/src/frontend/cmma.rs b/crates/burn-cube/src/frontend/cmma.rs
index f00e1dce3..5490ea130 100644
--- a/crates/burn-cube/src/frontend/cmma.rs
+++ b/crates/burn-cube/src/frontend/cmma.rs
@@ -86,7 +86,7 @@ impl Init for MatrixExpand {
 
 impl<C: CubePrimitive> Matrix<C> {
     /// Create a new matrix that is going to be used in the
-    /// [matrix-multiply and accumulate](execute) function.
+    /// [matrix-multiply and accumulate](execute()) function.
     ///
     /// You have to declare the shape used for the execution.
     /// The shape of the current matrix is determined using the [MatrixIdent].
@@ -129,11 +129,11 @@ pub fn fill<C: CubeType>(mat: &Matrix<C>, value: C) {
     unexpanded!()
 }
 
-/// Module containing the expand function for [fill].
+/// Module containing the expand function for [fill()].
 pub mod fill {
     use super::*;
 
-    /// Expand method of [fill].
+    /// Expand method of [fill()].
     pub fn __expand<C: CubeType>(
         context: &mut CubeContext,
         mat: MatrixExpand,
@@ -152,11 +152,11 @@ pub fn load<C: CubeType>(mat: &Matrix<C>, value: &Slice<'_, C>, stride: UInt) {
     unexpanded!()
 }
 
-/// Module containing the expand function for [load].
+/// Module containing the expand function for [load()].
 pub mod load {
     use super::*;
 
-    /// Expand method of [load].
+    /// Expand method of [load()].
     #[allow(unused_variables)]
     pub fn __expand<C: CubeType>(
         context: &mut CubeContext,
@@ -183,11 +183,11 @@ pub fn store<C: CubePrimitive>(
     unexpanded!()
 }
 
-/// Module containing the expand function for [store].
+/// Module containing the expand function for [store()].
 pub mod store {
     use super::*;
 
-    /// Expand method of [store].
+    /// Expand method of [store()].
     #[allow(unused_variables)]
     pub fn __expand<C: CubePrimitive>(
         context: &mut CubeContext,
@@ -216,11 +216,11 @@ pub fn execute<A: CubePrimitive, B: CubePrimitive, C: CubePrimitive, D: CubePrim
     unexpanded!()
 }
 
-/// Module containing the expand function for [execute].
+/// Module containing the expand function for [execute()].
 pub mod execute {
     use super::*;
 
-    /// Expand method of [execute].
+    /// Expand method of [execute()].
     pub fn __expand<A: CubePrimitive, B: CubePrimitive, C: CubePrimitive, D: CubePrimitive>(
         context: &mut CubeContext,
         mat_a: MatrixExpand,
diff --git a/crates/burn-cube/src/frontend/subcube.rs b/crates/burn-cube/src/frontend/subcube.rs
index 041c91f6e..3f50154ab 100644
--- a/crates/burn-cube/src/frontend/subcube.rs
+++ b/crates/burn-cube/src/frontend/subcube.rs
@@ -25,11 +25,11 @@ pub fn subcube_sum<E: CubePrimitive>(value: E) -> E {
     unexpanded!()
 }
 
-/// Module containing the expand function for [subcube_sum].
+/// Module containing the expand function for [subcube_sum()].
 pub mod subcube_sum {
     use super::*;
 
-    /// Expand method of [subcube_sum].
+    /// Expand method of [subcube_sum()].
     pub fn __expand<E: CubePrimitive>(
         context: &mut CubeContext,
         elem: ExpandElement,
@@ -53,11 +53,11 @@ pub fn subcube_prod<E: CubePrimitive>(_elem: E) -> E {
     unexpanded!()
 }
 
-/// Module containing the expand function for [subcube_prod].
+/// Module containing the expand function for [subcube_prod()].
 pub mod subcube_prod {
     use super::*;
 
-    /// Expand method of [subcube_prod].
+    /// Expand method of [subcube_prod()].
     pub fn __expand<E: CubePrimitive>(
         context: &mut CubeContext,
         elem: ExpandElement,
@@ -81,11 +81,11 @@ pub fn subcube_max<E: CubePrimitive>(_elem: E) -> E {
     unexpanded!()
 }
 
-/// Module containing the expand function for [subcube_max].
+/// Module containing the expand function for [subcube_max()].
 pub mod subcube_max {
     use super::*;
 
-    /// Expand method of [subcube_max].
+    /// Expand method of [subcube_max()].
     pub fn __expand<E: CubePrimitive>(
         context: &mut CubeContext,
         elem: ExpandElement,
@@ -109,11 +109,11 @@ pub fn subcube_min<E: CubePrimitive>(_elem: E) -> E {
     unexpanded!()
 }
 
-/// Module containing the expand function for [subcube_min].
+/// Module containing the expand function for [subcube_min()].
 pub mod subcube_min {
     use super::*;
 
-    /// Expand method of [subcube_min].
+    /// Expand method of [subcube_min()].
     pub fn __expand<E: CubePrimitive>(
         context: &mut CubeContext,
         elem: ExpandElement,
@@ -137,11 +137,11 @@ pub fn subcube_all<E: CubePrimitive>(_elem: E) -> E {
     unexpanded!()
 }
 
-/// Module containing the expand function for [subcube_all].
+/// Module containing the expand function for [subcube_all()].
 pub mod subcube_all {
     use super::*;
 
-    /// Expand method of [subcube_all].
+    /// Expand method of [subcube_all()].
     pub fn __expand<E: CubePrimitive>(
         context: &mut CubeContext,
         elem: ExpandElement,
diff --git a/crates/burn-cube/src/runtime_tests/slice.rs b/crates/burn-cube/src/runtime_tests/slice.rs
index a8f4ab96c..c46c07fe3 100644
--- a/crates/burn-cube/src/runtime_tests/slice.rs
+++ b/crates/burn-cube/src/runtime_tests/slice.rs
@@ -30,7 +30,7 @@ pub fn test_slice_select<R: Runtime>(client: ComputeClient<R::Server, R::Channel
     let input = client.create(f32::as_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0]));
     let output = client.empty(core::mem::size_of::<f32>());
 
-    slice_select_launch::<F32, R>(
+    slice_select::launch::<F32, R>(
         client.clone(),
         CubeCount::Static(1, 1, 1),
         CubeDim::new(1, 1, 1),
@@ -48,7 +48,7 @@ pub fn test_slice_len<R: Runtime>(client: ComputeClient<R::Server, R::Channel>)
     let input = client.create(f32::as_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0]));
     let output = client.empty(core::mem::size_of::<u32>());
 
-    slice_len_launch::<F32, R>(
+    slice_len::launch::<F32, R>(
         client.clone(),
         CubeCount::Static(1, 1, 1),
         CubeDim::new(1, 1, 1),
@@ -66,7 +66,7 @@ pub fn test_slice_assign<R: Runtime>(client: ComputeClient<R::Server, R::Channel
     let input = client.create(f32::as_bytes(&[15.0]));
     let output = client.create(f32::as_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0]));
 
-    slice_assign_launch::<F32, R>(
+    slice_assign::launch::<F32, R>(
         client.clone(),
         CubeCount::Static(1, 1, 1),
         CubeDim::new(1, 1, 1),