coreylowman · coreylowman · Oct 19, 2022 · Oct 18, 2022 · Oct 18, 2022
diff --git a/src/gradients.rs b/src/gradients.rs
@@ -3,7 +3,7 @@
 use std::collections::HashMap;
 use std::{boxed::Box, vec::Vec};
 
-use crate::arrays::HasArrayType;
+use crate::arrays::{HasArrayData, HasArrayType};
 use crate::devices::{AllocateZeros, HasDevice};
 use crate::unique_id::{HasUniqueId, UniqueId};
 
@@ -264,7 +264,7 @@ pub trait GradientProvider {
     /// based on the associated data!
     fn gradient<P>(&mut self, p: &P) -> Option<Box<P::Array>>
     where
-        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice;
+        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice + HasArrayData;
 }
 
 /// Represents something that can be updated with [GradientProvider].

diff --git a/src/nn/mod.rs b/src/nn/mod.rs
@@ -126,6 +126,7 @@ mod npz_impls;
 
 #[cfg(test)]
 mod tests {
+    use crate::arrays::{HasArrayData, HasArrayType};
     use crate::gradients::{GradientProvider, Gradients};
     use crate::unique_id::HasUniqueId;
     use std::boxed::Box;
@@ -136,7 +137,7 @@ mod tests {
     impl GradientProvider for SimpleGradients {
         fn gradient<P>(&mut self, p: &P) -> Option<Box<P::Array>>
         where
-            P: HasUniqueId + crate::arrays::HasArrayType<Dtype = f32> + crate::devices::HasDevice,
+            P: HasUniqueId + HasArrayType<Dtype = f32> + crate::devices::HasDevice + HasArrayData,
         {
             self.0.remove(p)
         }

diff --git a/src/optim/adam.rs b/src/optim/adam.rs
@@ -99,7 +99,7 @@ impl<M> Adam<M> {
 impl<M> GradientProvider for Adam<M> {
     fn gradient<P>(&mut self, p: &P) -> Option<Box<P::Array>>
     where
-        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice,
+        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice + HasArrayData,
     {
         let mut g_t = self.gradients.remove(p)?;
         let m_t = self.moment1.mut_gradient(p);

diff --git a/src/optim/rmsprop.rs b/src/optim/rmsprop.rs
@@ -110,7 +110,7 @@ impl<M> RMSprop<M> {
 impl<M> GradientProvider for RMSprop<M> {
     fn gradient<P>(&mut self, p: &P) -> Option<Box<P::Array>>
     where
-        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice,
+        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice + HasArrayData,
     {
         let mut g_t = self.gradients.remove(p)?;
 

diff --git a/src/optim/sgd.rs b/src/optim/sgd.rs
@@ -119,7 +119,7 @@ impl<M> Sgd<M> {
 impl<M> GradientProvider for Sgd<M> {
     fn gradient<P>(&mut self, p: &P) -> Option<Box<P::Array>>
     where
-        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice,
+        P: HasUniqueId + HasArrayType<Dtype = f32> + HasDevice + HasArrayData,
     {
         let mut g_t = self.gradients.remove(p)?;
         match self.cfg.momentum {