coreylowman · coreylowman · Jan 17, 2023 · Jan 17, 2023
diff --git a/src/nn/npz_impls.rs b/src/nn/npz_impls.rs
@@ -333,7 +333,7 @@ impl<
 #[cfg(test)]
 mod tests {
     use crate::{
-        shapes::{Dtype, Rank1, Rank3, Shape},
+        shapes::*,
         tensor::{AsArray, SampleTensor, Tensor},
         tensor_ops::Device,
         tests::TestDevice,
@@ -344,7 +344,7 @@ mod tests {
     use tempfile::NamedTempFile;
 
     fn test_save_load<
-        S: Shape + Default,
+        S: ConstShape,
         E: Dtype,
         D: Device<E>,
         M: ResetParams<D, E> + Module<Tensor<S, E, D>> + SaveToNpz + LoadFromNpz,

diff --git a/src/shapes/mod.rs b/src/shapes/mod.rs
@@ -17,6 +17,7 @@ pub(crate) use replace_dim::{RemoveDimTo, ReplaceDimTo};
 pub(crate) use same_numel::HasSameNumelAs;
 
 pub use axes::{Axes2, Axes3, Axes4, Axes5, Axes6, Axis, HasAxes};
-pub use shape::{Const, Dim};
+pub use shape::{Const, ConstDim, Dim};
+pub use shape::{ConstShape, HasShape, Shape};
 pub use shape::{Dtype, HasDtype, HasUnitType, Unit};
-pub use shape::{HasShape, Rank0, Rank1, Rank2, Rank3, Rank4, Rank5, Rank6, Shape};
+pub use shape::{Rank0, Rank1, Rank2, Rank3, Rank4, Rank5, Rank6};
diff --git a/src/shapes/shape.rs b/src/shapes/shape.rs
@@ -43,6 +43,10 @@ pub trait Dim: 'static + Copy + Clone + std::fmt::Debug + Send + Sync + Eq + Par
     fn from_size(size: usize) -> Option<Self>;
 }
 
+/// Represents a single dimension where all
+/// instances are guarunteed to be the same size at compile time.
+pub trait ConstDim: Default + Dim {}
+
 impl Dim for usize {
     #[inline(always)]
     fn size(&self) -> usize {
@@ -72,6 +76,8 @@ impl<const M: usize> Dim for Const<M> {
     }
 }
 
+impl<const M: usize> ConstDim for Const<M> {}
+
 /// A collection of dimensions ([Dim]) that change how a multi-dimensional
 /// array is interacted with.
 pub trait Shape:
@@ -132,6 +138,9 @@ pub trait Shape:
     }
 }
 
+/// Represents a [Shape] that has all [ConstDim]s
+pub trait ConstShape: Default + Shape {}
+
 /// Represents something that has a [Shape].
 pub trait HasShape {
     type WithShape<New: Shape>: HasShape<Shape = New>;
@@ -182,6 +191,7 @@ impl<$($D: Dim, )*> Shape for ($($D, )*) {
         Some(($(Dim::from_size(concrete[$Idx])?, )*))
     }
 }
+impl<$($D: ConstDim, )*> ConstShape for ($($D, )*) { }
     };
 }
 
@@ -203,6 +213,7 @@ impl Shape for () {
         Some(())
     }
 }
+impl ConstShape for () {}
 
 shape!((D1 0), rank=1, all=Axis);
 shape!((D1 0, D2 1), rank=2, all=Axes2);

diff --git a/src/tensor/storage_traits.rs b/src/tensor/storage_traits.rs
@@ -2,7 +2,7 @@ use rand::distributions::Distribution;
 use rand_distr::{Standard, StandardNormal};
 
 use crate::{
-    shapes::{Dtype, HasDtype, HasShape, HasUnitType, Shape, Unit},
+    shapes::{ConstShape, Dtype, HasDtype, HasShape, HasUnitType, Shape, Unit},
     unique_id::unique_id,
 };
 
@@ -98,12 +98,12 @@ pub trait ZerosTensor<E: Unit>: DeviceStorage {
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>> = dev.zeros();
     /// ```
-    fn zeros<S: Shape + Default>(&self) -> Tensor<S, E, Self> {
+    fn zeros<S: ConstShape>(&self) -> Tensor<S, E, Self> {
         self.try_zeros_like::<S>(&Default::default()).unwrap()
     }
 
     /// Fallible version of [ZerosTensor::zeros]
-    fn try_zeros<S: Shape + Default>(&self) -> Result<Tensor<S, E, Self>, Self::Err> {
+    fn try_zeros<S: ConstShape>(&self) -> Result<Tensor<S, E, Self>, Self::Err> {
         self.try_zeros_like::<S>(&Default::default())
     }
 
@@ -146,12 +146,12 @@ pub trait OnesTensor<E: Unit>: DeviceStorage {
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>> = dev.ones();
     /// ```
-    fn ones<S: Shape + Default>(&self) -> Tensor<S, E, Self> {
+    fn ones<S: ConstShape>(&self) -> Tensor<S, E, Self> {
         self.try_ones_like::<S>(&Default::default()).unwrap()
     }
 
     /// Fallible version of [OnesTensor::ones]
-    fn try_ones<S: Shape + Default>(&self) -> Result<Tensor<S, E, Self>, Self::Err> {
+    fn try_ones<S: ConstShape>(&self) -> Result<Tensor<S, E, Self>, Self::Err> {
         self.try_ones_like::<S>(&Default::default())
     }
 
@@ -188,25 +188,25 @@ pub trait OneFillStorage<E: Unit>: DeviceStorage {
 
 /// Constructs tensors filled with random values from a given distribution.
 pub trait SampleTensor<E: Unit>: DeviceStorage {
-    fn sample_uniform<S: Shape + Default>(&self) -> Tensor<S, E, Self>
+    fn sample_uniform<S: ConstShape>(&self) -> Tensor<S, E, Self>
     where
         Standard: Distribution<E>,
     {
         self.sample::<S, _>(Standard)
     }
 
-    fn sample_normal<S: Shape + Default>(&self) -> Tensor<S, E, Self>
+    fn sample_normal<S: ConstShape>(&self) -> Tensor<S, E, Self>
     where
         StandardNormal: Distribution<E>,
     {
         self.sample::<S, _>(StandardNormal)
     }
 
-    fn sample<S: Shape + Default, D: Distribution<E>>(&self, distr: D) -> Tensor<S, E, Self> {
+    fn sample<S: ConstShape, D: Distribution<E>>(&self, distr: D) -> Tensor<S, E, Self> {
         self.try_sample_like::<S, D>(&Default::default(), distr)
             .unwrap()
     }
-    fn try_sample<S: Shape + Default, D: Distribution<E>>(
+    fn try_sample<S: ConstShape, D: Distribution<E>>(
         &self,
         distr: D,
     ) -> Result<Tensor<S, E, Self>, Self::Err> {