Have scope detect used ops from environment

Signed-off-by: Ryan Nett <[email protected]>

Author: rnett

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/NameScope.java

@@ -17,14 +17,17 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import org.tensorflow.ExecutionEnvironment;
+import org.tensorflow.Graph;
 
 /**
  * A class to manage scoped (hierarchical) names for operators.
  *
  * <p>{@code NameScope} manages hierarchical names where each component in the hierarchy is
- * separated by a forward slash {@code '/'}. For instance, {@code nn/Const_72} or {@code
- * nn/gradient/assign/init}. Each scope is a subtree in this hierarchy.
+ * separated by a forward slash {@code '/'}. For instance, {@code nn/Const_72} or {@code nn/gradient/assign/init}. Each
+ * scope is a subtree in this hierarchy.
  *
  * <p>Use {@code NameScope} to group related operations within a hierarchy, which for example lets
  * tensorboard coalesce nodes for better graph visualizations.
@@ -50,6 +53,45 @@ NameScope withName(String name) {
     return new NameScope(opPrefix, name, ids);
   }
 
+  private static final Pattern NAME_PATTERN = Pattern.compile("(.+)_(\\d+)", Pattern.DOTALL);
+
+  /**
+   * "Import" used names from a graph.  Useful when adding to a loaded graph.
+   */
+  NameScope withUsedFrom(ExecutionEnvironment env) {
+
+    if (env instanceof Graph) {
+      ((Graph) env).operations().forEachRemaining(op -> {
+        if (op.name().startsWith(opPrefix != null ? opPrefix : "")) {
+          String name = op.name();
+
+          if (opPrefix != null) {
+            name = name.substring(opPrefix.length() + 1);
+          }
+
+          if (!name.contains("/")) {
+            Matcher matcher = NAME_PATTERN.matcher(name);
+            if (matcher.find()) {
+              String realName = matcher.group(1);
+              int num = Integer.parseInt(matcher.group(2)) + 1;
+
+              if (!(ids.containsKey(realName) && ids.get(realName) > num)) {
+                ids.put(realName, num);
+              }
+            } else {
+              if (!ids.containsKey(name)) {
+                ids.put(name, 1);
+              } else {
+                ids.put(name, ids.get(name) + 1);
+              }
+            }
+          }
+        }
+      });
+    }
+    return this;
+  }
+
   String makeOpName(String name) {
     checkPattern(NAME_REGEX, name);
     // Override with opName if it exists.
@@ -120,15 +162,22 @@ private String fullyQualify(String name) {
   // instance mapped to the next available numeric suffix for it.
   private final Map<String, Integer> ids;
 
+  static boolean isValidName(String name) {
+    if (name == null) {
+      return false;
+    }
+    return NAME_REGEX.matcher(name).matches();
+  }
+
   private static void checkPattern(Pattern pattern, String name) {
     if (name == null) {
       throw new IllegalArgumentException("Names cannot be null");
     }
     if (!pattern.matcher(name).matches()) {
       throw new IllegalArgumentException(
-          String.format(
-              "invalid name: '%s' does not match the regular expression %s",
-              name, NAME_REGEX.pattern()));
+              String.format(
+                      "invalid name: '%s' does not match the regular expression %s",
+                      name, NAME_REGEX.pattern()));
     }
   }
 

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/Scope.java

@@ -83,7 +83,7 @@ public final class Scope {
    * @param env The execution environment used by the scope.
    */
   public Scope(ExecutionEnvironment env) {
-    this(env, new NameScope(), new ArrayList<>(), DeviceSpec.newBuilder().build());
+    this(env, new NameScope().withUsedFrom(env), new ArrayList<>(), DeviceSpec.newBuilder().build());
   }
 
   /**
@@ -106,7 +106,7 @@ public ExecutionEnvironment env() {
    * @throws IllegalArgumentException if the name is invalid
    */
   public Scope withSubScope(String childScopeName) {
-    return new Scope(env, nameScope.withSubScope(childScopeName), controlDependencies, deviceSpec);
+    return new Scope(env, nameScope.withSubScope(childScopeName).withUsedFrom(env), controlDependencies, deviceSpec);
   }
 
   /**
@@ -141,7 +141,8 @@ public Scope withName(String opName) {
    * @throws IllegalArgumentException if the name is invalid
    */
   public Scope withNameAsSubScope(String defaultName){
-    return new Scope(env, nameScope.withSubScope(nameScope.makeOpName(defaultName)), controlDependencies, deviceSpec);
+    return new Scope(env, nameScope.withSubScope(nameScope.makeOpName(defaultName)).withUsedFrom(env),
+            controlDependencies, deviceSpec);
   }
 
   /**
@@ -181,8 +182,12 @@ public String makeOpName(String defaultName) {
     return nameScope.makeOpName(defaultName);
   }
 
+  public static boolean isValidOpName(String name) {
+    return NameScope.isValidName(name);
+  }
+
   private Scope(
-      ExecutionEnvironment env, NameScope nameScope, Iterable<Op> controlDependencies, DeviceSpec deviceSpec) {
+          ExecutionEnvironment env, NameScope nameScope, Iterable<Op> controlDependencies, DeviceSpec deviceSpec) {
     this.env = env;
     this.nameScope = nameScope;
     this.controlDependencies = controlDependencies;