Add WithOps, use for KotlinOps

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

Author: rnett

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

@@ -18,6 +18,7 @@
 package org.tensorflow.op;
 
 import java.nio.charset.Charset;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import org.tensorflow.ConcreteFunction;
@@ -330,7 +331,7 @@
  * }
  * }</pre>
  */
-public final class Ops {
+public final class Ops implements WithOps {
   public final NnOps nn;
 
   public final SummaryOps summary;
@@ -7906,42 +7907,51 @@ public <T extends TType> ZerosLike<T> zerosLike(Operand<T> x) {
     return ZerosLike.create(scope, x);
   }
 
+  @Override
+  public Ops tf() {
+    return this;
+  }
+
   /**
-   * Returns an API that builds operations with the provided name prefix.
-   *
-   * @see {@link Scope#withSubScope(String)}
+   * {@inheritDoc}
    */
+  @Override
   public Ops withSubScope(String childScopeName) {
     return new Ops(scope.withSubScope(childScopeName));
   }
 
   /**
-   * Returns an API that uses the provided name for an op.
-   *
-   * @see {@link Scope#withName(String)}
+   * {@inheritDoc}
    */
+  @Override
   public Ops withName(String opName) {
     return new Ops(scope.withName(opName));
   }
 
   /**
-   * Returns an API that places the created operations on the device(s) matching the provided spec.
-   *
-   * @see {@link Scope#withDevice(DeviceSpec)}
+   * {@inheritDoc}
    */
+  @Override
   public Ops withDevice(DeviceSpec deviceSpec) {
     return new Ops(scope.withDevice(deviceSpec));
   }
 
   /**
-   * Returns an API that adds operations to the graph with the provided control dependencies.
-   *
-   * @see {@link Scope#withControlDependencies(Iterable<Op<?>>)}
+   * {@inheritDoc}
    */
+  @Override
   public Ops withControlDependencies(Iterable<Op> controls) {
     return new Ops(scope.withControlDependencies(controls));
   }
 
+  /**
+   * {@inheritDoc}
+   */
+  @Override
+  public Ops withControlDependencies(Op... controls) {
+    return withControlDependencies(Arrays.asList(controls));
+  }
+
   /**
    * Returns the current {@link Scope scope} of this API
    */

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

@@ -16,10 +16,12 @@
 package org.tensorflow;
 
 import org.tensorflow.op.Op;
+import org.tensorflow.op.Ops;
 import org.tensorflow.op.Scope;
+import org.tensorflow.op.WithOps;
 
 /** Defines an environment for creating and executing TensorFlow {@link Operation}s. */
-public interface ExecutionEnvironment {
+public interface ExecutionEnvironment extends WithOps {
 
   enum Types {
     GRAPH,
@@ -86,4 +88,9 @@ default boolean isGraph() {
    * prevent name collisions.
    */
   Scope baseScope();
+
+  @Override
+  default Ops tf(){
+    return Ops.create(this);
+  }
 }

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

@@ -0,0 +1,78 @@
+/*
+  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 java.util.Arrays;
+import org.tensorflow.DeviceSpec;
+
+/**
+ * A context that provides a TensorFlow op builder.
+ */
+public interface WithOps {
+
+  /**
+   * Get the op builder for this context.
+   */
+  Ops tf();
+
+  /**
+   * Returns an API that builds operations with the provided name prefix.
+   *
+   * @see Scope#withSubScope(String)
+   */
+  default WithOps withSubScope(String childScopeName) {
+    return tf().withSubScope(childScopeName);
+  }
+
+  /**
+   * Returns an API that uses the provided name for an op.
+   *
+   * @see Scope#withName(String)
+   */
+  default WithOps withName(String opName) {
+    return tf().withName(opName);
+  }
+
+  /**
+   * Returns an API that places the created operations on the device(s) matching the provided spec.
+   *
+   * @see Scope#withDevice(DeviceSpec)
+   */
+  default WithOps withDevice(DeviceSpec deviceSpec) {
+    return tf().withDevice(deviceSpec);
+  }
+
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * @see Scope#withControlDependencies(Iterable)
+   */
+  default WithOps withControlDependencies(Iterable<Op> controls){
+    return tf().withControlDependencies(controls);
+  }
+
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * @see Scope#withControlDependencies(Iterable)
+   */
+  default WithOps withControlDependencies(Op... controls){
+    return withControlDependencies(Arrays.asList(controls));
+  }
+
+}

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/Names.java

@@ -53,6 +53,8 @@ public class Names {
   public static final ClassName OperationBuilder = ClassName.get(TensorflowPackage, "OperationBuilder");
   public static final TypeName IterableOp = ParameterizedTypeName.get(ClassName.get(Iterable.class), Op);
 
+  public static final ClassName WithOps = ClassName.get(OpPackage, "WithOps");
+
   public static final ClassName Operand = ClassName.get(TensorflowPackage, "Operand");
   public static final ClassName Output = ClassName.get(TensorflowPackage, "Output");
 

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

@@ -15,11 +15,14 @@
  */
 package org.tensorflow.processor.operator;
 
+import com.squareup.javapoet.ArrayTypeName;
+import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.FieldSpec;
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.List;
 import javax.lang.model.element.Modifier;
 import org.tensorflow.Names;
@@ -108,6 +111,7 @@ protected TypeSpec buildTopClass(OpsSpec spec) {
 
     TypeSpec.Builder opsBuilder =
         TypeSpec.classBuilder("Ops")
+            .addSuperinterface(Names.WithOps)
             .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
             .addJavadoc(
                 "An API for building operations as {@link $T Op}s\n<p>\n"
@@ -146,52 +150,65 @@ protected TypeSpec buildTopClass(OpsSpec spec) {
 
     opsBuilder.addMethod(ctorBuilder.build());
 
+    opsBuilder.addMethod(MethodSpec
+        .methodBuilder("tf")
+        .addModifiers(Modifier.PUBLIC)
+        .addAnnotation(Override.class)
+        .returns(Names.Ops)
+        .addStatement("return this")
+        .build()
+    );
+
     opsBuilder.addMethod(
         MethodSpec.methodBuilder("withSubScope")
             .addModifiers(Modifier.PUBLIC)
+            .addAnnotation(Override.class)
             .addParameter(Names.String, "childScopeName")
             .returns(Names.Ops)
             .addStatement("return new $T(scope.withSubScope(childScopeName))", Names.Ops)
-            .addJavadoc(
-                "Returns an API that builds operations with the provided name prefix.\n"
-                    + "\n@see {@link $T#withSubScope(String)}\n",
-                Names.Scope)
+            .addJavadoc("{@inheritDoc}")
             .build());
 
     opsBuilder.addMethod(
         MethodSpec.methodBuilder("withName")
             .addModifiers(Modifier.PUBLIC)
+            .addAnnotation(Override.class)
             .addParameter(Names.String, "opName")
             .returns(Names.Ops)
             .addStatement("return new Ops(scope.withName(opName))")
-            .addJavadoc(
-                "Returns an API that uses the provided name for an op.\n\n"
-                    + "@see {@link $T#withName(String)}\n",
-                Names.Scope)
+            .addJavadoc("{@inheritDoc}")
             .build());
 
     opsBuilder.addMethod(
         MethodSpec.methodBuilder("withDevice")
             .addModifiers(Modifier.PUBLIC)
+            .addAnnotation(Override.class)
             .addParameter(Names.DeviceSpec, "deviceSpec")
             .returns(Names.Ops)
             .addStatement("return new Ops(scope.withDevice(deviceSpec))")
-            .addJavadoc(
-                "Returns an API that places the created operations on the device(s) matching the provided spec.\n\n"
-                    + "@see {@link $T#withDevice(DeviceSpec)}\n",
-                Names.Scope)
+            .addJavadoc("{@inheritDoc}")
             .build());
 
     opsBuilder.addMethod(
         MethodSpec.methodBuilder("withControlDependencies")
             .addModifiers(Modifier.PUBLIC)
+            .addAnnotation(Override.class)
             .addParameter(Names.IterableOp, "controls")
             .returns(Names.Ops)
             .addStatement("return new Ops(scope.withControlDependencies(controls))")
-            .addJavadoc(
-                "Returns an API that adds operations to the graph with the provided control dependencies.\n\n"
-                    + "@see {@link $T#withControlDependencies(Iterable<Op<?>>)}\n",
-                Names.Scope)
+            .addJavadoc("{@inheritDoc}")
+            .build());
+
+    opsBuilder.addMethod(
+        MethodSpec.methodBuilder("withControlDependencies")
+            .addModifiers(Modifier.PUBLIC)
+            .addAnnotation(Override.class)
+            .addParameter(ArrayTypeName.of(Names.Op), "controls")
+            .varargs()
+            .returns(Names.Ops)
+            .addStatement("return withControlDependencies($T.asList(controls))", ClassName.get(
+                Arrays.class))
+            .addJavadoc("{@inheritDoc}")
             .build());
 
     opsBuilder.addField(