From 0a33aa363d2a040275342ca099669684b8d2c79b Mon Sep 17 00:00:00 2001
From: nathaniel <nathaniel.simard.42@gmail.com>
Date: Fri, 12 Jul 2024 09:25:45 -0400
Subject: [PATCH] Fix cube docs

---
 crates/burn-cube-macros/src/codegen_function/launch.rs | 4 +---
 crates/burn-cube-macros/src/codegen_trait/mod.rs       | 2 +-
 crates/burn-cube-macros/src/lib.rs                     | 8 ++++++--
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/crates/burn-cube-macros/src/codegen_function/launch.rs b/crates/burn-cube-macros/src/codegen_function/launch.rs
index e37b52c45..7f035a55b 100644
--- a/crates/burn-cube-macros/src/codegen_function/launch.rs
+++ b/crates/burn-cube-macros/src/codegen_function/launch.rs
@@ -463,8 +463,7 @@ pub fn codegen_launch(sig: &syn::Signature) -> TokenStream {
     let kernel = codegen.gen_kernel_struct();
     let compile = codegen.gen_compile_impl(&ident_expand);
     let (inputs, output) = (codegen.fn_inputs, codegen.fn_output);
-    let doc =
-        format!("Launch the kernel [{ident}] with the provided argument on the given runtime.");
+    let doc = format!("Launch the kernel [{ident}()] on the given runtime.");
 
     quote::quote! {
         #kernel
@@ -472,7 +471,6 @@ pub fn codegen_launch(sig: &syn::Signature) -> TokenStream {
 
         #[allow(clippy::too_many_arguments)]
         #[doc = #doc]
-        /// Launch the kernel.
         pub fn launch #generics (
             client: ComputeClient<R::Server, R::Channel>,
             cube_count: CubeCount<R::Server>,
diff --git a/crates/burn-cube-macros/src/codegen_trait/mod.rs b/crates/burn-cube-macros/src/codegen_trait/mod.rs
index 4810bc1c6..bba751a47 100644
--- a/crates/burn-cube-macros/src/codegen_trait/mod.rs
+++ b/crates/burn-cube-macros/src/codegen_trait/mod.rs
@@ -103,7 +103,7 @@ fn register_expand(
     quote::quote! (
         #expand {
             #[cube]
-            pub #func
+            #func
             #func_expand
         }
     )
diff --git a/crates/burn-cube-macros/src/lib.rs b/crates/burn-cube-macros/src/lib.rs
index 616fdbbcc..afcdd1bf4 100644
--- a/crates/burn-cube-macros/src/lib.rs
+++ b/crates/burn-cube-macros/src/lib.rs
@@ -143,7 +143,11 @@ fn codegen_cube(
         return Err(code);
     }
 
-    let launch_doc = if launch { "and launch function " } else { "" };
+    let launch_doc = if launch {
+        "and launch functions "
+    } else {
+        "function "
+    };
 
     let launch = if launch {
         codegen_launch(&func.sig)
@@ -153,7 +157,7 @@ fn codegen_cube(
 
     let mod_name = &func.sig.ident;
     let vis = &func.vis;
-    let doc = format!("Module containing the expand method {launch_doc}of {mod_name}.");
+    let doc = format!("Module containing the expand {launch_doc}of {mod_name}.");
 
     Ok(quote::quote! {
         #[allow(dead_code)]