Repository: coreylowman/dfdx (language: rust)

Issue:
`reshape` does not work on tensors of type `usize`
When I attempt to run this line:
```
let upsample_gather: Tensor<Rank1<{ SIZE * 2 }>, usize, Cuda> = upsample_gather.reshape();
```
I get this error:
```
error[E0599]: the method `reshape` exists for struct `Tensor<(Const<SIZE>, Const<2>), usize, Cuda>`, but its trait bounds were not satisfied
  --> src\upsample.rs:28:83
   |
28 |         let upsample_gather: Tensor<Rank1<{ SIZE * 2 }>, usize, Cuda> = upsample_gather.reshape();
   |                                                                                         ^^^^^^^ method cannot be called on `Tensor<(Const<SIZE>, Const<2>), usize, Cuda>` due to unsatisfied trait bounds
   |
  ::: C:\Users\heruf\.cargo\registry\src\index.crates.io-6f17d22bba15001f\dfdx-0.11.0\src\tensor\tensor_impls.rs:32:1
   |
32 | pub struct Tensor<S: Shape, E: Unit, D: DeviceStorage, T = NoneTape> {
   | -------------------------------------------------------------------- doesn't satisfy `_: ReshapeTo`
   |
  ::: C:\Users\heruf\.cargo\registry\src\index.crates.io-6f17d22bba15001f\dfdx-0.11.0\src\tensor\cuda\device.rs:14:1
   |
14 | pub struct Cuda {
   | --------------- doesn't satisfy `_: ReshapeKernel<usize>`
   |
   = note: the following trait bounds were not satisfied:
           `usize: Dtype`
           which is required by `dfdx::tensor::Tensor<(Const<SIZE>, Const<2>), usize, dfdx::tensor::Cuda>: dfdx::tensor_ops::ReshapeTo`
           `dfdx::tensor::Cuda: tensor_ops::reshape_to::ReshapeKernel<usize>`
           which is required by `dfdx::tensor::Tensor<(Const<SIZE>, Const<2>), usize, dfdx::tensor::Cuda>: dfdx::tensor_ops::ReshapeTo`
```
Before that line, `upsample_gather` has type `Tensor<(Const<SIZE>, Const<2>), usize, Cuda>`.

It appears `Dtype` is only implemented for `f32` and `f64`. I'm not sure what the exact intent of `Dtype` is. If it's supposed to cover all the data types that can go into a tensor, then it should be implemented for all the other types that tensors now support. If it's intended to tag specifically floating point types, then `reshape` should be updated to not require `Dtype`, because `reshape` doesn't do any arithmetic on the contents of the tensor, so shouldn't need to care what data type it stores.


Discussion / hints:
Dtype is intended to be any type that can do a backward pass with gradients (which is for now floating point numbers), and Unit is intended to be any type that a tensor can hold.

I see the need though, will think about this
Notably, during inference we can do forward pass without requirement of doing backward pass. Additionally for quantization, we need to support some operations with integer dtypes
For my immediate use case, I just needed to flatten a 2D tensor into a 1D tensor, so I was able to work around it by doing `device.tensor(upsample_gather.as_vec())`. But that's very specific to this case. And would also break back propagation, if that mattered, but since this is just the tensor for defining a gather operation, it shouldn't need back prop.
Resolved by #602.
I don't think it's actually resolved - the reshape traits & methods on tensor still require Dtype don't they?
Whoops, forgot to change that, this can be resolved by changing `Dtype` to `Unit` on lines 8 and 76 of `reshape_to/mod.rs`
This is a larger issue with all the data movement kernels - I could see people wanting to use broadcast/permute/etc on non float tensors when gradients aren't being tracked.

From the kernel trait level we could just make the top level `Unit` and then add `where E: Dtype` for backward. Then for the tensor methods/traits we'd have one impl for Dtype that calls fwd & bwd, and one impl for Unit that just calls fwd

Fix patch:
diff --git a/examples/07-custom-module.rs b/examples/07-custom-module.rs
index 717c3e222..9a3500977 100644
--- a/examples/07-custom-module.rs
+++ b/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.
diff --git a/src/nn/linear.rs b/src/nn/linear.rs
index f666bd590..a4f071f34 100644
--- a/src/nn/linear.rs
+++ b/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>;
 
diff --git a/src/nn/tensor_collection/collection.rs b/src/nn/tensor_collection/collection.rs
index e587af2e9..3277f565d 100644
--- a/src/nn/tensor_collection/collection.rs
+++ b/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>>(
diff --git a/src/shapes/shape.rs b/src/shapes/shape.rs
index 55662a83c..df77aeeca 100644
--- a/src/shapes/shape.rs
+++ b/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 {
diff --git a/src/tensor_ops/matmul/cuda_kernel.rs b/src/tensor_ops/matmul/cuda_kernel.rs
index a513d95fd..253822666 100644
--- a/src/tensor_ops/matmul/cuda_kernel.rs
+++ b/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]],
                 )?;
diff --git a/src/tensor_ops/reshape_to/mod.rs b/src/tensor_ops/reshape_to/mod.rs
index 606878257..63038a670 100644
--- a/src/tensor_ops/reshape_to/mod.rs
+++ b/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();
diff --git a/src/tensor_ops/tri.rs b/src/tensor_ops/tri.rs
index ad542ffbb..0e7660316 100644
--- a/src/tensor_ops/tri.rs
+++ b/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)
     }
 


Test patch: