[Breaking] Rewrite of nn to enable runtime layer sizes, proc macro declarations, and more

Cargo.toml

@@ -1,87 +1,11 @@
-[package]
-name = "dfdx"
-version = "0.13.0"
-edition = "2021"
-license = "MIT OR Apache-2.0"
-rust-version = "1.65"
+[workspace]
+members = ["dfdx-core", "dfdx-derives", "dfdx"]
+resolver = "2"
 
-description = "Ergonomic auto differentiation in Rust, with pytorch like apis."
-homepage = "https://github.com/coreylowman/dfdx"
-documentation = "https://docs.rs/dfdx"
-repository = "https://github.com/coreylowman/dfdx"
-readme = "README.md"
-
-keywords = [
-    "deep-learning",
-    "neural-network",
-    "backprop",
-    "tensor",
-    "autodiff",
-]
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
-[package.metadata.docs.rs]
-features = ["nightly", "numpy", "safetensors", "cuda", "ci-check"]
-
-[dependencies]
-no-std-compat = { version = "0.4.1", default-features = false, features = [ "alloc", "compat_hash" ], optional = true }
-spin = { version = "0.9.8", default-features = false, features = ["spin_mutex", "rwlock", "portable_atomic"], optional = true }
+[workspace.dependencies]
+num-traits = { version = "0.2.15", default-features = false }
+safetensors = { version = "0.3.3", default-features = false }
+memmap2 = { version = "0.5", default-features = false }
 rand = { version = "0.8.5", default-features = false, features = ["std_rng"] }
 rand_distr = { version = "0.4.3", default-features = false }
-zip = { version = "0.6.6", default-features = false, optional = true }
-cudarc = { version = "0.9.13", default-features = false, optional = true, features = ["driver", "cublas", "nvrtc"] }
-num-traits = { version = "0.2.15", default-features = false }
-safetensors = { version = "0.3", default-features = false, optional = true }
-memmap2 = { version = "0.5", default-features = false, optional = true }
-half = { version = "2.3.1", optional = true, features = ["num-traits", "rand_distr"] }
-gemm = { version = "0.15.4", default-features = false, optional = true }
-rayon = { version = "1.7.0", optional = true }
-libm = "0.2.7"
-
-[dev-dependencies]
-tempfile = "3.3.0"
-mnist = "0.5.0"
-indicatif = "0.17.3"
-
-[build-dependencies]
-glob = { version = "0.3.1", optional = true }
-
-[features]
-default = ["std", "fast-alloc", "cpu"]
-nightly = ["half?/use-intrinsics", "gemm?/nightly"]
-
-std = ["cudarc?/std", "rand_distr/std_math", "gemm?/std"]
-fast-alloc = ["std"]
-no-std = ["no-std-compat", "dep:spin", "cudarc?/no-std", "num-traits/libm"]
-
-cpu = ["dep:gemm", "dep:rayon"]
-cuda = ["dep:cudarc", "dep:glob"]
-cudnn = ["cuda", "cudarc?/cudnn"]
-
-f16 = ["dep:half", "cudarc?/f16"]
-
-numpy = ["dep:zip", "std"]
-safetensors = ["dep:safetensors", "std", "dep:memmap2"]
-
-test-f16 = ["f16"]
-test-amp-f16 = ["f16"]
-test-f64 = []
-test-integrations = []
-ci-check = ["cudarc?/ci-check"]
-
-[[bench]]
-name = "batchnorm2d"
-harness = false
-
-[[bench]]
-name = "conv2d"
-harness = false
-
-[[bench]]
-name = "sum"
-harness = false
-
-[[bench]]
-name = "softmax"
-harness = false
+libm = "0.2.7"

dfdx-core/Cargo.toml

@@ -0,0 +1,72 @@
+[package]
+name = "dfdx-core"
+version = "0.13.0"
+edition = "2021"
+license = "MIT OR Apache-2.0"
+rust-version = "1.65"
+
+description = "Ergonomic auto differentiation in Rust, with pytorch like apis."
+homepage = "https://github.com/coreylowman/dfdx"
+documentation = "https://docs.rs/dfdx"
+repository = "https://github.com/coreylowman/dfdx"
+readme = "README.md"
+
+keywords = [
+    "deep-learning",
+    "neural-network",
+    "backprop",
+    "tensor",
+    "autodiff",
+]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[package.metadata.docs.rs]
+features = ["nightly", "numpy", "safetensors", "cuda", "ci-check"]
+
+[dependencies]
+no-std-compat = { version = "0.4.1", default-features = false, features = [ "alloc", "compat_hash" ], optional = true }
+spin = { version = "0.9.8", default-features = false, features = ["spin_mutex", "rwlock", "portable_atomic"], optional = true }
+rand = { workspace = true }
+rand_distr = { workspace = true }
+zip = { version = "0.6.6", default-features = false, optional = true }
+cudarc = { version = "0.9.13", default-features = false, optional = true, features = ["driver", "cublas", "nvrtc"] }
+num-traits = { workspace = true }
+safetensors = { workspace = true, optional = true }
+memmap2 = { workspace = true, optional = true }
+half = { version = "2.3.1", optional = true, features = ["num-traits", "rand_distr"] }
+gemm = { version = "0.15.4", default-features = false, optional = true }
+rayon = { version = "1.7.0", optional = true }
+libm = { workspace = true }
+
+[dev-dependencies]
+tempfile = "3.3.0"
+mnist = "0.5.0"
+indicatif = "0.17.3"
+
+[build-dependencies]
+glob = { version = "0.3.1", optional = true }
+
+[features]
+default = ["std", "fast-alloc", "cpu"]
+nightly = ["half?/use-intrinsics", "gemm?/nightly"]
+
+std = ["cudarc?/std", "rand_distr/std_math", "gemm?/std"]
+no-std = ["no-std-compat", "dep:spin", "cudarc?/no-std", "num-traits/libm"]
+
+cpu = ["dep:gemm", "dep:rayon"]
+fast-alloc = ["std"]
+
+cuda = ["dep:cudarc", "dep:glob"]
+cudnn = ["cuda", "cudarc?/cudnn"]
+
+f16 = ["dep:half", "cudarc?/f16"]
+
+numpy = ["dep:zip", "std"]
+safetensors = ["dep:safetensors", "std", "dep:memmap2"]
+
+test-f16 = ["f16"]
+test-amp-f16 = ["f16"]
+test-f64 = []
+test-integrations = []
+ci-check = ["cudarc?/ci-check"]

src/data/arange.rs → dfdx-core/src/data/arange.rs

@@ -11,15 +11,15 @@ pub trait Arange<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E> {
     ///
     /// Const sized tensor:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::Arange};
+    /// # use dfdx_core::{prelude::*, data::Arange};
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank1<5>, f32, _> = dev.arange(Const::<5>);
     /// assert_eq!(t.array(), [0.0, 1.0, 2.0, 3.0, 4.0]);
     /// ```
     ///
     /// Runtime sized tensor:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::Arange};
+    /// # use dfdx_core::{prelude::*, data::Arange};
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<(usize, ), f32, _> = dev.arange(5);
     /// assert_eq!(t.as_vec(), [0.0, 1.0, 2.0, 3.0, 4.0]);

src/data/batch.rs → dfdx-core/src/data/batch.rs

@@ -80,14 +80,14 @@ pub trait IteratorBatchExt: Iterator {
     ///
     /// Const batches:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<[usize; 5]> = (0..12).batch_exact(Const::<5>).collect();
     /// assert_eq!(&items, &[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]);
     /// ```
     ///
     /// Runtime batches:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<Vec<usize>> = (0..12).batch_exact(5).collect();
     /// assert_eq!(&items, &[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]);
     /// ```
@@ -104,7 +104,7 @@ pub trait IteratorBatchExt: Iterator {
     ///
     /// Example:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<Vec<usize>> = (0..12).batch_with_last(5).collect();
     /// assert_eq!(&items, &[vec![0, 1, 2, 3, 4], vec![5, 6, 7, 8, 9], vec![10, 11]]);
     /// ```

src/data/collate.rs → dfdx-core/src/data/collate.rs

@@ -93,7 +93,7 @@ pub trait IteratorCollateExt: Iterator {
     ///
     /// Example implementations:
     /// ```rust
-    /// # use dfdx::data::IteratorCollateExt;
+    /// # use dfdx_core::data::IteratorCollateExt;
     /// let data = [[('a', 'b'); 10], [('c', 'd'); 10], [('e', 'f'); 10]];
     /// // we use collate to transform each batch:
     /// let mut iter = data.into_iter().collate();

src/data/one_hot_encode.rs → dfdx-core/src/data/one_hot_encode.rs

@@ -16,7 +16,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Const class labels and const n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = [0, 1, 2, 1, 1];
     /// let probs: Tensor<Rank2<5, 3>, f32, _> = dev.one_hot_encode(Const::<3>, class_labels);
@@ -31,7 +31,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Runtime class labels and const n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = [0, 1, 2, 1, 1];
     /// let probs: Tensor<(Const<5>, usize), f32, _> = dev.one_hot_encode(3, class_labels);
@@ -46,7 +46,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Const class labels and runtime n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = std::vec![0, 1, 2, 1, 1];
     /// let probs: Tensor<(usize, Const<3>), f32, _> = dev.one_hot_encode(Const, class_labels);
@@ -61,7 +61,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Runtime both:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = std::vec![0, 1, 2, 1, 1];
     /// let probs: Tensor<(usize, usize), f32, _> = dev.one_hot_encode(3, class_labels);

src/data/stack.rs → dfdx-core/src/data/stack.rs

@@ -29,7 +29,7 @@ pub trait IteratorStackExt: Iterator {
     ///
     /// Example implementations:
     /// ```rust
-    /// # use dfdx::{data::IteratorStackExt, prelude::*};
+    /// # use dfdx_core::{data::IteratorStackExt, prelude::*};
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank1<3>, f32, _> = dev.zeros();
     /// let data = [[a.clone(), a.clone(), a]];

src/lib.rs → dfdx-core/src/lib.rs

@@ -59,7 +59,7 @@
 //! There are two options for this currently, with more planned to be added in the future:
 //!
 //! 1. [tensor::Cpu] - for tensors stored on the heap
 //! 2. [tensor::Cuda] - for tensors stored in GPU memory
 //!
 //! Both devices implement [Default], you can also create them with a certain seed
 //! and ordinal.
@@ -67,7 +67,7 @@
 //! Here's how you might use a device:
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! let dev: Cpu = Default::default();
 //! let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
 //! ```
@@ -85,8 +85,8 @@
 //! | Unary Operations | `a.sqrt()` | `a.sqrt()` | `a.sqrt()` |
 //! | Binary Operations | `a + b` | `a + b` | `a + b` |
 //! | gemm/gemv | [tensor_ops::matmul] | `a @ b` | `a @ b` |
 //! | 2d Convolution | [tensor_ops::TryConv2D] | - | `torch.conv2d` |
 //! | 2d Transposed Convolution | [tensor_ops::TryConvTrans2D] | - | `torch.conv_transpose2d` |
 //! | Slicing | [tensor_ops::slice] | `a[...]` | `a[...]` |
 //! | Select | [tensor_ops::SelectTo] | `a[...]` | `torch.select` |
 //! | Gather | [tensor_ops::GatherTo] | `np.take` | `torch.gather` |
@@ -100,78 +100,6 @@
 //! | Concat | [tensor_ops::TryConcat] | `np.concatenate` | `torch.concat` |
 //!
 //! and **much much more!**
-//!
-//! # Neural networks
-//!
-//! *See [nn] for more information.*
-//!
-//! Neural networks are composed of building blocks that you can chain together. In
-//! dfdx, sequential neural networks are represents by **tuples**! For example,
-//! the following two networks are identical:
-//!
-//! | dfdx | pytorch |
-//! | --- | --- |
-//! | `(Linear<3, 5>, ReLU, Linear<5, 10>)` | `nn.Sequential(nn.Linear(3, 5), nn.ReLU(), nn.Linear(5, 10))` |
-//! | `((Conv2D<3, 2, 1>, Tanh), Conv2D<3, 2, 1>)` | `nn.Sequential(nn.Sequential(nn.Conv2d(3, 2, 1), nn.Tanh()), nn.Conv2d(3, 2, 1))`
-//!
-//! To build a neural network, you of course need a device:
-//!
-//! ```rust
-//! # use dfdx::prelude::*;
-//! let dev: Cpu = Default::default();
-//! type Model = (Linear<3, 5>, ReLU, Linear<5, 10>);
-//! let model = dev.build_module::<Model, f32>();
-//! ```
-//!
-//! Note two things:
-//! 1. We are using [nn::DeviceBuildExt] to instantiate the model
-//! 2. We **need** to pass a dtype (in this case f32) to create the model.
-//!
-//! You can then pass tensors into the model with [nn::Module::forward()]:
-//!
-//! ```rust
-//! # use dfdx::prelude::*;
-//! # let dev: Cpu = Default::default();
-//! # type Model = (Linear<3, 5>, ReLU, Linear<5, 10>);
-//! # let model = dev.build_module::<Model, f32>();
-//! // tensor with runtime batch dimension of 10
-//! let x: Tensor<(usize, Const<3>), f32, _> = dev.sample_normal_like(&(10, Const));
-//! let y = model.forward(x);
-//! ```
-//!
-//! # Optimizers and Gradients
-//!
-//! *See [optim] for more information*
-//!
-//! dfdx supports a number of the standard optimizers:
-//!
-//! | Optimizer | dfdx | pytorch |
-//! | --- | --- | --- |
-//! | SGD | [optim::Sgd] | `torch.optim.SGD` |
-//! | Adam | [optim::Adam] | torch.optim.Adam` |
-//! | AdamW | [optim::Adam] with [optim::WeightDecay::Decoupled] | `torch.optim.AdamW` |
-//! | RMSprop | [optim::RMSprop] | `torch.optim.RMSprop` |
-//!
-//! You can use optimizers to optimize neural networks (or even tensors!). Here's
-//! a simple example of how to do this with [nn::ZeroGrads]:
-//! ```rust
-//! # use dfdx::{prelude::*, optim::*};
-//! # let dev: Cpu = Default::default();
-//! type Model = (Linear<3, 5>, ReLU, Linear<5, 10>);
-//! let mut model = dev.build_module::<Model, f32>();
-//! // 1. allocate gradients for the model
-//! let mut grads = model.alloc_grads();
-//! // 2. create our optimizer
-//! let mut opt = Sgd::new(&model, Default::default());
-//! // 3. trace gradients through forward pass
-//! let x: Tensor<Rank2<10, 3>, f32, _> = dev.sample_normal();
-//! let y = model.forward_mut(x.traced(grads));
-//! // 4. compute loss & run backpropagation
-//! let loss = y.square().mean();
-//! grads = loss.backward();
-//! // 5. apply gradients
-//! opt.update(&mut model, &grads);
-//! ```
 
 #![cfg_attr(all(feature = "no-std", not(feature = "std")), no_std)]
 #![allow(incomplete_features)]
@@ -185,19 +113,15 @@
 
 pub mod data;
 pub mod dtypes;
-pub mod feature_flags;
 pub mod losses;
-pub mod nn;
-pub mod optim;
+pub mod nn_traits;
 pub mod shapes;
 pub mod tensor;
 pub mod tensor_ops;
 
 /// Contains subset of all public exports.
 pub mod prelude {
     pub use crate::losses::*;
-    pub use crate::nn::builders::*;
-    pub use crate::optim::prelude::*;
     pub use crate::shapes::*;
     pub use crate::tensor::*;
     pub use crate::tensor_ops::*;
@@ -217,8 +141,8 @@
 
     #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
     {
         use std::arch::x86_64::{_MM_FLUSH_ZERO_ON, _MM_SET_FLUSH_ZERO_MODE};
         unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON) }
     }
 }
 
@@ -236,14 +160,14 @@
 
     #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
     {
         use std::arch::x86_64::{_MM_FLUSH_ZERO_OFF, _MM_SET_FLUSH_ZERO_MODE};
         unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_OFF) }
     }
 }
 
 #[cfg(test)]
 pub(crate) mod tests {
-    pub use num_traits::{Float, FromPrimitive, NumCast, Zero};
+    pub use num_traits::{Float, NumCast, Zero};
 
     #[cfg(not(feature = "cuda"))]
     pub type TestDevice = crate::tensor::Cpu;