Export ops.pbtxt in addition to ops.pb (#306)

Author: rnett

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

@@ -79,15 +79,18 @@ fi
 GEN_SRCS_DIR=src/gen/java
 mkdir -p $GEN_SRCS_DIR
 
-GEN_RESOURCE_DIR=src/gen/resources/org/tensorflow/op
+GEN_RESOURCE_DIR=src/gen/resources
 mkdir -p $GEN_RESOURCE_DIR
 
 if [[ -z "${SKIP_EXPORT:-}" ]]; then
   # Export op defs
   echo "Exporting Ops"
   $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_LIB \
+      $GEN_RESOURCE_DIR/ops.pb \
+      $GEN_RESOURCE_DIR/ops.pbtxt \
+      $BAZEL_SRCS/external/org_tensorflow/tensorflow/core/api_def/base_api \
+      src/bazel/api_def
 else
   echo "Skipping Op export"
 fi

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

@@ -45,15 +45,15 @@
       <artifactId>ndarray</artifactId>
       <version>${project.version}</version>
     </dependency>
-      <dependency>
-        <groupId>org.junit.jupiter</groupId>
-        <artifactId>junit-jupiter-api</artifactId>
-        <scope>test</scope>
-      </dependency>
-      <dependency>
-        <groupId>org.junit.jupiter</groupId>
-        <artifactId>junit-jupiter-engine</artifactId>
-        <scope>test</scope>
+    <dependency>
+      <groupId>org.junit.jupiter</groupId>
+      <artifactId>junit-jupiter-api</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.junit.jupiter</groupId>
+      <artifactId>junit-jupiter-engine</artifactId>
+      <scope>test</scope>
     </dependency>
     <dependency>
       <groupId>org.openjdk.jmh</groupId>
@@ -359,7 +359,7 @@
           <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>
+            <argument>${project.basedir}/src/gen/resources/ops.pb</argument>
           </arguments>
         </configuration>
       </plugin>
@@ -423,10 +423,10 @@
           </execution>
         </executions>
         <configuration>
-            <!-- Activate the use of TCP to transmit events to the plugin -->
-            <!-- disabled as it appears to cause intermittent test failures in GitHub Actions
-            <forkNode implementation="org.apache.maven.plugin.surefire.extensions.SurefireForkNodeFactory"/>
-            -->
+          <!-- Activate the use of TCP to transmit events to the plugin -->
+          <!-- disabled as it appears to cause intermittent test failures in GitHub Actions
+          <forkNode implementation="org.apache.maven.plugin.surefire.extensions.SurefireForkNodeFactory"/>
+          -->
           <additionalClasspathElements>
             <additionalClasspathElement>${project.build.directory}/${project.artifactId}-${project.version}-${native.classifier}.jar</additionalClasspathElement>
             <!-- Note: the following path is not accessible in deploying profile, so other libraries like

tensorflow-core/tensorflow-core-api/src/bazel/op_generator/op_export_main.cc

@@ -15,6 +15,8 @@
 
 #include <string>
 #include <vector>
+#include <iostream>
+#include <fstream>
 #include <stdio.h>
 
 #include "tensorflow/core/framework/op.h"
@@ -43,7 +45,9 @@ const char kUsageHeader[] =
     "The first argument is the location of the tensorflow binary built for TF-"
     "Java.\nFor example, `bazel-out/k8-opt/bin/external/org_tensorflow/tensorfl"
     "ow/libtensorflow_cc.so`.\n\n"
-    "Finally, the `--api_dirs` argument takes a list of comma-separated "
+    "The second and third arguments are the binary and text output files, respectively.\n"
+    "The text output file will not include ApiDefs, like tensorflow's ops.pbtxt.\n\n"
+    "Finally, the rest of the arguments are used as "
     "directories of API definitions can be provided to override default\n"
     "values found in the ops definitions. Directories are ordered by priority "
     "(the last having precedence over the first).\nFor example, `bazel-tensorf"
@@ -86,32 +90,45 @@ Status UpdateOpDefs(OpList* op_list, const std::vector<tensorflow::string>& api_
 // See usage header.
 // Writes an OpList proto to stdout, with each OpDef having its ApiDef in field 100
 int main(int argc, char* argv[]) {
-  tensorflow::string api_dirs_str;
-  std::vector<tensorflow::Flag> flag_list = {
-      tensorflow::Flag(
-          "api_dirs", &api_dirs_str,
-          "List of directories that contain the ops API definitions protos")};
-  tensorflow::string usage = tensorflow::java::kUsageHeader;
-  usage += tensorflow::Flags::Usage(
-      tensorflow::string(argv[0]) + " <ops library paths...>", flag_list);
-  bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list);
-  tensorflow::port::InitMain(usage.c_str(), &argc, &argv);
-  QCHECK(parsed_flags_ok && argc > 1) << usage;
-  std::vector<tensorflow::string> api_dirs = tensorflow::str_util::Split(
-      api_dirs_str, ",", tensorflow::str_util::SkipEmpty());
+  tensorflow::port::InitMain(tensorflow::java::kUsageHeader, &argc, &argv);
+  std::vector<tensorflow::string> api_dirs;
 
-  tensorflow::Env* env = tensorflow::Env::Default();
-  void* ops_libs_handles[50];
-  for (int i = 1; i < argc; ++i) {
-    TF_CHECK_OK(env->LoadDynamicLibrary(argv[1], &ops_libs_handles[i - 1]));
+  if(argc < 4) {
+      std::cerr << "Must specify <library_path> <binary_output> <text_output>" << "\n";
+      std::cerr << tensorflow::java::kUsageHeader;
+      return 1;
+  }
+  
+  for(int i = 4 ; i < argc ; i++){
+    api_dirs.push_back(argv[i]);
+  }
+
+  std::ofstream binary_output (argv[2], std::ios::out | std::ios::trunc | std::ios::binary);
+  std::ofstream text_output  (argv[3], std::ios::out | std::ios::trunc);
+
+  if(!binary_output.is_open()){
+    std::cerr << "Error opening file " << argv[2] << "\n";
+    return 1;
+  }
+
+  if(!text_output.is_open()){
+    std::cerr << "Error opening file " << argv[3] << "\n";
+    return 1;
   }
+
+  tensorflow::Env* env = tensorflow::Env::Default();
+  void* ops_libs_handles[1];
+  TF_CHECK_OK(env->LoadDynamicLibrary(argv[1], &ops_libs_handles[0]));
   tensorflow::OpList ops;
   tensorflow::OpRegistry::Global()->Export(false, &ops);
-  TF_CHECK_OK(tensorflow::java::UpdateOpDefs(&ops, api_dirs, env));
 
+  text_output << ops.DebugString();
+  text_output.close();
+
+  TF_CHECK_OK(tensorflow::java::UpdateOpDefs(&ops, api_dirs, env));
 
-  std::ostream & out = std::cout;
-  ops.SerializeToOstream(&out);
+  ops.SerializeToOstream(&binary_output);
+  binary_output.close();
 
   return 0;
 }

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

@@ -354,20 +354,20 @@ public final class Ops {
 
   public final SparseOps sparse;
 
-  public final BitwiseOps bitwise;
-
   public final TpuOps tpu;
 
+  public final BitwiseOps bitwise;
+
   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) {
@@ -385,13 +385,13 @@ private Ops(Scope scope) {
     random = new RandomOps(this);
     strings = new StringsOps(this);
     sparse = new SparseOps(this);
-    bitwise = new BitwiseOps(this);
     tpu = new TpuOps(this);
+    bitwise = new BitwiseOps(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);
   }
 
   /**