Add Operation Instances for accessing Outputs during graph construction

tensorflow · May 13, 2022 · a03c7e0 · a03c7e0
1 parent f9f2ed6
commit a03c7e0
Show file tree

Hide file tree

Showing 3 changed files with 95,230 additions and 3,873 deletions.
diff --git a/src/graph.rs b/src/graph.rs
@@ -1037,6 +1037,15 @@ impl Operation {
         }
     }
 
+    /// Returns the given output edge.
+    /// The index argument is the index into the current operation's output array,
+    pub fn output(&self, index: usize) -> Output {
+        crate::Output {
+            operation: self.clone(),
+            index: index as c_int,
+        }
+    }
+
     // TODO: Figure out what this does and document it.
     #[allow(missing_docs)]
     pub fn output_list_length(&self, arg_name: &str) -> Result<usize> {