impl Dtype for all Unit types except bool

examples/07-custom-module.rs

@@ -19,8 +19,10 @@ struct Mlp<const IN: usize, const INNER: usize, const OUT: usize, E: Dtype, D: D
 
 // TensorCollection lets you do several operations on Modules, including constructing them with
 // randomized parameters, and iterating through or mutating all tensors in a model.
-impl<const IN: usize, const INNER: usize, const OUT: usize, E: Dtype, D: Device<E>>
-    TensorCollection<E, D> for Mlp<IN, INNER, OUT, E, D>
+impl<const IN: usize, const INNER: usize, const OUT: usize, E, D: Device<E>> TensorCollection<E, D>
+    for Mlp<IN, INNER, OUT, E, D>
+where
+    E: Dtype + num_traits::Float,
 {
     // Type alias that specifies the how Mlp's type changes when using a different dtype and/or
     // device.

src/nn/linear.rs

@@ -56,8 +56,8 @@ impl<const I: usize, const O: usize, E: Dtype, D: DeviceStorage> NonMutableModul
 {
 }
 
-impl<const I: usize, const O: usize, E: Dtype, D: Device<E>> TensorCollection<E, D>
-    for Linear<I, O, E, D>
+impl<const I: usize, const O: usize, E: Dtype + num_traits::Float, D: Device<E>>
+    TensorCollection<E, D> for Linear<I, O, E, D>
 {
     type To<E2: Dtype, D2: Device<E2>> = Linear<I, O, E2, D2>;
 

src/nn/tensor_collection/collection.rs

@@ -23,7 +23,7 @@ use super::{ModuleField, ModuleFields, TensorField};
 ///     relu: ReLU,
 /// }
 ///
-/// impl<E: Dtype, D: Device<E>> TensorCollection<E, D> for Mlp<E, D> {
+/// impl<E: Dtype + num_traits::Float, D: Device<E>> TensorCollection<E, D> for Mlp<E, D> {
 ///     type To<E2: Dtype, D2: Device<E2>> = Mlp<E2, D2>;
 ///
 ///     fn iter_tensors<V: ModuleVisitor<Self, E, D>>(

src/shapes/shape.rs

@@ -53,7 +53,10 @@ pub trait HasUnitType {
     type Unit: Unit;
 }
 
-/// Represents a data type or element of an array.
+/// Represents a data type or element of an array that can have
+/// arithmatic operations applied to it. The main difference
+/// between [Dtype] and [Unit] is that [`bool`] is [Unit], but
+/// not [Dtype].
 pub trait Dtype:
     Unit
     + std::ops::Add<Self, Output = Self>
@@ -65,12 +68,23 @@ pub trait Dtype:
     + std::ops::MulAssign
     + std::ops::DivAssign
     + num_traits::FromPrimitive
-    + num_traits::Float
     + rand_distr::uniform::SampleUniform
 {
 }
 impl Dtype for f32 {}
 impl Dtype for f64 {}
+impl Dtype for i8 {}
+impl Dtype for i16 {}
+impl Dtype for i32 {}
+impl Dtype for i64 {}
+impl Dtype for i128 {}
+impl Dtype for isize {}
+impl Dtype for u8 {}
+impl Dtype for u16 {}
+impl Dtype for u32 {}
+impl Dtype for u64 {}
+impl Dtype for u128 {}
+impl Dtype for usize {}
 
 /// Represents something that has a [Dtype].
 pub trait HasDtype {

src/tensor_ops/matmul/cuda_kernel.rs

@@ -184,7 +184,7 @@ where
                 [lhs.strides[0], 0],
                 rhs.data.as_ref(),
                 [0, rhs.strides[0]],
-                E::zero(),
+                Default::default(),
                 &mut storage,
                 strides,
             )
@@ -260,7 +260,7 @@ where
                 [0, lhs.strides[0]],
                 rhs.data.as_ref(),
                 rhs.strides,
-                E::zero(),
+                Default::default(),
                 &mut storage,
                 [0, strides[0]],
             )?;
@@ -333,7 +333,7 @@ where
                 lhs.strides,
                 rhs.data.as_ref(),
                 rhs.strides,
-                E::zero(),
+                Default::default(),
                 &mut storage,
                 strides,
             )
@@ -408,7 +408,7 @@ where
                 lhs.strides,
                 rhs.data.as_ref(),
                 [0, rhs.strides[0], rhs.strides[1]],
-                E::zero(),
+                Default::default(),
                 &mut storage,
                 strides,
             )?;
@@ -486,7 +486,7 @@ where
                 lhs.strides,
                 rhs.data.as_ref(),
                 rhs.strides,
-                E::zero(),
+                Default::default(),
                 &mut storage,
                 strides,
             )?;
@@ -571,7 +571,7 @@ where
                     [lhs.strides[1], lhs.strides[2], lhs.strides[3]],
                     &rhs.data.slice(b * rhs.strides[0]..),
                     [rhs.strides[1], rhs.strides[2], rhs.strides[3]],
-                    E::zero(),
+                    Default::default(),
                     &mut storage.slice_mut(b * strides[0]..),
                     [strides[1], strides[2], strides[3]],
                 )?;
@@ -585,7 +585,7 @@ where
                     [lhs.strides[1], lhs.strides[2], lhs.strides[3]],
                     &rhs.data.slice(b * rhs.strides[0]..),
                     [rhs.strides[1], rhs.strides[2], rhs.strides[3]],
-                    E::zero(),
+                    Default::default(),
                     &mut storage.slice_mut(b * strides[0]..),
                     [strides[1], strides[2], strides[3]],
                 )?;

src/tensor_ops/reshape_to/mod.rs

@@ -119,6 +119,13 @@ mod tests {
         assert!(t.reshape_like(&(7,)).is_none());
     }
 
+    #[test]
+    fn test_unit_non_contiguous_reshapes() {
+        let dev: TestDevice = Default::default();
+        let t: Tensor<Rank2<2, 3>, usize, _> = dev.zeros::<Rank0>().broadcast();
+        let _: Tensor<Rank1<6>, usize, _> = t.reshape();
+    }
+
     #[test]
     fn test_valid_reshapes() {
         let dev: TestDevice = Default::default();

src/tensor_ops/tri.rs

@@ -40,7 +40,7 @@ where
     ) -> Result<Self, <Self as HasErr>::Err> {
         let out = self
             .device
-            .try_lower_tri_like(&self.shape, E::one(), diagonal)?;
+            .try_lower_tri_like(&self.shape, E::ONE, diagonal)?;
         self.try_mul(out)
     }
 
@@ -56,7 +56,7 @@ where
     ) -> Result<Self, <Self as HasErr>::Err> {
         let out = self
             .device
-            .try_upper_tri_like(&self.shape, E::one(), diagonal)?;
+            .try_upper_tri_like(&self.shape, E::ONE, diagonal)?;
         self.try_mul(out)
     }