util functions

src/core/graph.rs

@@ -16,7 +16,7 @@ use colored::Colorize;
 use itertools::Itertools;
 use petgraph::{stable_graph::StableGraph, visit::EdgeRef, Direction};
 
-use super::compiler_utils::ToIdsMut;
+use super::compiler_utils::{ToIds, ToIdsMut};
 
 pub type MainGraph = StableGraph<Box<dyn Operator>, Dependency>;
 pub use petgraph::stable_graph::NodeIndex;
@@ -82,6 +82,25 @@ impl Graph {
         self.tensors.get(&(id, ind))
     }
 
+    pub fn drop_tensors<T: ToIds>(&mut self, tensors: T) {
+        for id in tensors.to_ids() {
+            self.tensors.remove(&(id, 0));
+        }
+    }
+
+    pub fn keep_tensors<T: ToIds>(&mut self, tensors: T) {
+        for id in tensors.to_ids() {
+            self.no_delete.insert(id);
+        }
+    }
+
+    pub fn retrieve_tensors<T: ToIds>(&mut self, tensors: T) {
+        for id in tensors.to_ids() {
+            self.no_delete.insert(id);
+            self.to_retrieve.insert(id);
+        }
+    }
+
     pub fn set_tensor(&mut self, id: NodeIndex, ind: u8, tensor: Tensor) {
         self.tensors.insert((id, ind), tensor);
     }

src/core/graph_tensor.rs

@@ -34,20 +34,20 @@ impl<S: Shape> GraphTensor<S> {
 
     /// Mark this tensor to not be deleted
     pub fn keep(self) -> Self {
-        self.graph().no_delete.insert(self.id);
+        self.graph().keep_tensors(self.id);
         self
     }
 
     /// Mark this tensor to be retrieved later
     pub fn retrieve(self) -> Self {
         self.keep();
-        self.graph().to_retrieve.insert(self.id);
+        self.graph().retrieve_tensors(self.id);
         self
     }
 
     /// Remove this tensor's data from the graph.
     pub fn drop(&self) {
-        self.graph().tensors.remove(&(self.id, 0));
+        self.graph().drop_tensors(self.id);
     }
 
     #[allow(clippy::mut_from_ref)]