Add support for Tensor NIO

Author: karllessard

pom.xml

@@ -30,11 +30,40 @@
 
   <modules>
     <module>tensorflow-core</module>
-    <!--module>tensorflow-utils</module-->
+    <module>tensorflow-utils</module>
     <!--module>tensorflow-frameworks</module> TODO -->
     <!--module>tensorflow-starters</module> TODO -->
   </modules>
 
+  <properties>
+    <maven.compiler.source>1.8</maven.compiler.source>
+    <maven.compiler.target>1.8</maven.compiler.target>
+    <junit.version>4.12</junit.version>
+    <jmh.version>1.21</jmh.version>
+  </properties>
+
+  <dependencyManagement>
+    <dependencies>
+      <dependency>
+        <groupId>junit</groupId>
+        <artifactId>junit</artifactId>
+        <version>${junit.version}</version>
+      </dependency>
+	  <dependency>
+	    <groupId>org.openjdk.jmh</groupId>
+	    <artifactId>jmh-core</artifactId>
+	    <version>${jmh.version}</version>
+  	    <scope>test</scope>
+	  </dependency>
+	  <dependency>
+	    <groupId>org.openjdk.jmh</groupId>
+	    <artifactId>jmh-generator-annprocess</artifactId>
+	    <version>${jmh.version}</version>
+        <scope>test</scope>
+	  </dependency>
+    </dependencies>
+  </dependencyManagement>
+
   <!-- Two profiles are used:
        ossrh - deploys to ossrh/maven central
        bintray - deploys to bintray/jcenter. -->
@@ -64,6 +93,7 @@
       </distributionManagement>
     </profile>
   </profiles>
+
   <!-- http://central.sonatype.org/pages/requirements.html#developer-information -->
   <developers>
     <developer>
@@ -72,6 +102,7 @@
       <organizationUrl>http://www.tensorflow.org</organizationUrl>
     </developer>
   </developers>
+
   <build>
     <plugins>
       <!-- GPG signed components: http://central.sonatype.org/pages/apache-maven.html#gpg-signed-components -->

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

@@ -31,6 +31,11 @@
       <version>${project.version}</version>
       <optional>true</optional> <!-- for compilation only -->
     </dependency>
+    <dependency>
+      <groupId>org.tensorflow</groupId>
+      <artifactId>nio-utils</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>junit</groupId>
       <artifactId>junit</artifactId>
@@ -76,10 +81,6 @@
       <plugin>
         <artifactId>maven-compiler-plugin</artifactId>
         <version>3.8.0</version>
-        <configuration>
-          <source>1.7</source>
-          <target>1.7</target>
-        </configuration>
         <executions>
           <execution>
             <id>default-compile</id>
@@ -99,7 +100,7 @@
             </goals>
             <configuration>
               <includes>
-                <include>org/tensorflow/c_api/presets/*.java</include>
+                <include>org/tensorflow/internal/c_api/presets/*.java</include>
               </includes>
             </configuration>
           </execution>
@@ -150,6 +151,7 @@
             <compilerOption>${project.basedir}/src/main/native/server_jni.cc</compilerOption>
             <compilerOption>${project.basedir}/src/main/native/session_jni.cc</compilerOption>
             <compilerOption>${project.basedir}/src/main/native/tensorflow_jni.cc</compilerOption>
+            <compilerOption>${project.basedir}/src/main/native/tensor_buffers_jni.cc</compilerOption>
             <compilerOption>${project.basedir}/src/main/native/tensor_jni.cc</compilerOption>
             <compilerOption>${project.basedir}/src/main/native/utils_jni.cc</compilerOption>
           </compilerOptions>
@@ -201,7 +203,7 @@
             <configuration>
               <skip>${javacpp.parser.skip}</skip>
               <outputDirectory>${project.basedir}/src/gen/java</outputDirectory>
-              <classOrPackageName>org.tensorflow.c_api.presets.*</classOrPackageName>
+              <classOrPackageName>org.tensorflow.internal.c_api.presets.*</classOrPackageName>
             </configuration>
           </execution>
           <execution>
@@ -211,9 +213,9 @@
               <goal>build</goal>
             </goals>
             <configuration>
-              <outputDirectory>${project.build.directory}/native/org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/</outputDirectory>
+              <outputDirectory>${project.build.directory}/native/org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/</outputDirectory>
               <skip>${javacpp.compiler.skip}</skip>
-              <classOrPackageName>org.tensorflow.c_api.**</classOrPackageName>
+              <classOrPackageName>org.tensorflow.internal.c_api.**</classOrPackageName>
               <copyLibs>true</copyLibs>
               <copyResources>true</copyResources>
             </configuration>
@@ -224,6 +226,9 @@
         <artifactId>maven-surefire-plugin</artifactId>
         <version>2.22.0</version>
         <configuration>
+          <argLine>
+            -Djava.library.path=${project.build.directory}/native/org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}
+          </argLine>
           <additionalClasspathElements>${project.build.directory}/native/</additionalClasspathElements>
         </configuration>
       </plugin>
@@ -256,16 +261,16 @@
                 <!-- In case of successive builds for multiple platforms
                      without cleaning, ensures we only include files for
                      this platform. -->
-                <include>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/</include>
+                <include>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/</include>
               </includes>
               <classesDirectory>${project.build.directory}/native</classesDirectory>
               <excludes>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*.exp</exclude>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*.lib</exclude>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*.obj</exclude>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*mklml*</exclude>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*iomp5*</exclude>
-                <exclude>org/tensorflow/c_api/${javacpp.platform}${javacpp.platform.extension}/*msvcr120*</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*.exp</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*.lib</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*.obj</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*mklml*</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*iomp5*</exclude>
+                <exclude>org/tensorflow/internal/c_api/${javacpp.platform}${javacpp.platform.extension}/*msvcr120*</exclude>
               </excludes>
             </configuration>
           </execution>

tensorflow-core/tensorflow-core-api/src/bazel/api_def/api_def_KafkaDataset.pbtxt

@@ -0,0 +1,6 @@
+op {
+  graph_op_name: "KafkaDataset"
+  endpoint {
+    name: "data.KafkaDataset"
+  }
+}

tensorflow-core/tensorflow-core-api/src/bazel/op_generator/java_defs.h

@@ -97,22 +97,25 @@ class Type {
   static Type IterableOf(const Type& type) {
     return Interface("Iterable").add_parameter(type);
   }
+  static Type DataTypeOf(const Type& type) {
+    return Class("DataType", "org.tensorflow").add_parameter(type);
+  }
   static Type ForDataType(DataType data_type) {
     switch (data_type) {
       case DataType::DT_BOOL:
-        return Class("Boolean");
+        return Class("TBool", "org.tensorflow.types");
       case DataType::DT_STRING:
-        return Class("String");
+        return Class("TString", "org.tensorflow.types");
       case DataType::DT_FLOAT:
-        return Class("Float");
+        return Class("TFloat", "org.tensorflow.types");
       case DataType::DT_DOUBLE:
-        return Class("Double");
+        return Class("TDouble", "org.tensorflow.types");
       case DataType::DT_UINT8:
-        return Class("UInt8", "org.tensorflow.types");
+        return Class("TUInt8", "org.tensorflow.types");
       case DataType::DT_INT32:
-        return Class("Integer");
+        return Class("TInt32", "org.tensorflow.types");
       case DataType::DT_INT64:
-        return Class("Long");
+        return Class("TInt64", "org.tensorflow.types");
       case DataType::DT_RESOURCE:
         // TODO(karllessard) create a Resource utility class that could be
         // used to store a resource and its type (passed in a second argument).

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

@@ -124,24 +124,16 @@ void WriteSetAttrDirective(const AttributeSpec& attr, bool optional,
         .EndLine()
         .BeginBlock("for (int i = 0; i < " + array_name + ".length; ++i)")
         .Append(array_name + "[i] = ");
-    if (attr.type().kind() == Type::GENERIC) {
-      writer->Append("DataType.fromClass(" + var_name + ".get(i));");
-    } else {
-      writer->Append(var_name + ".get(i);");
-    }
+    writer->Append(var_name + ".get(i);");
     writer->EndLine()
         .EndBlock()
         .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ")
         .Append(array_name + ");")
         .EndLine();
   } else {
-    writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ");
-    if (attr.var().type().name() == "Class") {
-      writer->Append("DataType.fromClass(" + var_name + "));");
-    } else {
-      writer->Append(var_name + ");");
-    }
-    writer->EndLine();
+    writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ")
+        .Append(var_name + ");")
+        .EndLine();
   }
 }
 
@@ -179,7 +171,7 @@ void RenderSecondaryFactoryMethod(const OpSpec& op, const Type& op_class,
     if (attr.type().kind() == Type::GENERIC &&
         default_types.find(attr.type().name()) != default_types.end()) {
       factory_statement << default_types.at(attr.type().name()).name()
-                        << ".class";
+                        << ".DTYPE";
     } else {
       AddArgument(attr.var(), attr.description(), &factory, &factory_doc);
       factory_statement << attr.var().name();

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

@@ -148,7 +148,7 @@ std::pair<Type, Type> TypeResolver::TypesOf(const OpDef_AttrDef& attr_def,
     types = MakeTypePair(Type::Class("Boolean"), Type::Boolean());
 
   } else if (attr_type == "shape") {
-    types = MakeTypePair(Type::Class("Shape", "org.tensorflow"));
+    types = MakeTypePair(Type::Class("Shape", "org.tensorflow.nio.nd"));
 
   } else if (attr_type == "tensor") {
     types = MakeTypePair(Type::Class("Tensor", "org.tensorflow")
@@ -157,7 +157,7 @@ std::pair<Type, Type> TypeResolver::TypesOf(const OpDef_AttrDef& attr_def,
   } else if (attr_type == "type") {
     Type type = *iterable_out ? Type::Wildcard() : NextGeneric();
     if (IsRealNumbers(attr_def.allowed_values())) {
-      type.add_supertype(Type::Class("Number"));
+      type.add_supertype(Type::Class("TNumber", "org.tensorflow.types.family"));
     }
     types = MakeTypePair(type, Type::Enum("DataType", "org.tensorflow"));
 
@@ -305,7 +305,7 @@ AttributeSpec CreateAttribute(const OpDef_AttrDef& attr_def,
   bool iterable = false;
   std::pair<Type, Type> types = type_resolver->TypesOf(attr_def, &iterable);
   Type var_type = types.first.kind() == Type::GENERIC
-                      ? Type::ClassOf(types.first)
+                      ? Type::DataTypeOf(types.first)
                       : types.first;
   if (iterable) {
     var_type = Type::ListOf(var_type);

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

@@ -5,6 +5,9 @@
 import org.tensorflow.op.audio.DecodeWav;
 import org.tensorflow.op.audio.EncodeWav;
 import org.tensorflow.op.audio.Mfcc;
+import org.tensorflow.types.TFloat;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.TString;
 
 /**
  * An API for building {@code audio} operations as {@link Op Op}s
@@ -18,33 +21,6 @@ public final class AudioOps {
     this.scope = scope;
   }
 
-  /**
-   * Builds an {@link AudioSpectrogram} operation
-   *
-   * @param input Float representation of audio data.
-   * @param windowSize How wide the input window is in samples. For the highest efficiency
-   * @param stride How widely apart the center of adjacent sample windows should be.
-   * @param options carries optional attributes values
-   * @return a new instance of AudioSpectrogram
-   * @see org.tensorflow.op.audio.AudioSpectrogram
-   */
-  public AudioSpectrogram audioSpectrogram(Operand<Float> input, Long windowSize, Long stride,
-      AudioSpectrogram.Options... options) {
-    return AudioSpectrogram.create(scope, input, windowSize, stride, options);
-  }
-
-  /**
-   * Builds an {@link EncodeWav} operation
-   *
-   * @param audio 2-D with shape `[length, channels]`.
-   * @param sampleRate Scalar containing the sample frequency.
-   * @return a new instance of EncodeWav
-   * @see org.tensorflow.op.audio.EncodeWav
-   */
-  public EncodeWav encodeWav(Operand<Float> audio, Operand<Integer> sampleRate) {
-    return EncodeWav.create(scope, audio, sampleRate);
-  }
-
   /**
    * Builds an {@link DecodeWav} operation
    *
@@ -53,7 +29,7 @@ public EncodeWav encodeWav(Operand<Float> audio, Operand<Integer> sampleRate) {
    * @return a new instance of DecodeWav
    * @see org.tensorflow.op.audio.DecodeWav
    */
-  public DecodeWav decodeWav(Operand<String> contents, DecodeWav.Options... options) {
+  public DecodeWav decodeWav(Operand<TString> contents, DecodeWav.Options... options) {
     return DecodeWav.create(scope, contents, options);
   }
 
@@ -66,8 +42,35 @@ public DecodeWav decodeWav(Operand<String> contents, DecodeWav.Options... option
    * @return a new instance of Mfcc
    * @see org.tensorflow.op.audio.Mfcc
    */
-  public Mfcc mfcc(Operand<Float> spectrogram, Operand<Integer> sampleRate,
+  public Mfcc mfcc(Operand<TFloat> spectrogram, Operand<TInt32> sampleRate,
       Mfcc.Options... options) {
     return Mfcc.create(scope, spectrogram, sampleRate, options);
   }
+
+  /**
+   * Builds an {@link EncodeWav} operation
+   *
+   * @param audio 2-D with shape `[length, channels]`.
+   * @param sampleRate Scalar containing the sample frequency.
+   * @return a new instance of EncodeWav
+   * @see org.tensorflow.op.audio.EncodeWav
+   */
+  public EncodeWav encodeWav(Operand<TFloat> audio, Operand<TInt32> sampleRate) {
+    return EncodeWav.create(scope, audio, sampleRate);
+  }
+
+  /**
+   * Builds an {@link AudioSpectrogram} operation
+   *
+   * @param input Float representation of audio data.
+   * @param windowSize How wide the input window is in samples. For the highest efficiency
+   * @param stride How widely apart the center of adjacent sample windows should be.
+   * @param options carries optional attributes values
+   * @return a new instance of AudioSpectrogram
+   * @see org.tensorflow.op.audio.AudioSpectrogram
+   */
+  public AudioSpectrogram audioSpectrogram(Operand<TFloat> input, Long windowSize, Long stride,
+      AudioSpectrogram.Options... options) {
+    return AudioSpectrogram.create(scope, input, windowSize, stride, options);
+  }
 }