Move op generation to Java

.gitignore

@@ -16,6 +16,7 @@ __pycache__
 cmake_build/
 tensorflow/contrib/cmake/_build/
 .idea/**
+.run
 /build/
 [Bb]uild/
 /tensorflow/core/util/version_info.cc

tensorflow-core/tensorflow-core-api/build.sh

@@ -89,7 +89,7 @@ $BAZEL_BIN/java_op_generator \
 GEN_RESOURCE_DIR=src/gen/resources/org/tensorflow/op
 mkdir -p $GEN_RESOURCE_DIR
 
-# Generate Java operator wrappers
+# Export op defs
 $BAZEL_BIN/java_op_exporter \
     --api_dirs=$BAZEL_SRCS/external/org_tensorflow/tensorflow/core/api_def/base_api,src/bazel/api_def \
     $TENSORFLOW_LIB > $GEN_RESOURCE_DIR/ops.pb

tensorflow-core/tensorflow-core-api/pom.xml

@@ -328,6 +328,41 @@
           </dependency>
         </dependencies>
       </plugin>
+      <plugin>
+        <!--
+        An execution to execute the op class generator (which lives in tensorflow-core-generator).
+        Must be ran after build.sh, which generates the ops.pb file it reads.
+        Will be ran during the generate-sources phase.
+        -->
+        <groupId>org.codehaus.mojo</groupId>
+        <artifactId>exec-maven-plugin</artifactId>
+        <version>3.0.0</version>
+        <executions>
+          <execution>
+            <id>generate-ops</id>
+            <goals>
+              <goal>java</goal>
+            </goals>
+            <!--            <phase>generate-sources</phase>-->
+          </execution>
+        </executions>
+        <dependencies>
+          <dependency>
+            <groupId>org.tensorflow</groupId>
+            <artifactId>tensorflow-core-generator</artifactId>
+            <version>${project.version}</version>
+          </dependency>
+        </dependencies>
+        <configuration>
+          <includeProjectDependencies>false</includeProjectDependencies>
+          <includePluginDependencies>true</includePluginDependencies>
+          <mainClass>org.tensorflow.op.generator.OpGenerator</mainClass>
+          <arguments>
+            <argument>${project.basedir}/src/gen/java</argument>
+            <argument>${project.basedir}/src/gen/resources/org/tensorflow/op/ops.pb</argument>
+          </arguments>
+        </configuration>
+      </plugin>
       <plugin>
         <artifactId>maven-jar-plugin</artifactId>
         <version>3.1.0</version>

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/types/registry/TensorTypeRegistry.java

@@ -91,6 +91,7 @@ private static <T extends TType> void register(Class<T> type) {
 
   static {
     // TODO (karllessard) scan and registered automatically all annotated tensors types
+    // TODO use in generator?
     register(TBool.class);
     register(TFloat64.class);
     register(TFloat32.class);

tensorflow-core/tensorflow-core-generator/pom.xml

@@ -10,8 +10,8 @@
   <artifactId>tensorflow-core-generator</artifactId>
   <packaging>jar</packaging>
 
-  <name>TensorFlow Core Annotation Processor</name>
-  <description>Annotation processor for TensorFlow Java client</description>
+  <name>TensorFlow Core Generators</name>
+  <description>Code generators for TensorFlow Java client</description>
 
   <properties>
     <java.module.name>org.tensorflow.core.generator</java.module.name>
@@ -28,11 +28,22 @@
       <artifactId>javapoet</artifactId>
       <version>1.12.1</version>
     </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+      <version>${protobuf.version}</version>
+    </dependency>
     <dependency>
       <groupId>com.github.javaparser</groupId>
       <artifactId>javaparser-core</artifactId>
       <version>3.15.12</version>
     </dependency>
+    <!-- https://mvnrepository.com/artifact/org.commonmark/commonmark -->
+    <dependency>
+      <groupId>org.commonmark</groupId>
+      <artifactId>commonmark</artifactId>
+      <version>0.17.1</version>
+    </dependency>
   </dependencies>
 
   <build>

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

@@ -0,0 +1,73 @@
+/*
+  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;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.ParameterizedTypeName;
+import com.squareup.javapoet.TypeName;
+
+public class Names {
+
+  public static final String TensorflowPackage = "org.tensorflow";
+  public static final String OpPackage = TensorflowPackage + ".op";
+  public static final String TypesPackage = TensorflowPackage + ".types";
+
+  public static final ClassName Operator = ClassName.get(OpPackage + ".annotation", "Operator");
+  public static final ClassName Endpoint = ClassName.get(OpPackage + ".annotation", "Endpoint");
+
+  public static final ClassName TType = ClassName.get(TypesPackage + ".family", "TType");
+  public static final ClassName TString = ClassName.get(TypesPackage, "TString");
+  public static final ClassName TBool = ClassName.get(TypesPackage, "TBool");
+
+  public static final ClassName TNumber = ClassName.get(TypesPackage + ".family", "TNumber");
+
+  public static final ClassName TFloating = ClassName.get(TypesPackage + ".family", "TFloating");
+  public static final ClassName TBfloat16 = ClassName.get(TypesPackage, "TBfloat16");
+  public static final ClassName TFloat16 = ClassName.get(TypesPackage, "TFloat16");
+  public static final ClassName TFloat32 = ClassName.get(TypesPackage, "TFloat32");
+  public static final ClassName TFloat64 = ClassName.get(TypesPackage, "TFloat64");
+
+  public static final ClassName TIntegral = ClassName.get(TypesPackage + ".family", "TIntegral");
+  public static final ClassName TUint8 = ClassName.get(TypesPackage, "TUint8");
+  public static final ClassName TInt32 = ClassName.get(TypesPackage, "TInt32");
+  public static final ClassName TInt64 = ClassName.get(TypesPackage, "TInt64");
+
+  public static final TypeName Op = ClassName.get(OpPackage, "Op");
+  public static final ClassName RawOp = ClassName.get(OpPackage, "RawOp");
+  public static final ClassName Operation = ClassName.get(TensorflowPackage, "Operation");
+  public static final ClassName Operands = ClassName.get(OpPackage, "Operands");
+  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 Operand = ClassName.get(TensorflowPackage, "Operand");
+  public static final ClassName Output = ClassName.get(TensorflowPackage, "Output");
+
+  public static final ClassName Shape = ClassName.get(TensorflowPackage + ".ndarray", "Shape");
+  public static final ClassName Tensor = ClassName.get(TensorflowPackage, "Tensor");
+  public static final ClassName ConcreteFunction = ClassName.get(TensorflowPackage, "ConcreteFunction");
+
+  public static final ClassName Scope = ClassName.get(OpPackage, "Scope");
+  public static final TypeName DeviceSpec = ClassName.get(TensorflowPackage, "DeviceSpec");
+  public static final ClassName Ops = ClassName.get(OpPackage, "Ops");
+
+  public static final TypeName ExecutionEnvironment =
+      ClassName.get(TensorflowPackage, "ExecutionEnvironment");
+  public static final TypeName EagerSession = ClassName.get(TensorflowPackage, "EagerSession");
+
+  public static final TypeName String = ClassName.get(String.class);
+
+}