Store and allow getting native scope device when it has been set from Java

Signed-off-by: Ryan Nett <JNett96@gmail.com>

Author: rnett

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

@@ -38,12 +38,12 @@ public ExecutionEnvironment env() {
 
   @Override
   public GradientScope withSubScope(String childScopeName) {
-    return new GradientScope(nativeScope.NewSubScope(childScopeName), graph, childScopeName);
+    return new GradientScope(nativeScope.NewSubScope(childScopeName), graph, childScopeName, device);
   }
 
   @Override
   public GradientScope withName(String opName) {
-    return new GradientScope(nativeScope, graph, opName);
+    return new GradientScope(nativeScope, graph, opName, device);
   }
 
   @Override
@@ -53,7 +53,7 @@ public GradientScope withNameAsSubScope(String defaultName) {
 
   @Override
   public GradientScope withDevice(DeviceSpec deviceSpec) {
-    return new GradientScope(nativeScope.WithDevice(deviceSpec.toString()), graph);
+    return new GradientScope(nativeScope.WithDevice(deviceSpec.toString()), graph, deviceSpec.toString());
   }
 
   @Override
@@ -90,7 +90,7 @@ public GradientScope withControlDependencies(Iterable<Op> controls) {
           .put(new NativeOperation(((GraphOperation) op).getUnsafeNativeHandle().node()));
     }
 
-    return new GradientScope(nativeScope.WithControlDependencies(new NativeOperation(ops)), graph);
+    return new GradientScope(nativeScope.WithControlDependencies(new NativeOperation(ops)), graph, device);
   }
 
   @Override
@@ -108,7 +108,7 @@ public Scope withControlDependencyOps(Iterable<Operation> controls) {
           .put(new NativeOperation(((GraphOperation) op).getUnsafeNativeHandle().node()));
     }
 
-    return new GradientScope(nativeScope.WithControlDependencies(new NativeOperation(ops)), graph);
+    return new GradientScope(nativeScope.WithControlDependencies(new NativeOperation(ops)), graph, device);
   }
 
   @Override
@@ -121,25 +121,31 @@ public void onOpCreated(Operation op) {}
 
   @Override
   public String getDeviceString() {
-    throw new IllegalStateException("Can't get device string for gradient scope");
+    if (device == null) {
+      throw new UnsupportedOperationException("Can't get device string for gradient scope unless it has been explicitly set");
+    } else {
+      return device;
+    }
   }
 
   @Override
   public boolean isInit() {
     return false;
   }
 
-  GradientScope(TF_Scope nativeScope, Graph graph) {
-    this(nativeScope, graph, null);
+  GradientScope(TF_Scope nativeScope, Graph graph, String device) {
+    this(nativeScope, graph, null, device);
   }
 
-  private GradientScope(TF_Scope nativeScope, Graph graph, String opName) {
+  private GradientScope(TF_Scope nativeScope, Graph graph, String opName, String device) {
     this.graph = graph;
     this.nativeScope = nativeScope;
     this.opName = opName;
+    this.device = device;
   }
 
   private final Graph graph;
   private final TF_Scope nativeScope;
   private final String opName;
+  private final String device;
 }

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

@@ -54,7 +54,7 @@ public NativeStatus call(
         throw new IllegalStateException("No graph found for native gradient scope.");
       }
 
-      Scope nativeScope = new GradientScope(scope, g);
+      Scope nativeScope = new GradientScope(scope, g, null);
       Ops tf = new Ops(nativeScope);
 
       List<Output<?>> gradInputs = BaseGradientAdapter.fromNativeOutputs(g, grad_inputs);

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

@@ -60,7 +60,7 @@ public NativeStatus call(
         throw new IllegalStateException("No graph found for native gradient scope.");
       }
 
-      Scope nativeScope = new GradientScope(scope, g);
+      Scope nativeScope = new GradientScope(scope, g, null);
 
       Ops tf = new Ops(nativeScope);