NativeScope class

Signed-off-by: Ryan Nett <rnett@calpoly.edu>

Author: rnett

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/GraphOperation.java

@@ -331,7 +331,7 @@ public Set<GraphOperation> controlConsumers() {
   }
 
 
-  TF_Operation getUnsafeNativeHandle() {
+  public TF_Operation getUnsafeNativeHandle() {
     return unsafeNativeHandle;
   }
 

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

@@ -0,0 +1,40 @@
+/*
+  Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ ==============================================================================
+ */
+package org.tensorflow.op;
+
+import org.tensorflow.DeviceSpec;
+import org.tensorflow.ExecutionEnvironment;
+import org.tensorflow.OperationBuilder;
+
+public interface IScope {
+
+    ExecutionEnvironment env();
+
+    IScope withSubScope(String childScopeName);
+
+    IScope withName(String opName);
+
+    IScope withNameAsSubScope(String defaultName);
+
+    IScope withDevice(DeviceSpec deviceSpec);
+
+    String makeOpName(String defaultName);
+
+    IScope withControlDependencies(Iterable<Op> controls);
+
+    OperationBuilder apply(OperationBuilder builder);
+}

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

@@ -0,0 +1,94 @@
+/*
+  Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ ==============================================================================
+ */
+package org.tensorflow.op;
+
+import org.bytedeco.javacpp.PointerPointer;
+import org.tensorflow.*;
+import org.tensorflow.internal.c_api.TF_Operation;
+import org.tensorflow.internal.c_api.TF_Scope;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+public final class NativeScope implements IScope {
+
+    @Override
+    public ExecutionEnvironment env() {
+        return graph;
+    }
+
+    @Override
+    public NativeScope withSubScope(String childScopeName) {
+        return new NativeScope(nativeScope.NewSubScope(childScopeName), graph);
+    }
+
+    @Override
+    public NativeScope withName(String opName) {
+        return new NativeScope(nativeScope, graph, opName);
+    }
+
+    @Override
+    public NativeScope withNameAsSubScope(String defaultName) {
+        return withSubScope(opName);
+    }
+
+    @Override
+    public NativeScope withDevice(DeviceSpec deviceSpec) {
+        return new NativeScope(nativeScope.WithDevice(deviceSpec.toString()), graph);
+    }
+
+    @Override
+    public String makeOpName(String defaultName) {
+        String name = opName != null ? opName : defaultName;
+        return nativeScope.GetUniqueNameForOp(name);
+    }
+
+    @Override
+    public NativeScope withControlDependencies(Iterable<Op> controls) {
+        List<Op> controlDeps = StreamSupport.stream(controls.spliterator(), false).collect(Collectors.toList());
+        PointerPointer<TF_Operation> ops = new PointerPointer<TF_Operation>(controlDeps.size());
+
+        for(int i = 0 ; i < controlDeps.size() ; i++){
+            Operation op = controlDeps.get(i).op();
+            if(!(op instanceof GraphOperation))
+                throw new IllegalArgumentException("Can only add graph ops as control dependencies");
+            ops.put(i, (((GraphOperation) op).getUnsafeNativeHandle()));
+        }
+
+        return new NativeScope(nativeScope.WithControlDependencies(new TF_Operation(ops)), graph);
+    }
+
+    @Override
+    public OperationBuilder apply(OperationBuilder builder) {
+        return builder;
+    }
+
+    NativeScope(TF_Scope nativeScope, Graph graph){
+        this(nativeScope, graph, null);
+    }
+
+    private NativeScope(TF_Scope nativeScope, Graph graph, String opName){
+        this.graph = graph;
+        this.nativeScope = nativeScope;
+        this.opName = opName;
+    }
+
+    private final Graph graph;
+    private final TF_Scope nativeScope;
+    private final String opName;
+}

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

@@ -75,7 +75,7 @@
  *
  * <p>Scope objects are <b>not</b> thread-safe.
  */
-public final class Scope {
+public final class Scope implements IScope {
 
   /**
    * Create a new top-level scope.
@@ -89,6 +89,7 @@ public Scope(ExecutionEnvironment env) {
   /**
    * Returns the execution environment used by this scope.
    */
+  @Override
   public ExecutionEnvironment env() {
     return env;
   }
@@ -105,6 +106,7 @@ public ExecutionEnvironment env() {
    * @return a new subscope
    * @throws IllegalArgumentException if the name is invalid
    */
+  @Override
   public Scope withSubScope(String childScopeName) {
     return new Scope(env, nameScope.withSubScope(childScopeName), controlDependencies, deviceSpec);
   }
@@ -121,6 +123,7 @@ public Scope withSubScope(String childScopeName) {
    * @return a new Scope that uses opName for operations.
    * @throws IllegalArgumentException if the name is invalid
    */
+  @Override
   public Scope withName(String opName) {
     return new Scope(env, nameScope.withName(opName), controlDependencies, deviceSpec);
   }
@@ -140,6 +143,7 @@ public Scope withName(String opName) {
    * @return a new subscope
    * @throws IllegalArgumentException if the name is invalid
    */
+  @Override
   public Scope withNameAsSubScope(String defaultName){
     return new Scope(env, nameScope.withSubScope(nameScope.makeOpName(defaultName)), controlDependencies, deviceSpec);
   }
@@ -153,6 +157,7 @@ public Scope withNameAsSubScope(String defaultName){
    * @param deviceSpec device specification for an operator in the returned scope
    * @return a new Scope that uses opName for operations.
    */
+  @Override
   public Scope withDevice(DeviceSpec deviceSpec) {
     return new Scope(env, nameScope, controlDependencies, deviceSpec);
   }
@@ -177,6 +182,7 @@ public Scope withDevice(DeviceSpec deviceSpec) {
    * @return unique name for the operator.
    * @throws IllegalArgumentException if the default name is invalid.
    */
+  @Override
   public String makeOpName(String defaultName) {
     return nameScope.makeOpName(defaultName);
   }
@@ -198,6 +204,7 @@ private Scope(
    * @param controls control dependencies for ops created with the returned scope
    * @return a new scope with the provided control dependencies
    */
+  @Override
   public Scope withControlDependencies(Iterable<Op> controls) {
     for (Op control : controls) {
       env.checkInput(control);
@@ -211,6 +218,7 @@ public Scope withControlDependencies(Iterable<Op> controls) {
    *
    * @param builder OperationBuilder to add control inputs and device specification to
    */
+  @Override
   public OperationBuilder apply(OperationBuilder builder) {
     builder.setDevice(deviceSpec.toString());
     return applyControlDependencies(builder);