Move optimizer components to burn-optim (#3773)

* Move optim, grad_clipping and lr_scheduler to burn-optim

* Add publish workflow

* Fix docs link

Author: laggui

.github/workflows/publish.yml

@@ -256,7 +256,6 @@ jobs:
       - publish-burn-derive
       - publish-burn-tensor
       - publish-burn-vision
-      - publish-burn-collective
       # dev dependencies
       - publish-burn-autodiff
       - publish-burn-wgpu
@@ -286,6 +285,23 @@ jobs:
     secrets:
       CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
 
+  publish-burn-optim:
+    uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v4
+    needs:
+      - publish-burn-core
+      - publish-burn-collective
+      # dev dependencies
+      - publish-burn-autodiff
+      - publish-burn-wgpu
+      - publish-burn-tch
+      - publish-burn-ndarray
+      - publish-burn-candle
+      - publish-burn-remote
+    with:
+      crate: burn-optim
+    secrets:
+      CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
+
   publish-burn-train:
     uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v4
     needs:

Cargo.lock

@@ -50,7 +50,6 @@ std = [
 ]
 vision = ["burn-vision", "burn-dataset?/vision", "burn-common/network"]
 audio = ["burn-dataset?/audio"]
-collective = ["burn-collective"]
 
 # Custom deserializer for Record that is helpful for importing data, such as PyTorch pt files.
 record-item-custom-serde = ["thiserror"]
@@ -91,7 +90,6 @@ burn-dataset = { path = "../burn-dataset", version = "0.19.0", optional = true,
 burn-derive = { path = "../burn-derive", version = "0.19.0" }
 burn-tensor = { path = "../burn-tensor", version = "0.19.0", default-features = false }
 burn-vision = { path = "../burn-vision", version = "0.19.0", optional = true, default-features = false }
-burn-collective = { path = "../burn-collective", version = "0.19.0", optional = true, default-features = false }
 
 data-encoding = { workspace = true }
 uuid = { workspace = true }

crates/burn-core/Cargo.toml

@@ -18,16 +18,6 @@ pub mod config;
 #[cfg(feature = "std")]
 pub mod data;
 
-/// Optimizer module.
-pub mod optim;
-
-/// Learning rate scheduler module.
-#[cfg(feature = "std")]
-pub mod lr_scheduler;
-
-/// Gradient clipping module.
-pub mod grad_clipping;
-
 /// Module for the neural network module.
 pub mod module;
 
@@ -87,7 +77,6 @@ mod test_utils {
     use crate::module::Param;
     use burn_tensor::Tensor;
     use burn_tensor::backend::Backend;
-    use burn_tensor::module::linear;
 
     /// Simple linear module.
     #[derive(Module, Debug)]
@@ -110,23 +99,9 @@ mod test_utils {
                 bias: Some(Param::from_tensor(bias)),
             }
         }
-
-        pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
-            linear(
-                input,
-                self.weight.val(),
-                self.bias.as_ref().map(|b| b.val()),
-            )
-        }
     }
 }
 
-/// Type alias for the learning rate.
-///
-/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it
-/// can be used for constant learning rate.
-pub type LearningRate = f64; // We could potentially change the type.
-
 pub mod prelude {
     //! Structs and macros used by most projects. Add `use
     //! burn::prelude::*` to your code to quickly get started with

crates/burn-core/src/lib.rs

@@ -1,7 +1,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 #![warn(missing_docs)]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
-// #![recursion_limit = "135"]
 
 //! Burn neural network module.
 

crates/burn-nn/src/lib.rs

@@ -0,0 +1,91 @@
+[package]
+authors = ["nathanielsimard <[email protected]>"]
+categories = ["science", "no-std", "embedded", "wasm"]
+description = "Optimizer building blocks for the Burn deep learning framework"
+documentation = "https://docs.rs/burn-optim"
+edition.workspace = true
+keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"]
+license.workspace = true
+name = "burn-optim"
+readme.workspace = true
+repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-optim"
+version.workspace = true
+
+[lints]
+workspace = true
+
+[features]
+default = [
+    "std",
+    "burn-core/default",
+]
+doc = [
+    "std",
+    # Doc features
+    "burn-core/doc",
+]
+std = [
+    "burn-core/std",
+    "num-traits/std",
+    "serde/std",
+    "log",
+]
+
+collective = ["burn-collective"]
+
+test-cuda = [
+    "burn-cuda/default",
+] # To use cuda during testing, default uses ndarray.
+test-rocm = [
+    "burn-rocm/default",
+] # To use hip during testing, default uses ndarray.
+test-tch = [
+    "burn-tch/default",
+] # To use tch during testing, default uses ndarray.
+test-wgpu = [
+    "burn-wgpu/default",
+] # To use wgpu during testing, default uses ndarray.
+test-vulkan = [
+    "test-wgpu",
+    "burn-wgpu/vulkan",
+] # To use wgpu-spirv during testing, default uses ndarray.
+test-metal = [
+    "test-wgpu",
+    "burn-wgpu/metal",
+] # To use wgpu-spirv during testing, default uses ndarray.
+
+# Memory checks are disabled by default
+test-memory-checks = ["burn-fusion/memory-checks"]
+
+[dependencies]
+
+# ** Please make sure all dependencies support no_std when std is disabled **
+burn-core = { path = "../burn-core", version = "0.19.0", default-features = false }
+burn-collective = { path = "../burn-collective", version = "0.19.0", optional = true, default-features = false }
+
+num-traits = { workspace = true }
+derive-new = { workspace = true }
+log = { workspace = true, optional = true }
+serde = { workspace = true, features = ["derive"] }
+
+# The same implementation of HashMap in std but with no_std support (only alloc crate is needed)
+hashbrown = { workspace = true, features = ["serde"] } # no_std compatible
+
+# FOR TESTING
+burn-cuda = { path = "../burn-cuda", version = "0.19.0", optional = true, default-features = false }
+burn-rocm = { path = "../burn-rocm", version = "0.19.0", optional = true, default-features = false }
+burn-remote = { path = "../burn-remote", version = "0.19.0", default-features = false, optional = true }
+burn-router = { path = "../burn-router", version = "0.19.0", default-features = false, optional = true }
+burn-tch = { path = "../burn-tch", version = "0.19.0", optional = true }
+burn-wgpu = { path = "../burn-wgpu", version = "0.19.0", optional = true, default-features = false }
+burn-fusion = { path = "../burn-fusion", version = "0.19.0", optional = true }
+
+[dev-dependencies]
+burn-nn = { path = "../burn-nn", version = "0.19.0" }
+burn-ndarray = { path = "../burn-ndarray", version = "0.19.0" }
+burn-autodiff = { path = "../burn-autodiff", version = "0.19.0" }
+rstest = { workspace = true }
+
+[package.metadata.docs.rs]
+features = ["doc"]
+rustdoc-args = ["--cfg", "docsrs"]

crates/burn-optim/Cargo.toml

@@ -0,0 +1,3 @@
+# Burn Optimizers
+
+Core building blocks for Burn optimizers.

crates/burn-optim/README.md

@@ -1,7 +1,7 @@
-use crate as burn;
+use burn_core as burn;
 
-use crate::{config::Config, tensor::Tensor};
-use burn_tensor::backend::Backend;
+use burn::tensor::backend::Backend;
+use burn::{config::Config, tensor::Tensor};
 
 /// Gradient Clipping provides a way to mitigate exploding gradients
 #[derive(Config, Debug)]
@@ -91,7 +91,7 @@ impl GradientClipping {
 mod tests {
     use super::*;
     use crate::TestBackend;
-    use crate::tensor::Tensor;
+    use burn::tensor::Tensor;
 
     #[test]
     fn test_clip_by_value() {

crates/burn-core/src/grad_clipping/base.rs renamed to crates/burn-optim/src/grad_clipping/base.rs

@@ -0,0 +1,62 @@
+#![cfg_attr(not(feature = "std"), no_std)]
+#![warn(missing_docs)]
+#![cfg_attr(docsrs, feature(doc_auto_cfg))]
+
+//! Burn optimizers.
+
+#[macro_use]
+extern crate derive_new;
+
+extern crate alloc;
+
+/// Optimizer module.
+pub mod optim;
+pub use optim::*;
+
+/// Gradient clipping module.
+pub mod grad_clipping;
+
+/// Learning rate scheduler module.
+#[cfg(feature = "std")]
+pub mod lr_scheduler;
+
+/// Type alias for the learning rate.
+///
+/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it
+/// can be used for constant learning rate.
+pub type LearningRate = f64; // We could potentially change the type.
+
+/// Backend for test cases
+#[cfg(all(
+    test,
+    not(feature = "test-tch"),
+    not(feature = "test-wgpu"),
+    not(feature = "test-cuda"),
+    not(feature = "test-rocm")
+))]
+pub type TestBackend = burn_ndarray::NdArray<f32>;
+
+#[cfg(all(test, feature = "test-tch"))]
+/// Backend for test cases
+pub type TestBackend = burn_tch::LibTorch<f32>;
+
+#[cfg(all(test, feature = "test-wgpu"))]
+/// Backend for test cases
+pub type TestBackend = burn_wgpu::Wgpu;
+
+#[cfg(all(test, feature = "test-cuda"))]
+/// Backend for test cases
+pub type TestBackend = burn_cuda::Cuda;
+
+#[cfg(all(test, feature = "test-rocm"))]
+/// Backend for test cases
+pub type TestBackend = burn_rocm::Rocm;
+
+/// Backend for autodiff test cases
+#[cfg(test)]
+pub type TestAutodiffBackend = burn_autodiff::Autodiff<TestBackend>;
+
+#[cfg(all(test, feature = "test-memory-checks"))]
+mod tests {
+    burn_fusion::memory_checks!();
+}