Update framework

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

Author: rnett

tensorflow-framework/src/main/java/org/tensorflow/framework/data/Dataset.java

@@ -29,7 +29,6 @@
 import org.tensorflow.framework.data.impl.TensorSliceDataset;
 import org.tensorflow.framework.data.impl.TextLineDataset;
 import org.tensorflow.ndarray.Shape;
-import org.tensorflow.op.Op;
 import org.tensorflow.op.Ops;
 import org.tensorflow.types.family.TType;
 
@@ -254,7 +253,7 @@ public DatasetIterator makeInitializeableIterator() {
    * <pre>
    *     try (Session session = new Session(graph) {
    *         // Immediately run initializers
-   *         session.run(tf.init());
+   *         session.initialize();
    *     }
    * </pre>
    *
@@ -264,8 +263,8 @@ public DatasetIterator makeInitializeableIterator() {
    */
   public DatasetIterator makeOneShotIterator() {
     DatasetIterator iterator = makeInitializeableIterator();
-    Op initializer = iterator.makeInitializer(this);
-    if (tf.scope().env().isGraph()) tf.initAdd(initializer);
+    // TODO should pass the scope instead
+    tf.scope().env().registerInitOp(iterator.makeInitializer(this).op());
     return iterator;
   }
 

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/AdaDelta.java

@@ -15,6 +15,7 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -23,8 +24,6 @@
 import org.tensorflow.op.train.ApplyAdadelta;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-
 /**
  * Optimizer that implements the Adadelta algorithm.
  *

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/AdaGrad.java

@@ -15,16 +15,15 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
 import org.tensorflow.op.Op;
-import org.tensorflow.op.train.ApplyAdagrad;
 import org.tensorflow.op.core.Variable;
+import org.tensorflow.op.train.ApplyAdagrad;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-
 /**
  * Optimizer that implements the Adagrad algorithm.
  *
@@ -43,8 +42,8 @@ public class AdaGrad extends Optimizer {
   public static final float LEARNING_RATE_DEFAULT = 0.001f;
   public static final float INITIAL_ACCUMULATOR_DEFAULT = 0.01f;
 
-  private static final ApplyAdagrad.Options[] opts = new ApplyAdagrad.Options[]{
-          ApplyAdagrad.updateSlots(true), ApplyAdagrad.useLocking(true)};
+  private static final ApplyAdagrad.Options[] opts =
+      new ApplyAdagrad.Options[] {ApplyAdagrad.updateSlots(true), ApplyAdagrad.useLocking(true)};
 
   private final float learningRate;
 
@@ -135,7 +134,8 @@ protected void createSlots(List<Output<? extends TType>> variables) {
    */
   private <T extends TType> void createAdaGradSlot(Output<T> v) {
     Operand<T> initializer =
-        tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
+        tf.fill(
+            tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
     createSlot(v.asOutput(), ACCUMULATOR, initializer);
   }
 

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/AdaGradDA.java

@@ -15,20 +15,18 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
+import java.util.Optional;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.op.Op;
-import org.tensorflow.op.core.Assign;
 import org.tensorflow.op.core.Variable;
 import org.tensorflow.op.train.ApplyAdagradDa;
 import org.tensorflow.types.TInt64;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-import java.util.Optional;
-
 /**
  * Optimizer that implements the Adagrad Dual-Averaging algorithm.
  *
@@ -188,9 +186,8 @@ protected void createSlots(List<Output<? extends TType>> variables) {
     for (Output<? extends TType> v : variables) {
       createAdaGradDASlot(v);
     }
-    globalStep = tf.withName("adagrad-da-global-step").variable(Shape.scalar(), TInt64.class);
-    Assign<TInt64> globalStepInitializer = tf.assign(globalStep, tf.constant(0L));
-    graph.addInitializer(globalStepInitializer);
+    globalStep = tf.initScope().withName("adagrad-da-global-step").variable(Shape.scalar(), TInt64.class);
+    tf.initScope().assign(globalStep, tf.constant(0L));
   }
 
   /**
@@ -200,10 +197,12 @@ protected void createSlots(List<Output<? extends TType>> variables) {
    * @param <T> the datatype of the variable.
    */
   private <T extends TType> void createAdaGradDASlot(Output<T> v) {
-    Operand<T> initializer = tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(0.0f), v.type()));
+    Operand<T> initializer =
+        tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(0.0f), v.type()));
     createSlot(v.asOutput(), ACCUMULATOR, initializer);
     Operand<T> sqInitializer =
-        tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
+        tf.fill(
+            tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
     createSlot(v.asOutput(), SQUARED_ACCUMULATOR, sqInitializer);
   }
 

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Adam.java

@@ -15,6 +15,8 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
+import java.util.Optional;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -30,9 +32,6 @@
 import org.tensorflow.types.TFloat32;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-import java.util.Optional;
-
 /**
  * Optimizer that implements the Adam algorithm.
  *
@@ -190,12 +189,10 @@ protected void createSlots(List<Output<? extends TType>> variables) {
     for (Output<? extends TType> v : variables) {
       createAdamSlot(v.asOutput());
     }
-    betaOnePower = tf.withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> betaOnePowerInit = tf.assign(betaOnePower, tf.constant(betaOne));
-    graph.addInitializer(betaOnePowerInit);
-    betaTwoPower = tf.withName("beta2_power").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> betaTwoPowerInit = tf.assign(betaTwoPower, tf.constant(betaTwo));
-    graph.addInitializer(betaTwoPowerInit);
+    betaOnePower = tf.initScope().withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(betaOnePower, tf.constant(betaOne));
+    betaTwoPower = tf.initScope().withName("beta2_power").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(betaTwoPower, tf.constant(betaTwo));
   }
 
   /** {@inheritDoc} */

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Adamax.java

@@ -1,5 +1,7 @@
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
+import java.util.Optional;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -12,9 +14,6 @@
 import org.tensorflow.types.TFloat32;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-import java.util.Optional;
-
 /**
  * Optimizer that implements the Adamax algorithm.
  *
@@ -135,9 +134,8 @@ protected void createSlots(List<Output<? extends TType>> variables) {
     for (Output<? extends TType> v : variables) {
       createAdamaxSlot(v.asOutput());
     }
-    betaOnePower = tf.withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> betaOnePowerInit = tf.assign(betaOnePower, tf.constant(betaOne));
-    ((Graph) tf.scope().env()).addInitializer(betaOnePowerInit);
+    betaOnePower = tf.initScope().withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(betaOnePower, tf.constant(betaOne));
   }
 
   /**

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Ftrl.java

@@ -1,5 +1,6 @@
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -8,8 +9,6 @@
 import org.tensorflow.op.train.ApplyFtrl;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-
 /**
  * Optimizer that implements the FTRL algorithm.
  *
@@ -230,7 +229,8 @@ protected void createSlots(List<Output<? extends TType>> variables) {
    */
   private <T extends TType> void createFtrlSlot(Output<T> v) {
     Operand<T> initializer =
-        tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
+        tf.fill(
+            tf.shape(v), tf.dtypes.cast(tf.constant(initialAccumulatorValue), v.type()));
     createSlot(v.asOutput(), ACCUMULATOR, initializer);
     Operand<T> linearInitializer =
         tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(0.0f), v.type()));

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Momentum.java

@@ -15,6 +15,7 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -23,8 +24,6 @@
 import org.tensorflow.op.train.ApplyMomentum;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-
 /**
  * Stochastic gradient descent plus momentum, either nesterov or traditional.
  *
@@ -125,7 +124,8 @@ protected void createSlots(List<Output<? extends TType>> variables) {
    * @param <T> the data type of the variable
    */
   private <T extends TType> void createMomentumSlot(Output<T> v) {
-    Operand<T> initializer = tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(0.0f), v.type()));
+    Operand<T> initializer =
+        tf.fill(tf.shape(v), tf.dtypes.cast(tf.constant(0.0f), v.type()));
     createSlot(v.asOutput(), MOMENTUM, initializer);
   }
 

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Nadam.java

@@ -1,5 +1,7 @@
 package org.tensorflow.framework.optimizers;
 
+import java.util.List;
+import java.util.Optional;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Output;
@@ -12,9 +14,6 @@
 import org.tensorflow.types.TInt64;
 import org.tensorflow.types.family.TType;
 
-import java.util.List;
-import java.util.Optional;
-
 /**
  * Nadam Optimizer that implements the NAdam algorithm.
  *
@@ -140,17 +139,14 @@ protected void createSlots(List<Output<? extends TType>> variables) {
     for (Output<? extends TType> v : variables) {
       createNadamSlot(v.asOutput());
     }
-    betaOnePower = tf.withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> betaOnePowerInit = tf.assign(betaOnePower, tf.constant(betaOne));
-    ((Graph) tf.scope().env()).addInitializer(betaOnePowerInit);
+    betaOnePower = tf.initScope().withName("beta1_power").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(betaOnePower, tf.constant(betaOne));
 
-    betaTwoPower = tf.withName("beta2_power").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> betaTwoPowerInit = tf.assign(betaTwoPower, tf.constant(betaTwo));
-    ((Graph) tf.scope().env()).addInitializer(betaTwoPowerInit);
+    betaTwoPower = tf.initScope().withName("beta2_power").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(betaTwoPower, tf.constant(betaTwo));
 
-    momentum = tf.withName("momentum").variable(Shape.scalar(), TFloat32.class);
-    Assign<TFloat32> momentumInit = tf.assign(momentum, tf.constant(1.0F));
-    ((Graph) tf.scope().env()).addInitializer(momentumInit);
+    momentum = tf.initScope().withName("momentum").variable(Shape.scalar(), TFloat32.class);
+    tf.initScope().assign(momentum, tf.constant(1.0F));
   }
 
   /**

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Optimizer.java

@@ -15,21 +15,19 @@
  */
 package org.tensorflow.framework.optimizers;
 
+import java.util.*;
+import java.util.stream.Collectors;
 import org.tensorflow.Graph;
 import org.tensorflow.Operand;
 import org.tensorflow.Operation;
 import org.tensorflow.Output;
 import org.tensorflow.op.Op;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.Scope;
-import org.tensorflow.op.core.Assign;
 import org.tensorflow.op.core.NoOp;
 import org.tensorflow.op.core.Variable;
 import org.tensorflow.types.family.TType;
 
-import java.util.*;
-import java.util.stream.Collectors;
-
 /** Base class for gradient optimizers. */
 public abstract class Optimizer {
 
@@ -41,6 +39,7 @@ public abstract class Optimizer {
   protected final Graph graph;
   /** The ops builder for the graph. */
   protected final Ops tf;
+
   /** Top level map key is the variable name, lower level map key is the slot name. */
   private final Map<String, Map<String, Variable<?>>> slots;
 
@@ -221,9 +220,10 @@ private <T extends TType> Optional<Variable<T>> getSlot(String varName, String s
   protected <T extends TType> void createSlot(
       Output<T> variable, String slotName, Operand<T> initializer) {
     Variable<T> slot =
-        tf.withName(createName(variable, slotName)).variable(variable.shape(), variable.type());
-    Assign<T> slotInit = tf.assign(slot, initializer);
-    graph.addInitializer(slotInit);
+        tf.initScope()
+            .withName(createName(variable, slotName))
+            .variable(variable.shape(), variable.type());
+    tf.initScope().assign(slot, initializer);
     String varName = variable.op().name();
     Map<String, Variable<? extends TType>> variables =
         slots.computeIfAbsent(slotName, (k) -> new HashMap<>());