Add Scope::with_kernel_label

src/scope.rs

@@ -102,6 +102,7 @@ pub struct Scope {
     op_names: Rc<RefCell<HashMap<String, i32>>>,
     device: String,
     control_deps: Vec<Operation>,
+    kernel_label: String,
 }
 
 impl Scope {
@@ -117,6 +118,7 @@ impl Scope {
             op_names: Rc::new(RefCell::new(HashMap::new())),
             device: "".to_string(),
             control_deps: Vec::new(),
+            kernel_label: "".to_string(),
         }
     }
 
@@ -160,6 +162,7 @@ impl Scope {
             },
             device: self.device.clone(),
             control_deps: self.control_deps.clone(),
+            kernel_label: self.kernel_label.clone(),
         }
     }
 
@@ -174,6 +177,7 @@ impl Scope {
             op_names: self.op_names.clone(),
             device: self.device.clone(),
             control_deps: self.control_deps.clone(),
+            kernel_label: self.kernel_label.clone(),
         }
     }
 
@@ -213,6 +217,7 @@ impl Scope {
             op_names: self.op_names.clone(),
             device: device.to_string(),
             control_deps: self.control_deps.clone(),
+            kernel_label: self.kernel_label.clone(),
         }
     }
 
@@ -233,6 +238,7 @@ impl Scope {
                 .chain(control_deps.iter())
                 .cloned()
                 .collect(),
+            kernel_label: self.kernel_label.clone(),
         }
     }
 
@@ -247,6 +253,22 @@ impl Scope {
             op_names: self.op_names.clone(),
             device: self.device.clone(),
             control_deps: vec![],
+            kernel_label: self.kernel_label.clone(),
+        }
+    }
+
+    /// Return a new scope. All ops created with the new scope will have
+    /// kernel_label as the value for their '_kernel' attribute.
+    pub fn with_kernel_label(&self, kernel_label: &str) -> Scope {
+        Scope {
+            graph: self.graph.clone(),
+            name: self.name.clone(),
+            children_names: self.children_names.clone(),
+            op_name: self.op_name.clone(),
+            op_names: self.op_names.clone(),
+            device: self.device.clone(),
+            control_deps: self.control_deps.clone(),
+            kernel_label: kernel_label.to_string(),
         }
     }
 
@@ -263,6 +285,9 @@ impl Scope {
         for control_dep in &self.control_deps {
             nd.add_control_input(control_dep);
         }
+        if !self.kernel_label.is_empty() {
+            nd.set_attr_string("_kernel", &self.kernel_label)?;
+        }
         f(&mut nd)?;
         Ok(nd.finish()?)
     }
@@ -333,6 +358,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn kernel_label() {
+        assert_eq!(
+            Scope::new_root_scope()
+                .with_kernel_label("foo")
+                .new_operation("NoOp", |_| Ok(()))
+                .unwrap()
+                .get_attr_string("_kernel")
+                .unwrap(),
+            "foo"
+        );
+    }
+
     #[test]
     fn control_dependencies() {
         let mut scope = Scope::new_root_scope();