Rebase fixes

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

Author: rnett

tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java

@@ -106,7 +106,6 @@
 import org.tensorflow.op.core.IdentityN;
 import org.tensorflow.op.core.If;
 import org.tensorflow.op.core.ImmutableConst;
-import org.tensorflow.op.core.Init;
 import org.tensorflow.op.core.InitializeTable;
 import org.tensorflow.op.core.InitializeTableFromTextFile;
 import org.tensorflow.op.core.InplaceAdd;
@@ -295,6 +294,12 @@
 import org.tensorflow.op.core.VariableShape;
 import org.tensorflow.op.core.Where;
 import org.tensorflow.op.core.While;
+import org.tensorflow.op.core.XlaConvV2;
+import org.tensorflow.op.core.XlaDotV2;
+import org.tensorflow.op.core.XlaSetDynamicDimensionSize;
+import org.tensorflow.op.core.XlaSpmdFullToShardShape;
+import org.tensorflow.op.core.XlaSpmdShardToFullShape;
+import org.tensorflow.op.core.XlaVariadicSort;
 import org.tensorflow.op.core.Zeros;
 import org.tensorflow.op.core.ZerosLike;
 import org.tensorflow.types.TBool;
@@ -366,20 +371,20 @@ public final class Ops {
 
   public final SparseOps sparse;
 
-  public final TpuOps tpu;
-
   public final BitwiseOps bitwise;
 
+  public final TpuOps tpu;
+
   public final MathOps math;
 
   public final AudioOps audio;
 
   public final SignalOps signal;
 
-  public final QuantizationOps quantization;
-
   public final TrainOps train;
 
+  public final QuantizationOps quantization;
+
   private final Scope scope;
 
   private Ops(Scope scope) {
@@ -397,13 +402,13 @@ private Ops(Scope scope) {
     random = new RandomOps(this);
     strings = new StringsOps(this);
     sparse = new SparseOps(this);
-    tpu = new TpuOps(this);
     bitwise = new BitwiseOps(this);
+    tpu = new TpuOps(this);
     math = new MathOps(this);
     audio = new AudioOps(this);
     signal = new SignalOps(this);
-    quantization = new QuantizationOps(this);
     train = new TrainOps(this);
+    quantization = new QuantizationOps(this);
   }
 
   /**
@@ -1952,14 +1957,15 @@ public Constant<TFloat32> constant(Shape shape, FloatDataBuffer data) {
   }
 
   /**
-   * Creates a scalar of {@code type}, with the value of {@code number}. {@code number} may be truncated if it does not
-   *  fit in the target type.
+   * Creates a scalar of {@code type}, with the value of {@code number}. {@code number} may be
+   *  truncated if it does not fit in the target type.
    *
-   * @param type the type of tensor to create.  Must be concrete (i.e. not {@link org.tensorflow.types.family.TFloating})
+   * @param type the type of tensor to create. Must be concrete (i.e. not {@link
+   *      org.tensorflow.types.family.TFloating})
    * @param number the value of the tensor
    * @return a constant of the passed type
-   * @throws IllegalArgumentException if the type is abstract (i.e. {@link org.tensorflow.types.family.TFloating}) or
-   *  unknown.
+   * @throws IllegalArgumentException if the type is abstract (i.e. {@link
+   *      org.tensorflow.types.family.TFloating}) or unknown.
    */
   public <T extends TNumber> Constant<T> constant(Class<T> type, Number number) {
     return Constant.tensorOf(scope, type, number);
@@ -1994,11 +2000,12 @@ public <T extends TType> Constant<T> constant(Class<T> type, Shape shape, ByteDa
   }
 
   /**
-   * Create a constant by making an immutable copy of {@code tensor}. {@code tensor} may be closed afterwards without
-   *  issue.
+   * Create a constant by making an immutable copy of {@code tensor}. {@code tensor} may be closed
+   *  afterwards without issue.
    *
    *  <p>Note: this endpoint cannot be simply called {@code constant} since it will conflict with
-   *  other endpoints accepting an NdArray in parameter {e.g. {@link #tensorOf(Scope, FloatNdArray)}}.
+   *  other endpoints accepting an NdArray in parameter {e.g. {@link #tensorOf(Scope,
+   *  FloatNdArray)}}.
    *
    * @param tensor a Tensor holding the constant value
    * @return a constant of the same data type as `tensor`
@@ -2008,8 +2015,8 @@ public <T extends TType> Constant<T> constantOf(T tensor) {
   }
 
   /**
-   * Creates a scalar of the same type as {@code toMatch}, with the value of {@code number}. {@code number} may be
-   *  truncated if it does not fit in the target type.
+   * Creates a scalar of the same type as {@code toMatch}, with the value of {@code number}. {@code
+   *  number} may be truncated if it does not fit in the target type.
    *
    * @param toMatch the operand providing the target type
    * @param number the value of the tensor
@@ -2993,80 +3000,6 @@ public <T extends TType> ImmutableConst<T> immutableConst(Class<T> dtype, Shape
     return ImmutableConst.create(scope, dtype, shape, memoryRegionName);
   }
 
-  /**
-   * Factory method to create an operation executing all initializers of a graph.
-   *
-   *  <p>All initializers added to a graph via
-   *  {@link org.tensorflow.op.core.Init#add(Scope, Op) tf.initAdd} are grouped together as a single
-   *  unit of computation in the graph. This operation must then be added to any graph using one or
-   *  more {@link Variable variables} and executed once before running the graph so the variable
-   *  states are initialized properly.</p>
-   *
-   *  <p>When the graph is built by the same process that is running the session, the initializers
-   *  can be invoked by executing this single endpoint. For example:</p>
-   *  <pre>{@code
-   *  try (Graph g = new Graph()) {
-   *    Variable<TInt32> x = tf.variable(tf.constant(10));  // initAdd is called implicitly
-   *    Variable<TInt32> y = tf.variable(tf.constant(20));  // idem
-   *    Add<TInt32> z = tf.math.add(x, y);
-   *
-   *    try (Session s = new Session(g)) {
-   *      s.run(tf.init());  // initialize all variables
-   *
-   *      try (TInt32 t = (TInt32)s.runner().fetch(z).run().get(0)) {
-   *        assertEquals(30, t.data().getInt());
-   *      }
-   *    }
-   *  }
-   *  }</pre>
-   *
-   *  <p>When the graph is built by a separate process, the initializers can be invoked by running
-   *  the init op by its name, which defaults to {@link org.tensorflow.op.core.Init#DEFAULT_NAME}.
-   *  For example:</p>
-   *  <pre>{@code
-   *  // Building the model
-   *  try (Graph g = new Graph()) {
-   *    Variable<TInt32> x = tf.variable(tf.constant(10));  // initAdd is called implicitly
-   *    Variable<TInt32> y = tf.variable(tf.constant(20));  // idem
-   *    Add<TInt32> z = tf.withName("z").math.add(x, y);
-   *
-   *    tf.init();  // add variables initializers to the graph, as Init.DEFAULT_NAME
-   *    // ...exporting graph as a saved model...
-   *  }
-   *
-   *  ...
-   *
-   *  // Running the model
-   *  try (SavedModelBundle model = SavedModelBundle.load("/path/to/model", "train")) {
-   *    model.session().run(Init.DEFAULT_NAME);
-   *
-   *    try (TInt32 t = (TInt32)s.runner().fetch("z").run().get(0)) {
-   *      assertEquals(30, t.data().getInt());
-   *    }
-   *  }
-   *  }</pre>
-   *
-   * @return an op grouping all initializers added to the graph
-   * @throws IllegalArgumentException if the execution environment in scope is not a graph
-   */
-  public Init init() {
-    return Init.create(scope);
-  }
-
-  /**
-   * Register an op as an initializer of the graph.
-   *
-   *  <p>Registered initializers are then grouped as a single unit of computation by adding
-   *  and executing an {@link org.tensorflow.op.core.Init#create(Scope) init} operation from a graph
-   *  session.  This is a no-op if executed in an eager session.
-   *
-   * @param initializer
-   * @see org.tensorflow.op.core.Init#create(Scope) init
-   */
-  public void initAdd(Op initializer) {
-    Init.add(scope, initializer);
-  }
-
   /**
    * Table initializer that takes two tensors for keys and values respectively.
    *
@@ -7948,9 +7881,9 @@ public VarIsInitializedOp varIsInitializedOp(Operand<? extends TType> resource)
 
   /**
    * Factory method to create a new Variable with it's initializer.
-   *  <p>
-   *  Only supported on Graph sessions as the {@link org.tensorflow.op.core.Assign} op
-   *  does not work in an EagerSession.
+   *
+   *  <p>Only supported on Graph sessions as the {@link org.tensorflow.op.core.Assign} op does not
+   *  work in an EagerSession.
    *
    * @param init The op to use to initialise this variable.
    * @param options carries optional attributes values
@@ -8143,6 +8076,26 @@ public Ops withSubScope(String childScopeName) {
     return new Ops(scope.withSubScope(childScopeName));
   }
 
+  /**
+   * Returns an API that builds init operations.
+   * <p>Init operations will be initialized at session creation, will have their inputs (and control inputs) made init ops as well, and are never used as control dependencies.
+   * Additionally, this scope drops all of its control dependencies.  If an input can not be made an init op (i.e. a Placeholder), will error on op creation.
+   */
+  public Ops initScope() {
+    return new Ops(scope.initScope());
+  }
+
+  /**
+   * Make {@code op} an init operation, doing the same for all of it's inputs (and control inputs).
+   * <p>Init operations will be initialized at session creation, will have their inputs (and control inputs) made init ops as well, and are never used as control dependencies.
+   * Additionally, this scope drops all of its control dependencies.  If an input can not be made an init op (i.e. a Placeholder), will error on op creation.
+   * @throws IllegalArgumentException if the op or one of its inputs can't be made an init op.
+   */
+  public <T extends Operand> T liftToInitScope(T op) {
+    scope.env().registerInitOp(op.op());
+    return op;
+  }
+
   /**
    * Returns an API that uses the provided name for an op.
    *

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

@@ -115,6 +115,7 @@ default ExecutionEnvironment initEnv() {
    *
    * <p><b>Should generally only be used internally, prefer {@link
    * org.tensorflow.op.Ops#initScope()}.</b>
+   *
    * @throws IllegalArgumentException if the op or one of its inputs can't be made an init op.
    */
   void registerInitOp(Operation op);

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

@@ -145,6 +145,7 @@ public void close() {
    * Execute the graph's initializers.
    *
    * <p>This runs any ops that have been created with an init scope.
+   *
    * @throws IllegalStateException if the session has already been initialized
    */
   public synchronized void initialize() {

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/op/generator/ClassGenerator.java

@@ -259,8 +259,10 @@ void buildClass() {
 
       if (iterable) {
         mode = RenderMode.LIST_OPERAND;
-        if (!isStateSubclass) {builder.addSuperinterface(
-            ParameterizedTypeName.get(ClassName.get(Iterable.class), operandType));}
+        if (!isStateSubclass) {
+          builder.addSuperinterface(
+              ParameterizedTypeName.get(ClassName.get(Iterable.class), operandType));
+        }
       } else {
         mode = RenderMode.OPERAND;
         if (!isStateSubclass) {
@@ -315,26 +317,27 @@ void buildClass() {
       buildInterfaceImpl();
     }
 
-    if (!isStateSelector) {// add op name field
-    builder.addField(
-        FieldSpec.builder(
-                TypeResolver.STRING, OP_NAME_FIELD,
+    if (!isStateSelector) { // add op name field
+      builder.addField(
+          FieldSpec.builder(
+                  TypeResolver.STRING,
+                  OP_NAME_FIELD,
                   Modifier.PUBLIC,
                   Modifier.STATIC,
                   Modifier.FINAL)
               .addJavadoc("$L", "The name of this op, as known by TensorFlow core engine")
               .initializer("$S", op.getName())
               .build());
 
-    // add output fields
-    if (op.getOutputArgCount() > 0) {
-      for (ArgDef output : op.getOutputArgList()) {
-        builder.addField(
-            resolver.typeOf(output).listIfIterable().javaType,
-            getJavaName(output),
-            Modifier.PRIVATE);
+      // add output fields
+      if (op.getOutputArgCount() > 0) {
+        for (ArgDef output : op.getOutputArgList()) {
+          builder.addField(
+              resolver.typeOf(output).listIfIterable().javaType,
+              getJavaName(output),
+              Modifier.PRIVATE);
+        }
       }
-    }
 
       buildConstructor();
     }
@@ -411,8 +414,7 @@ private void buildOptionsClass() {
       }
 
       // add the field
-      optionsBuilder.addField(
-          field.build());
+      optionsBuilder.addField(field.build());
     }
 
     // add a private constructor
@@ -503,8 +505,7 @@ private void buildFactoryMethods() {
     Set<TypeVariableName> typeVars = new LinkedHashSet<>(typeParams);
 
     body.addStatement(
-        "$T opBuilder = scope.opBuilder($L, $S)", Names.OperationBuilder,
-        OP_NAME_FIELD, className);
+        "$T opBuilder = scope.opBuilder($L, $S)", Names.OperationBuilder, OP_NAME_FIELD, className);
 
     List<String> functionArgs = new ArrayList<>();
     List<String> iterableFunctionArgs = new ArrayList<>();
@@ -599,9 +600,7 @@ private void buildFactoryMethods() {
       }
 
       factoryBuilder.addParameter(
-          ParameterSpec.builder(
-                  ArrayTypeName.of(optionsClassName), "options")
-              .build());
+          ParameterSpec.builder(ArrayTypeName.of(optionsClassName), "options").build());
       paramTags.put("options", CodeBlock.of("$L", "carries optional attribute values"));
       factoryBuilder.varargs();
 
@@ -619,7 +618,7 @@ private void buildFactoryMethods() {
       body.endControlFlow();
 
       body.endControlFlow();
-}
+    }
 
     body.addStatement(
         "return new $L(opBuilder.build())", typeParams.isEmpty() ? className : (className + "<>"));
@@ -826,14 +825,16 @@ private void buildInterfaceImpl() {
 
       if (isStateSelector) {
         asOutput.addModifiers(Modifier.ABSTRACT);
-      } else {if (uncheckedCast) {
-        asOutput.addAnnotation(
-            AnnotationSpec.builder(SuppressWarnings.class)
-                .addMember("value", "$S", "unchecked")
-                .build());
-        asOutput.addCode("return ($T) $L;", outputType, getJavaName(output));
       } else {
-        asOutput.addCode("return $L;", getJavaName(output));}
+        if (uncheckedCast) {
+          asOutput.addAnnotation(
+              AnnotationSpec.builder(SuppressWarnings.class)
+                  .addMember("value", "$S", "unchecked")
+                  .build());
+          asOutput.addCode("return ($T) $L;", outputType, getJavaName(output));
+        } else {
+          asOutput.addCode("return $L;", getJavaName(output));
+        }
       }
 
       builder.addMethod(asOutput.build());

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/processor/operator/OperatorProcessor.java

@@ -26,7 +26,6 @@
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterSpec;
-import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeVariableName;
@@ -40,8 +39,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.regex.Pattern;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.Filer;
@@ -575,8 +572,10 @@ private static TypeSpec buildTopClass(OpsSpec spec) {
             .returns(T)
             .addStatement("scope.env().registerInitOp(op.op())")
             .addStatement("return op")
-            .addJavadoc("Make {@code op} an init operation, doing the same for all of it's inputs (and control inputs).\n" + initScopeComment +
-                "\n@throws IllegalArgumentException if the op or one of its inputs can't be made an init op.")
+            .addJavadoc(
+                "Make {@code op} an init operation, doing the same for all of it's inputs (and control inputs).\n"
+                    + initScopeComment
+                    + "\n@throws IllegalArgumentException if the op or one of its inputs can't be made an init op.")
             .build());
 
     opsBuilder.addMethod(