Add wrapped_as_base_mut impl for other modules

elftausend · Dec 16, 2023 · df3f0bd · df3f0bd
1 parent bd65c7e
commit df3f0bd
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 1 deletion.
diff --git a/Cargo.toml b/Cargo.toml
@@ -49,7 +49,7 @@ min-cl = { path="../min-cl", optional=true }
 
 [features]
 # default = ["cpu", "autograd", "macro"]
-default = ["cpu", "static-api", "lazy"]
+default = ["cpu", "static-api", "lazy", "autograd", "graph", "fork"]
 
 cpu = []
 opencl = ["dep:min-cl", "cpu", "cached"]

diff --git a/src/modules/autograd.rs b/src/modules/autograd.rs
@@ -34,6 +34,11 @@ impl<Mods: WrappedData> WrappedData for Autograd<Mods> {
     fn wrapped_as_base<'a, T, Base: HasId + PtrType>(&self, wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
         self.modules.wrapped_as_base(wrap)
     }
+
+    #[inline]
+    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(&self, wrap: &'a mut Self::Wrap<T, Base>) -> &'a mut Base {
+        self.modules.wrapped_as_base_mut(wrap)
+    }
 }
 
 impl<Mods: Module<D>, D: Device> Module<D> for Autograd<Mods> {

diff --git a/src/modules/cached.rs b/src/modules/cached.rs
@@ -32,6 +32,11 @@ impl<Mods: WrappedData, SD: Device> WrappedData for CachedModule<Mods, SD> {
     fn wrapped_as_base<'a, T, Base: HasId + PtrType>(&self, wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
         self.modules.wrapped_as_base(wrap)
     }
+
+    #[inline]
+    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(&self, wrap: &'a mut Self::Wrap<T, Base>) -> &'a mut Base {
+        self.modules.wrapped_as_base_mut(wrap)
+    }
 }
 
 impl<Mods: Module<D>, D: Device> Module<D> for Cached<Mods> {

diff --git a/src/modules/fork.rs b/src/modules/fork.rs
@@ -61,6 +61,11 @@ impl<Mods: WrappedData> WrappedData for Fork<Mods> {
     fn wrapped_as_base<'a, T, Base: HasId + PtrType>(&self, wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
         self.modules.wrapped_as_base(wrap)
     }
+
+    #[inline]
+    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(&self, wrap: &'a mut Self::Wrap<T, Base>) -> &'a mut Base {
+        self.modules.wrapped_as_base_mut(wrap)
+    }
 }
 
 impl<Mods: Module<D>, D: Device> Module<D> for Fork<Mods> {

diff --git a/src/modules/graph.rs b/src/modules/graph.rs
@@ -33,6 +33,11 @@ impl<Mods: WrappedData> WrappedData for Graph<Mods> {
     fn wrapped_as_base<'a, T, Base: HasId + PtrType>(&self, wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
         self.modules.wrapped_as_base(wrap)
     }
+
+    #[inline]
+    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(&self, wrap: &'a mut Self::Wrap<T, Base>) -> &'a mut Base {
+        self.modules.wrapped_as_base_mut(wrap)
+    }
 }
 
 impl<Mods: Module<D>, D: Device> Module<D> for Graph<Mods> {