Making f32 default dtype for Tensor, updating examples/docstrings

README.md

@@ -11,7 +11,7 @@ Ergonomics & safety focused deep learning in Rust.
 
 Features at a glance:
 1. Tensor library with shapes up to 6d!
-2. Shapes with both compile and runtime sized dimensions. (e.g. `Tensor<(usize, Const<10>), f32>` and `Tensor<Rank2<5, 10>, f32>`)
+2. Shapes with both compile and runtime sized dimensions. (e.g. `Tensor<(usize, Const<10>)>` and `Tensor<Rank2<5, 10>>`)
 3. A large library of tensor operations (including `matmul`, `conv2d`, and much more).
     1. All tensor operations shape and type checked at compile time!!
 4. Ergonomic neural network building blocks (like `Linear`, `Conv2D`, and `Transformer`).
@@ -65,8 +65,8 @@ type Mlp = (
 fn main() {
     let dev: Cpu = Default::default();
     let mlp: Mlp = dev.build_module();
-    let x: Tensor<Rank1<10>, f32> = dev.zeros();
-    let y /*: Tensor<Rank1<2>, f32>*/ = mlp.forward(x);
+    let x: Tensor<Rank1<10>> = dev.zeros();
+    let y /*: Tensor<Rank1<2>>*/ = mlp.forward(x);
     println!("{:?}", y);
     mlp.save("checkpoint.npz")?;
 }
@@ -81,7 +81,7 @@ let mut sgd = Sgd::new(SgdConfig {
     momentum: Some(Momentum::Nesterov(0.9))
 });
 
-let loss: Tensor<Rank0, f32, _, OwnedTape<_>> = ...
+let loss: Tensor<Rank0, _, _, OwnedTape<_>> = ...
 
 // run backprop to get the gradients
 let gradients = loss.backward();
@@ -90,13 +90,13 @@ sgd.update(&mut model, gradients);
 
 3. 💡 Tensors can be converted to and from normal rust arrays
 ```rust
-let t0: Tensor<Rank0, f32> = dev.tensor(0.0);
+let t0: Tensor<Rank0> = dev.tensor(0.0);
 assert_eq!(t0.array(), &0.0);
 
-let t1 /*: Tensor<Rank1<3>, f32>*/ = dev.tensor([1.0, 2.0, 3.0]);
+let t1 /*: Tensor<Rank1<3>>*/ = dev.tensor([1.0, 2.0, 3.0]);
 assert_eq!(t1.array(), [1.0, 2.0, 3.0]);
 
-let t2: Tensor<Rank2<2, 3>, f32> = dev.randn();
+let t2: Tensor<Rank2<2, 3>> = dev.randn();
 assert_ne!(t2.array(), [[0.0; 3]; 2]);
 ```
 

examples/01-tensor.rs

@@ -15,16 +15,18 @@ fn main() {
     // easily create tensors using the `TensorFromArray::tensor` method of devices
     // tensors are generic over the:
     // 1. Shape (in this case a rank 1 (1 dimension) array with 5 elements)
-    // 2. Data type (in this case f32 values)
-    // 3. The device they are stored on (in this case the Cpu)
-    let _: Tensor<Rank1<5>, f32> = dev.tensor([1.0, 2.0, 3.0, 4.0, 5.0]);
+    // 2. Data type (in this case the default of `f32`)
+    // 3. The device they are stored on (in this case the default of `Cpu`)
+    // 4. A tape - see examples/04-gradients.rs
+    let _: Tensor<Rank1<5>> = dev.tensor([1.0, 2.0, 3.0, 4.0, 5.0]);
 
     // You can also use [ZerosTensor::zeros] and [OnesTensor::ones] to create tensors
     // filled with the corresponding values.
-    let _: Tensor<Rank2<2, 3>, f32> = dev.zeros();
-    let _: Tensor<Rank3<1, 2, 3>, f32> = dev.ones();
+    let _: Tensor<Rank2<2, 3>> = dev.zeros();
+    let _: Tensor<Rank3<1, 2, 3>> = dev.ones();
 
     // each of the creation methods also supports specifying the shape on the function
+    // note to change the dtype we specify the dtype as the 2nd generic parameter
     let _: Tensor<Rank2<2, 3>, f64> = dev.zeros();
     let _ = dev.ones::<Rank2<2, 3>>();
 

examples/03-nn.rs

@@ -18,23 +18,23 @@ fn main() {
 
     // Modules act on tensors using either:
     // 1. `Module::forward`, which does not mutate the module
-    let _: Tensor<Rank1<2>, f32> = m.forward(dev.zeros::<Rank1<4>>());
+    let _: Tensor<Rank1<2>> = m.forward(dev.zeros::<Rank1<4>>());
 
     // 2. `ModuleMut::forward_mut()`, which may mutate the module
-    let _: Tensor<Rank1<2>, f32> = m.forward_mut(dev.zeros::<Rank1<4>>());
+    let _: Tensor<Rank1<2>> = m.forward_mut(dev.zeros::<Rank1<4>>());
 
     // most of them can also act on many different shapes of tensors
     // here we see that Linear can also accept batches of inputs
     // Note: the Rank2 with a batch size of 10 in the input
     //       AND the output
-    let _: Tensor<Rank2<10, 2>, f32> = m.forward(dev.zeros::<Rank2<10, 4>>());
+    let _: Tensor<Rank2<10, 2>> = m.forward(dev.zeros::<Rank2<10, 4>>());
 
     // you can also combine multiple modules with tuples
     let mlp: (Linear<4, 2>, ReLU, Linear<2, 1>) = dev.build_module();
 
     // and of course forward passes the input through each module sequentially:
     let x = dev.randn::<Rank1<4>>();
-    let a: Tensor<Rank1<1>, f32> = mlp.forward(x.clone());
+    let a: Tensor<Rank1<1>> = mlp.forward(x.clone());
     let b = mlp.2.forward(mlp.1.forward(mlp.0.forward(x)));
     assert_eq!(a.array(), b.array());
 }

examples/07-custom-module.rs

@@ -54,12 +54,12 @@ impl<const IN: usize, const INNER: usize, const OUT: usize> GradientUpdate<Cpu,
 }
 
 // impl Module for single item
-impl<const IN: usize, const INNER: usize, const OUT: usize> nn::Module<Tensor<Rank1<IN>, f32, Cpu>>
+impl<const IN: usize, const INNER: usize, const OUT: usize> nn::Module<Tensor<Rank1<IN>>>
     for Mlp<IN, INNER, OUT>
 {
-    type Output = Tensor<Rank1<OUT>, f32, Cpu>;
+    type Output = Tensor<Rank1<OUT>>;
 
-    fn forward(&self, x: Tensor<Rank1<IN>, f32, Cpu>) -> Self::Output {
+    fn forward(&self, x: Tensor<Rank1<IN>>) -> Self::Output {
         let x = self.l1.forward(x);
         let x = self.relu.forward(x);
         self.l2.forward(x)

examples/09-tensor-permute.rs

@@ -7,7 +7,7 @@ use dfdx::tensor_ops::PermuteTo;
 fn main() {
     let dev: Cpu = Default::default();
 
-    let a: Tensor<Rank3<3, 5, 7>, f32> = dev.zeros();
+    let a: Tensor<Rank3<3, 5, 7>> = dev.zeros();
 
     // permuting is as easy as just expressing the desired shape
     // note that we are reversing the order of the axes here!
@@ -19,7 +19,7 @@ fn main() {
     // Just like broadcast/reduce there are times when
     // type inference is impossible because of ambiguities.
     // You can specify axes explicitly to get aroudn this.
-    let c: Tensor<Rank3<1, 1, 1>, f32, _> = dev.zeros();
+    let c: Tensor<Rank3<1, 1, 1>> = dev.zeros();
     let _ = c.permute::<_, Axes3<1, 0, 2>>();
     // NOTE: fails with "Multiple impls satisfying..."
     // let _ = c.permute::<Rank3<1, 1, 1>, _>();

examples/10-tensor-index.rs

@@ -9,7 +9,7 @@ use dfdx::{
 fn main() {
     let dev: Cpu = Default::default();
 
-    let a: Tensor<Rank3<4, 2, 3>, f32> = dev.tensor([
+    let a: Tensor<Rank3<4, 2, 3>> = dev.tensor([
         [[0.00, 0.01, 0.02], [0.10, 0.11, 0.12]],
         [[1.00, 1.01, 1.02], [1.10, 1.11, 1.12]],
         [[2.00, 2.01, 2.02], [2.10, 2.11, 2.12]],

examples/11-multi-headed.rs

@@ -16,5 +16,5 @@ fn main() {
     let m: SplitInto<(Linear<1, 3>, Linear<1, 5>)> = dev.build_module();
 
     // when we forward data through, we get a tuple back!
-    let _: (Tensor<Rank1<3>, f32>, Tensor<Rank1<5>, f32>) = m.forward(dev.tensor([1.0]));
+    let _: (Tensor<Rank1<3>>, Tensor<Rank1<5>>) = m.forward(dev.tensor([1.0]));
 }

examples/nightly-conv-net.rs

@@ -18,12 +18,12 @@ fn main() {
     let m: Model = dev.build_module();
 
     // single image forward
-    let x: Tensor<Rank3<3, 28, 28>, f32> = dev.randn();
-    let _: Tensor<Rank1<10>, f32> = m.forward(x);
+    let x: Tensor<Rank3<3, 28, 28>> = dev.randn();
+    let _: Tensor<Rank1<10>> = m.forward(x);
 
     // batched image forward
-    let x: Tensor<Rank4<32, 3, 28, 28>, f32> = dev.randn();
-    let _: Tensor<Rank2<32, 10>, f32> = m.forward(x);
+    let x: Tensor<Rank4<32, 3, 28, 28>> = dev.randn();
+    let _: Tensor<Rank2<32, 10>> = m.forward(x);
 }
 
 #[cfg(not(feature = "nightly"))]

examples/nightly-transformer.rs

@@ -8,9 +8,9 @@ fn main() {
     let dev: Cpu = Default::default();
     let t: Transformer<16, 4, 3, 3, 8> = dev.build_module();
 
-    let src: Tensor<Rank3<4, 12, 16>, f32> = dev.randn();
-    let tgt: Tensor<Rank3<4, 6, 16>, f32> = dev.randn();
-    let _: Tensor<Rank3<4, 6, 16>, f32, _, _> = t.forward((src.trace(), tgt));
+    let src: Tensor<Rank3<4, 12, 16>> = dev.randn();
+    let tgt: Tensor<Rank3<4, 6, 16>> = dev.randn();
+    let _: Tensor<Rank3<4, 6, 16>, _, _, _> = t.forward((src.trace(), tgt));
 }
 
 #[cfg(not(feature = "nightly"))]

src/lib.rs

@@ -1,6 +1,6 @@
 //! Ergonomics & safety focused deep learning in Rust. Main features include:
 //! 1. Tensor library with shapes up to 6d!
-//! 2. Shapes with both compile and runtime sized dimensions. (e.g. `Tensor<(usize, Const<10>), f32>` and `Tensor<Rank2<5, 10>, f32>`)
+//! 2. Shapes with both compile and runtime sized dimensions. (e.g. `Tensor<(usize, Const<10>)>` and `Tensor<Rank2<5, 10>>`)
 //! 3. A large library of tensor operations (including `matmul`, `conv2d`, and much more).
 //!     a. All tensor operations shape and type checked at compile time!!
 //! 4. Ergonomic neural network building blocks (like `Linear`, `Conv2D`, and `Transformer`).
@@ -15,7 +15,7 @@
 //! # use dfdx::prelude::*;
 //! let dev: Cpu = Default::default();
 //! let x = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
-//! let y: Tensor<Rank2<2, 3>, f32> = dev.ones();
+//! let y: Tensor<Rank2<2, 3>> = dev.ones();
 //! ```
 //!
 //! 2. Neural networks are built with types. Tuples are sequential models. See [crate::nn].
@@ -40,16 +40,16 @@
 //! # use dfdx::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let mlp: Linear<5, 2> = dev.build_module();
-//! let x: Tensor<Rank1<5>, f32> = dev.zeros();
-//! let y = mlp.forward(x); // compiler infers that `y` must be `Tensor<Rank1<2>, f32>`
+//! let x: Tensor<Rank1<5>> = dev.zeros();
+//! let y = mlp.forward(x); // compiler infers that `y` must be `Tensor<Rank1<2>>`
 //! ```
 //!
 //! 5. Trace gradients using [crate::tensor::Tensor::trace()]
 //! ```rust
 //! # use dfdx::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! # let model: Linear<10, 5> = dev.build_module();
-//! # let y_true: Tensor<Rank1<5>, f32> = dev.randn().softmax();
+//! # let y_true: Tensor<Rank1<5>> = dev.randn().softmax();
 //! // tensors default to not having a tape
 //! let x: Tensor<Rank1<10>, f32, Cpu, NoneTape> = dev.zeros();
 //!

src/nn/add_into.rs

@@ -17,7 +17,7 @@ use super::{Module, ModuleMut, ResetParams};
 /// let model: AddInto<(Linear<2, 5>, Linear<3, 5>)> = dev.build_module();
 /// let a = dev.zeros::<Rank1<2>>();
 /// let b = dev.zeros::<Rank1<3>>();
-/// let _: Tensor<Rank1<5>, f32> = model.forward((a, b));
+/// let _: Tensor<Rank1<5>> = model.forward((a, b));
 /// ```
 #[derive(Debug, Default, Clone)]
 pub struct AddInto<T>(pub T);

src/nn/layer_norm.rs

@@ -17,7 +17,7 @@ use super::{Module, ModuleMut, ResetParams};
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let model: LayerNorm1D<5> = dev.build_module();
-/// let _: Tensor<Rank1<5>, f32> = model.forward(dev.zeros::<Rank1<5>>());
+/// let _: Tensor<Rank1<5>> = model.forward(dev.zeros::<Rank1<5>>());
 /// ```
 #[derive(Debug, Clone)]
 pub struct LayerNorm1D<const M: usize, D: Device<f32> = Cpu> {

src/nn/linear.rs

@@ -19,9 +19,9 @@ use super::module::{Module, ModuleMut, ResetParams};
 /// # let dev: Cpu = Default::default();
 /// let model: Linear<5, 2> = dev.build_module();
 /// // single item forward
-/// let _: Tensor<Rank1<2>, f32> = model.forward(dev.zeros::<Rank1<5>>());
+/// let _: Tensor<Rank1<2>> = model.forward(dev.zeros::<Rank1<5>>());
 /// // batched forward
-/// let _: Tensor<Rank2<10, 2>, f32> = model.forward(dev.zeros::<Rank2<10, 5>>());
+/// let _: Tensor<Rank2<10, 2>> = model.forward(dev.zeros::<Rank2<10, 5>>());
 /// ```
 #[derive(Debug, Clone)]
 pub struct Linear<const I: usize, const O: usize, D: Device<f32> = Cpu> {

src/nn/pool_global.rs

@@ -14,8 +14,8 @@ use super::{Module, NonMutableModule, ZeroSizedModule};
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let m: AvgPoolGlobal = Default::default();
-/// let _: Tensor<Rank1<5>, f32> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
-/// let _: Tensor<Rank2<10, 5>, f32> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
+/// let _: Tensor<Rank1<5>> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
+/// let _: Tensor<Rank2<10, 5>> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
 /// ```
 #[derive(Clone, Copy, Default)]
 pub struct AvgPoolGlobal;
@@ -32,8 +32,8 @@ pub struct AvgPoolGlobal;
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let m: MaxPoolGlobal = Default::default();
-/// let _: Tensor<Rank1<5>, f32> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
-/// let _: Tensor<Rank2<10, 5>, f32> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
+/// let _: Tensor<Rank1<5>> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
+/// let _: Tensor<Rank2<10, 5>> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
 /// ```
 #[derive(Clone, Copy, Default)]
 pub struct MaxPoolGlobal;
@@ -50,8 +50,8 @@ pub struct MaxPoolGlobal;
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let m: MinPoolGlobal = Default::default();
-/// let _: Tensor<Rank1<5>, f32> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
-/// let _: Tensor<Rank2<10, 5>, f32> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
+/// let _: Tensor<Rank1<5>> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
+/// let _: Tensor<Rank2<10, 5>> = m.forward(dev.zeros::<Rank4<10, 5, 16, 8>>());
 /// ```
 #[derive(Clone, Copy, Default)]
 pub struct MinPoolGlobal;

src/nn/repeated.rs

@@ -13,7 +13,7 @@ use super::{Module, ModuleMut, ResetParams};
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let model: Repeated<(Linear<10, 10>, ReLU), 5> = dev.build_module();
-/// let out: Tensor<Rank1<10>, f32> = model.forward(dev.zeros());
+/// let out: Tensor<Rank1<10>> = model.forward(dev.zeros());
 /// ```
 #[derive(Debug, Clone)]
 pub struct Repeated<T, const N: usize> {

src/nn/split_into.rs

@@ -16,7 +16,7 @@ use super::{Module, ModuleMut, ResetParams};
 /// # use dfdx::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let model: SplitInto<(Linear<5, 3>, Linear<5, 7>)> = dev.build_module();
-/// let _: (Tensor<Rank1<3>, f32>, Tensor<Rank1<7>, f32>) = model.forward(dev.zeros::<Rank1<5>>());
+/// let _: (Tensor<Rank1<3>>, Tensor<Rank1<7>>) = model.forward(dev.zeros::<Rank1<5>>());
 /// ```
 #[derive(Debug, Default, Clone)]
 pub struct SplitInto<T>(pub T);

src/optim/adam/mod.rs

@@ -56,14 +56,14 @@ impl Default for AdamConfig<f32> {
 /// Constructing using default:
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let mut opt: Adam<Model> = Default::default();
 /// ```
 ///
 /// Changing using new
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let mut opt: Adam<Model> = Adam::new(AdamConfig {
 ///     lr: 1e-2,
 ///     betas: [0.5, 0.25],

src/optim/rmsprop/mod.rs

@@ -64,14 +64,14 @@ impl Default for RMSpropConfig<f32> {
 /// Constructing using default:
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let mut opt: RMSprop<Model> = Default::default();
 /// ```
 ///
 /// Constructing using new:
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let rmsprop: RMSprop<Model> = RMSprop::new(RMSpropConfig {
 ///     lr: 1e-3,
 ///     alpha: 0.5,

src/optim/sgd/mod.rs

@@ -95,14 +95,14 @@ impl Default for SgdConfig<f32> {
 /// Constructing using default:
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let mut opt: Sgd<Model> = Default::default();
 /// ```
 ///
 /// Constructing using new:
 /// ```rust
 /// # use dfdx::{prelude::*, optim::*};
-/// # type Model = Tensor<Rank0, f32>;
+/// # type Model = Tensor<Rank0>;
 /// let mut opt: Sgd<Model> = Sgd::new(SgdConfig {
 ///     lr: 1e-3,
 ///     momentum: Some(Momentum::Classic(0.5)),