Move nn components to burn-nn (tracel-ai#3740)

Author: laggui

.github/workflows/publish.yml

@@ -223,16 +223,45 @@ jobs:
     secrets:
       CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
 
+  publish-burn-collective:
+    uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v3
+    with:
+      crate: burn-collective
+    needs:
+      - publish-burn-common
+      - publish-burn-tensor
+      - publish-burn-communication
+      # dev dependencies
+      - publish-burn-wgpu
+      - publish-burn-ndarray
+      - publish-burn-cuda
+    secrets:
+      CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
+
+  publish-burn-communication:
+    uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v3
+    with:
+      crate: burn-communication
+    needs:
+      - publish-burn-common
+      - publish-burn-tensor
+    secrets:
+      CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
+
   publish-burn-core:
     uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v3
     needs:
       - publish-burn-dataset
       - publish-burn-common
       - publish-burn-derive
       - publish-burn-tensor
+      - publish-burn-vision
+      - publish-burn-collective
+      # dev dependencies
       - publish-burn-autodiff
       - publish-burn-wgpu
       - publish-burn-tch
+      - publish-burn-cuda
       - publish-burn-ndarray
       - publish-burn-candle
       - publish-burn-remote
@@ -241,6 +270,22 @@ jobs:
     secrets:
       CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
 
+  publish-burn-nn:
+    uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v3
+    needs:
+      - publish-burn-core
+      # dev dependencies
+      - publish-burn-autodiff
+      - publish-burn-wgpu
+      - publish-burn-tch
+      - publish-burn-ndarray
+      - publish-burn-candle
+      - publish-burn-remote
+    with:
+      crate: burn-nn
+    secrets:
+      CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }}
+
   publish-burn-train:
     uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v3
     needs:

Cargo.lock

@@ -15,7 +15,6 @@ version.workspace = true
 workspace = true
 
 [features]
-dataset = ["burn-dataset"]
 default = [
     "std",
     "burn-common/default",
@@ -26,12 +25,13 @@ doc = [
     "std",
     "dataset",
     "audio",
-    "vision",
     # Doc features
     "burn-common/doc",
     "burn-dataset/doc",
     "burn-tensor/doc",
 ]
+dataset = ["burn-dataset"]
+
 network = ["burn-common/network"]
 sqlite = ["burn-dataset?/sqlite"]
 sqlite-bundled = ["burn-dataset?/sqlite-bundled"]

crates/burn-core/Cargo.toml

@@ -1,14 +1,15 @@
 # Burn Core
 
-This crate should be used with [burn](https://github.com/tracel-ai/burn).
+This crate should be used with [burn](https://github.com/tracel-ai/burn). It contains the core
+traits and components for building and training deep learning models with Burn.
 
 [![Current Crates.io Version](https://img.shields.io/crates/v/burn-core.svg)](https://crates.io/crates/burn-core)
 [![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-core/blob/master/README.md)
 
 ## Feature Flags
 
-This crate can be used without the standard library (`#![no_std]`) with `alloc` by disabling
-the default `std` feature.
+This crate can be used without the standard library (`#![no_std]`) with `alloc` by disabling the
+default `std` feature.
 
 - `std` - enables the standard library. Enabled by default.
 - `experimental-named-tensor` - enables experimental named tensor.

crates/burn-core/README.md

@@ -31,9 +31,6 @@ pub mod grad_clipping;
 /// Module for the neural network module.
 pub mod module;
 
-/// Neural network module.
-pub mod nn;
-
 /// Module for the recorder.
 pub mod record;
 
@@ -47,7 +44,6 @@ pub use tensor::Tensor;
 pub mod vision;
 
 extern crate alloc;
-extern crate core;
 
 /// Backend for test cases
 #[cfg(all(
@@ -84,6 +80,47 @@ mod tests {
     burn_fusion::memory_checks!();
 }
 
+#[cfg(test)]
+mod test_utils {
+    use crate as burn;
+    use crate::module::Module;
+    use crate::module::Param;
+    use burn_tensor::Tensor;
+    use burn_tensor::backend::Backend;
+    use burn_tensor::module::linear;
+
+    /// Simple linear module.
+    #[derive(Module, Debug)]
+    pub struct SimpleLinear<B: Backend> {
+        pub weight: Param<Tensor<B, 2>>,
+        pub bias: Option<Param<Tensor<B, 1>>>,
+    }
+
+    impl<B: Backend> SimpleLinear<B> {
+        pub fn new(in_features: usize, out_features: usize, device: &B::Device) -> Self {
+            let weight = Tensor::random(
+                [out_features, in_features],
+                burn_tensor::Distribution::Default,
+                device,
+            );
+            let bias = Tensor::random([out_features], burn_tensor::Distribution::Default, device);
+
+            Self {
+                weight: Param::from_tensor(weight),
+                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
@@ -97,7 +134,6 @@ pub mod prelude {
     pub use crate::{
         config::Config,
         module::Module,
-        nn,
         tensor::{
             Bool, Device, ElementConversion, Float, Int, RangesArg, Shape, Tensor, TensorData,
             backend::Backend, cast::ToElement, s,

crates/burn-core/src/lib.rs

@@ -63,7 +63,7 @@ macro_rules! module {
 /// parameter B. This will be used by the [derive](burn_derive::Module) attribute to generate the code
 /// necessary to optimize and train the module on any backend.
 ///
-/// ```no_run
+/// ```rust, ignore
 /// // Not necessary when using the burn crate directly.
 /// use burn_core as burn;
 ///

crates/burn-core/src/module/base.rs

@@ -287,7 +287,7 @@ fn qr_decomposition<B: Backend>(
 mod tests {
     use super::*;
 
-    use crate::tensor::{ElementConversion, TensorData};
+    use burn_tensor::{ElementConversion, TensorData};
     use num_traits::Pow;
 
     pub type TB = burn_ndarray::NdArray<f32>;

crates/burn-core/src/nn/initializer.rs renamed to crates/burn-core/src/module/initializer.rs

@@ -1,12 +1,14 @@
 mod base;
 mod display;
+mod initializer;
 mod param;
 mod quantize;
 #[cfg(feature = "std")]
 mod reinit;
 
 pub use base::*;
 pub use display::*;
+pub use initializer::*;
 pub use param::*;
 pub use quantize::*;
 

crates/burn-core/src/module/mod.rs

@@ -34,7 +34,7 @@ fn new_mapper<T, F: Fn(T) -> T + Send + Sync + 'static>(func: F) -> Mapper<T> {
 /// # Laziness
 ///
 /// The initialization of parameters can be lazy when created using
-/// [uninitialized](Self::uninitialized), which can be done using an [initializer](crate::nn::Initializer).
+/// [uninitialized](Self::uninitialized), which can be done using an [initializer](crate::module::Initializer).
 ///
 /// This reduces the amount of allocations done when loading a model for inference without having
 /// to create a custom initialization function only for inference.

crates/burn-core/src/module/param/base.rs

@@ -1,15 +1,14 @@
 use alloc::{format, string::ToString};
 use core::{fmt::Display, marker::PhantomData};
 
+use crate as burn;
 use crate::{
-    self as burn,
     module::{
         AutodiffModule, Content, Devices, Module, ModuleDisplay, ModuleDisplayDefault,
         ModuleMapper, ModuleVisitor,
     },
-    record::Record,
+    record::{PrecisionSettings, Record},
 };
-use burn::record::PrecisionSettings;
 use burn_tensor::{
     BasicAutodiffOps, BasicOps, Tensor,
     backend::{AutodiffBackend, Backend},