Fix cube docs

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>,

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
         }
     )

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)]