Replace withInitScope with liftToInitScope

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

Author: rnett

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-generator/src/main/java/org/tensorflow/processor/operator/OperatorProcessor.java

@@ -566,33 +566,17 @@ private static TypeSpec buildTopClass(OpsSpec spec) {
             .addJavadoc("Returns an API that builds init operations.\n" + initScopeComment)
             .build());
 
+    TypeVariableName T = TypeVariableName.get("T").withBounds(Names.Operand);
     opsBuilder.addMethod(
-        MethodSpec.methodBuilder("withInitScope")
-            .addTypeVariable(TypeVariableName.get("R"))
+        MethodSpec.methodBuilder("liftToInitScope")
+            .addTypeVariable(T)
             .addModifiers(Modifier.PUBLIC)
-            .addParameter(
-                ParameterizedTypeName.get(
-                    ClassName.get(Function.class), Names.Ops, TypeVariableName.get("R")),
-                "block")
-            .returns(TypeVariableName.get("R"))
-            .addStatement("return block.apply(initScope())")
-            .addJavadoc(
-                "Call {@code block} with an init scope.\n"
-                    + initScopeComment
-                    + "\n@see #initScope()")
-            .build());
-
-    opsBuilder.addMethod(
-        MethodSpec.methodBuilder("withInitScope")
-            .addModifiers(Modifier.PUBLIC)
-            .addParameter(
-                ParameterizedTypeName.get(ClassName.get(Consumer.class), Names.Ops), "block")
-            .returns(TypeName.VOID)
-            .addStatement("block.accept(initScope())")
-            .addJavadoc(
-                "Call {@code block} with an init scope.\n"
-                    + initScopeComment
-                    + "\n@see #initScope()")
+            .addParameter(T, "op")
+            .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.")
             .build());
 
     opsBuilder.addMethod(