Add Graph::import_graph_def_with_return_outputs

Author: adamcrume

src/expr.rs

@@ -94,7 +94,7 @@ pub trait ExprImpl<T: TensorType>: Display + Debug {
     /// rather than creating child operations itself.
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status>;
 
@@ -113,7 +113,7 @@ impl<T: TensorType> ExprImpl<T> for T {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        _children: &[Rc<Operation>],
+                        _children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("Const", &id_gen()));
@@ -197,11 +197,11 @@ macro_rules! impl_bin_op {
         vec![Box::new(self.left.clone()), Box::new(self.right.clone())]
       }
 
-      fn create_operation(&self, graph: &mut Graph, children: &[Rc<Operation>],
+      fn create_operation(&self, graph: &mut Graph, children: &[Operation],
           id_gen: &mut FnMut() -> String) -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation($tf_op, &id_gen()));
-        nd.add_input(Output {operation: &children[0], index: 0});
-        nd.add_input(Output {operation: &children[1], index: 0});
+        nd.add_input(Output {operation: children[0].clone(), index: 0});
+        nd.add_input(Output {operation: children[1].clone(), index: 0});
         nd.finish()
       }
 
@@ -287,18 +287,18 @@ impl<T: TensorType> ExprImpl<T> for TruncateDiv<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("TruncateDiv", &id_gen()));
         nd.add_input(Output {
-            operation: &children[0],
-            index: 0,
-        });
+                         operation: children[0].clone(),
+                         index: 0,
+                     });
         nd.add_input(Output {
-            operation: &children[1],
-            index: 0,
-        });
+                         operation: children[1].clone(),
+                         index: 0,
+                     });
         nd.finish()
     }
 
@@ -351,14 +351,14 @@ impl<T: TensorType> ExprImpl<T> for Neg<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("Neg", &id_gen()));
         nd.add_input(Output {
-            operation: &children[0],
-            index: 0,
-        });
+                         operation: children[0].clone(),
+                         index: 0,
+                     });
         nd.finish()
     }
 
@@ -409,7 +409,7 @@ impl<T: TensorType> ExprImpl<T> for Variable<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        _children: &[Rc<Operation>],
+                        _children: &[Operation],
                         _id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("Variable", &self.name));
@@ -469,7 +469,7 @@ impl<T: TensorType> ExprImpl<T> for Placeholder<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        _children: &[Rc<Operation>],
+                        _children: &[Operation],
                         _id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("Placeholder", &self.name));
@@ -523,18 +523,18 @@ impl<T: TensorType> ExprImpl<T> for Assign<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         let mut nd = try!(graph.new_operation("Assign", &id_gen()));
         nd.add_input(Output {
-            operation: &children[0],
-            index: 0,
-        });
+                         operation: children[0].clone(),
+                         index: 0,
+                     });
         nd.add_input(Output {
-            operation: &children[1],
-            index: 0,
-        });
+                         operation: children[1].clone(),
+                         index: 0,
+                     });
         nd.finish()
     }
 
@@ -563,7 +563,7 @@ pub trait AnyExpr: Debug {
     /// rather than creating child operations itself.
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status>;
 
@@ -586,7 +586,7 @@ impl<T: TensorType> AnyExpr for Expr<T> {
 
     fn create_operation(&self,
                         graph: &mut Graph,
-                        children: &[Rc<Operation>],
+                        children: &[Operation],
                         id_gen: &mut FnMut() -> String)
                         -> Result<Operation, Status> {
         self.expr.create_operation(graph, children, id_gen)
@@ -620,7 +620,7 @@ impl Hash for Key {
 #[derive(Debug)]
 pub struct Compiler<'l> {
     graph: &'l mut Graph,
-    operations: HashMap<Key, Rc<Operation>>,
+    operations: HashMap<Key, Operation>,
     next_id: i32,
 }
 
@@ -635,12 +635,12 @@ impl<'l> Compiler<'l> {
     }
 
     /// Compiles the expression.
-    pub fn compile<T: TensorType>(&mut self, expr: Expr<T>) -> Result<Rc<Operation>, Status> {
+    pub fn compile<T: TensorType>(&mut self, expr: Expr<T>) -> Result<Operation, Status> {
         self.compile_any(Box::new(expr))
     }
 
     /// Compiles the expression.
-    pub fn compile_any(&mut self, expr: Box<AnyExpr>) -> Result<Rc<Operation>, Status> {
+    pub fn compile_any(&mut self, expr: Box<AnyExpr>) -> Result<Operation, Status> {
         let mut child_operations = vec![];
         for child in expr.children() {
             let key = Key(child.clone_box());
@@ -661,7 +661,7 @@ impl<'l> Compiler<'l> {
                                                id
                                            });
         self.next_id = next_id;
-        let operation = Rc::new(try!(result));
+        let operation = result?;
         self.operations.insert(Key(expr), operation.clone());
         Ok(operation)
     }

src/graph.rs

@@ -248,7 +248,7 @@ impl Graph {
     ///   * `output` is not in `graph`.
     pub fn tensor_shape(&self, output: Output) -> Result<Shape> {
         let mut status = Status::new();
-        let n = try!(self.num_dims(output));
+        let n = try!(self.num_dims(output.clone()));
         if n == -1 {
             return Ok(Shape(None));
         }
@@ -283,6 +283,36 @@ impl Graph {
             status.into_result()
         }
     }
+
+    /// Import the graph serialized in `graph_def`.
+    ///
+    /// `num_return_outputs` must be the number of return outputs added (i.e. the
+    /// result of TF_ImportGraphDefOptionsNumReturnOutputs()).  If
+    /// `num_return_outputs` is non-zero, `return_outputs` must be of length
+    /// `num_return_outputs`. Otherwise it can be null.
+    pub fn import_graph_def_with_return_outputs(&mut self,
+                                                graph_def: &[u8],
+                                                options: &ImportGraphDefOptions)
+                                                -> Result<Vec<Output>> {
+        let buf = Buffer::from(graph_def);
+        let mut status = Status::new();
+        let mut c_return_outputs = Vec::new();
+        let n = options.num_return_outputs();
+        unsafe {
+            c_return_outputs.set_len(n);
+            tf::TF_GraphImportGraphDefWithReturnOutputs(self.gimpl.inner,
+                                                        buf.inner(),
+                                                        options.inner,
+                                                        c_return_outputs.as_mut_ptr(),
+                                                        n as c_int,
+                                                        status.inner());
+        }
+        status.into_result()?;
+        Ok(c_return_outputs
+               .iter()
+               .map(|x| Output::from_c(self, x))
+               .collect())
+    }
 }
 
 impl GraphTrait for Graph {
@@ -324,7 +354,7 @@ impl<'a> Iterator for OperationIter<'a> {
 
 /// An `Operation` is a node in a `Graph`.
 /// It is a computation which accepts inputs and produces outputs.
-#[derive(Debug)]
+#[derive(Debug,Clone)]
 pub struct Operation {
     inner: *mut tf::TF_Operation,
     gimpl: Arc<GraphImpl>,
@@ -556,22 +586,32 @@ impl<'a> Input<'a> {
 
 /// A `Output` is one end of a graph edge.
 /// It holds an operation and an index into the outputs of that operation.
-#[derive(Debug,Copy,Clone)]
-pub struct Output<'a> {
+#[derive(Debug,Clone)]
+pub struct Output {
     /// Operation the edge connects to.
-    pub operation: &'a Operation,
+    pub operation: Operation,
 
     /// Index into either the outputs of the operation.
     pub index: c_int,
 }
 
-impl<'a> Output<'a> {
+impl Output {
     fn to_c(&self) -> tf::TF_Output {
         tf::TF_Output {
             oper: self.operation.inner,
             index: self.index,
         }
     }
+
+    fn from_c(graph: &Graph, output: &tf::TF_Output) -> Self {
+        Output {
+            operation: Operation {
+                inner: output.oper,
+                gimpl: graph.gimpl.clone(),
+            },
+            index: output.index,
+        }
+    }
 }
 
 ////////////////////////

src/session.rs

@@ -301,11 +301,11 @@ mod tests {
         let y = {
             let mut nd = g.new_operation("Mul", "y").unwrap();
             nd.add_input(Output {
-                             operation: &two,
+                             operation: two,
                              index: 0,
                          });
             nd.add_input(Output {
-                             operation: &x,
+                             operation: x.clone(),
                              index: 0,
                          });
             nd.finish().unwrap()