Run Java API Import for TF2.6.0

Author: karllessard

CONTRIBUTING.md

@@ -21,7 +21,7 @@ the `dev` profile in your Maven command to use those artifacts instead of buildi
 Modifying the native op generation code (not the annotation processor) or the JavaCPP configuration (not the abstract Pointers) will require a
 complete build could be required to reflect the changes, otherwise `-Pdev` should be fine.
 
-## JDK 16+
+### JDK 16+
 
 If you're using JDK 16+, you need to add some exports for the formatter plugin:
 
@@ -98,6 +98,63 @@ Failed to execute goal org.apache.maven.plugins:maven-surefire-plugin:2.22.0:tes
 This is because the native code crashed (i.e. because of a segfault), and it should have created a dump file somewhere in the project that you can use
 to tell what caused the issue.
 
+## Upgrading TensorFlow Version
+
+To upgrade the version of TensorFlow that is embedded within TensorFlow Java, please follow carefully these steps.
+
+### Upgrading TensorFlow Runtime Library
+
+You can upgrade the version of the TensorFlow library by updating the archive downloadeded in the Bazel
+[workspace](https://github.com/tensorflow/java/blob/master/tensorflow-core/tensorflow-core-api/WORKSPACE#L19) at build time. Make sure to
+update the `urls`, `sha256` and `strip_prefix` fields of the `org_tensorflow` archive rule to reflect the values for the new version.
+
+### Ops Classification
+
+After building with the version of TensorFlow, you might notice that a lot of new operations appeared in the `org.tensorflow.ops.core`
+package of the [generated sources](https://github.com/tensorflow/java/tree/master/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core) of
+the `tensorflow-core-api` module. Many of these ops must be reclassified manually after running this initial build.
+
+The actual classification process is a bit arbitrary and based on the good jugement of the developer. The reason is that most ops in Python
+are being wrapped by a higher-level API and therefore are left unclassified, while in Java they are exposed and can be used directly by
+the users.
+
+For classifying an op, a `api_def` proto must be added to the `tensorflow-core-api` [folder](https://github.com/tensorflow/java/tree/master/tensorflow-core/tensorflow-core-api/src/bazel/api_def)
+for this purpose, redefining optionally its endpoints or its visibility.
+
+Writing these protos and trying the guess the right location for each new operation can become a tedious job so an utility program called `java_api_import`
+has been created to help you with this task. This utility is available under the `bazel-bin` folder of `tensorflow-core-api` after the
+initial build. Here is how to invoke it:
+
+```
+cd tensorflow-core/tensorflow-core-api
+./bazel-bin/java_api_import \
+  --java_api_dir=src/bazel/api_def \
+  --tf_src_dir=bazel-tensorflow-core-api/external/org_tensorflow \
+  --tf_lib_path=bazel-bin/external/org_tensorflow/tensorflow/libtensorflow_cc.<version>.<ext>
+```
+
+For each new operation detected (i.e. any operation that does not have a valid `api_def` proto yet), the utility will suggest you some possible
+package names that can be a good match for its classification (unless a "perfect match" has been found in the Python code, in which case the utility
+will automatically classify the op). It is also possible to enter manually the name of the package to use, and the package can have multiple levels (e.g. `linalg.sparse`). The utility
+application will then take care to write the `api_def` proto for each operation classified.
+
+Make sure to erase completely the generated source folder of the `tensorflow-core-api` module before rerunning the build so you can see
+if your ops have been classified properly.
+
+#### Ops Kernel Upgrade
+
+Some operations might be just an upgrade of another existing operations. For instance, there are many version of the `BatchMatMul` kernel (V1, V2, V3...).
+When you see that a new op is just an upgrade from another other one, make sure that the latest version has a valid endpoint and that all other
+previous versions of this operation are marked as `VISIBILITY: SKIP`.
+
+### Java Protos Classification
+
+TensorFlow Java distributes a large number proto definitions found in the TensorFlow Runtime Library as Java classes. Again, new protos might not
+be classified properly since they may be lacking the `option java_*` statements at the beginning of their definition. If you notice in the
+[generated protos](https://github.com/tensorflow/java/tree/master/tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/proto) of the `tensorflow-core-api`
+that some new proto classes seems to be in the wrong package, create a Bazel patch at this effect to add the missing options.
+See [existing patches](https://github.com/tensorflow/java/blob/master/tensorflow-core/tensorflow-core-api/external/tensorflow-proto.patch) for examples.
+
 ## Contributing
 
 ### Formatting

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

@@ -1,6 +1,4 @@
 op {
   graph_op_name: "BatchMatMulV2"
-  endpoint {
-    name: "train.BatchMatMul"
-  }
+  visibility: SKIP
 }

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

@@ -0,0 +1,6 @@
+op {
+  graph_op_name: "BatchMatMulV3"
+  endpoint {
+    name: "train.BatchMatMul"
+  }
+}

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

@@ -0,0 +1,3 @@
+op {
+  graph_op_name: "CompositeTensorVariantFromComponents"
+}

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

@@ -0,0 +1,3 @@
+op {
+  graph_op_name: "CompositeTensorVariantToComponents"
+}

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

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

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

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

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

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

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

@@ -0,0 +1,3 @@
+op {
+  graph_op_name: "Window"
+}

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

@@ -0,0 +1,6 @@
+op {
+  graph_op_name: "XlaRemoveDynamicDimensionSize"
+  endpoint {
+    name: "xla.RemoveDynamicDimensionSize"
+  }
+}

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

@@ -59,7 +59,6 @@
 import org.tensorflow.op.core.BarrierTakeMany;
 import org.tensorflow.op.core.Batch;
 import org.tensorflow.op.core.BatchFunction;
-import org.tensorflow.op.core.BatchMatMulV3;
 import org.tensorflow.op.core.BatchToSpace;
 import org.tensorflow.op.core.BatchToSpaceNd;
 import org.tensorflow.op.core.Bitcast;
@@ -213,7 +212,6 @@
 import org.tensorflow.op.core.Slice;
 import org.tensorflow.op.core.Snapshot;
 import org.tensorflow.op.core.SpaceToBatchNd;
-import org.tensorflow.op.core.SparseSegmentSumGrad;
 import org.tensorflow.op.core.Split;
 import org.tensorflow.op.core.SplitV;
 import org.tensorflow.op.core.Squeeze;
@@ -296,7 +294,6 @@
 import org.tensorflow.op.core.VariableShape;
 import org.tensorflow.op.core.Where;
 import org.tensorflow.op.core.While;
-import org.tensorflow.op.core.XlaRemoveDynamicDimensionSize;
 import org.tensorflow.op.core.Zeros;
 import org.tensorflow.op.core.ZerosLike;
 import org.tensorflow.types.TBool;
@@ -842,42 +839,6 @@ public BatchFunction batchFunction(Iterable<Operand<?>> inTensors,
     return BatchFunction.create(scope, inTensors, capturedTensors, f, numBatchThreads, maxBatchSize, batchTimeoutMicros, Tout, options);
   }
 
-  /**
-   * Multiplies slices of two tensors in batches.
-   *  Multiplies all slices of {@code Tensor} {@code x} and {@code y} (each slice can be
-   *  viewed as an element of a batch), and arranges the individual results
-   *  in a single output tensor of the same batch size. Each of the
-   *  individual slices can optionally be adjointed (to adjoint a matrix
-   *  means to transpose and conjugate it) before multiplication by setting
-   *  the {@code adj_x} or {@code adj_y} flag to {@code True}, which are by default {@code False}.
-   *  <p>The input tensors {@code x} and {@code y} are 2-D or higher with shape {@code [..., r_x, c_x]}
-   *  and {@code [..., r_y, c_y]}.
-   *  <p>The output tensor is 2-D or higher with shape {@code [..., r_o, c_o]}, where:
-   *  <pre>
-   *  r_o = c_x if adj_x else r_x
-   *  c_o = r_y if adj_y else c_y
-   *  </pre>
-   *  <p>It is computed as:
-   *  <pre>
-   *  output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :])
-   *  </pre>
-   *  <p><em>NOTE</em>: {@code BatchMatMulV3} supports broadcasting in the batch dimensions. More
-   *  about broadcasting
-   *   <a href="http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html">here</a> .
-   *
-   * @param <V> data type for {@code output} output
-   * @param x 2-D or higher with shape {@code [..., r_x, c_x]}.
-   * @param y 2-D or higher with shape {@code [..., r_y, c_y]}.
-   * @param Tout If not spcified, Tout is the same type to input type.
-   * @param options carries optional attribute values
-   * @param <V> data type for {@code BatchMatMulV3} output and operands
-   * @return a new instance of BatchMatMulV3
-   */
-  public <V extends TType> BatchMatMulV3<V> batchMatMulV3(Operand<? extends TType> x,
-      Operand<? extends TType> y, Class<V> Tout, BatchMatMulV3.Options... options) {
-    return BatchMatMulV3.create(scope, x, y, Tout, options);
-  }
-
   /**
    * BatchToSpace for 4-D tensors of type T.
    *  This is a legacy version of the more general BatchToSpaceND.
@@ -5856,25 +5817,6 @@ public <T extends TType> SpaceToBatchNd<T> spaceToBatchNd(Operand<T> input,
     return SpaceToBatchNd.create(scope, input, blockShape, paddings);
   }
 
-  /**
-   * Computes gradients for SparseSegmentSum.
-   *  Returns tensor &quot;output&quot; with same shape as grad, except for dimension 0 whose
-   *  value is output_dim0.
-   *
-   * @param <T> data type for {@code output} output
-   * @param grad gradient propagated to the SparseSegmentSum op.
-   * @param indices indices passed to the corresponding SparseSegmentSum op.
-   * @param segmentIds segment_ids passed to the corresponding SparseSegmentSum op.
-   * @param outputDim0 dimension 0 of &quot;data&quot; passed to SparseSegmentSum op.
-   * @param <T> data type for {@code SparseSegmentSumGrad} output and operands
-   * @return a new instance of SparseSegmentSumGrad
-   */
-  public <T extends TNumber> SparseSegmentSumGrad<T> sparseSegmentSumGrad(Operand<T> grad,
-      Operand<? extends TNumber> indices, Operand<? extends TNumber> segmentIds,
-      Operand<TInt32> outputDim0) {
-    return SparseSegmentSumGrad.create(scope, grad, indices, segmentIds, outputDim0);
-  }
-
   /**
    * Splits a tensor into {@code num_split} tensors along one dimension.
    *
@@ -8099,24 +8041,6 @@ public While whileOp(Iterable<Operand<?>> input, ConcreteFunction cond, Concrete
     return While.create(scope, input, cond, body, options);
   }
 
-  /**
-   * Inverse of XlaSetDynamicDimensionSize. Make an xla bounded
-   *  <pre>
-   *      dynamic dimension into a static dimension. The bound of the size of
-   *      dimension `dim_index` becomes the static dimension size.
-   *  </pre>
-   *
-   * @param <T> data type for {@code output} output
-   * @param input the input value
-   * @param dimIndex the dimIndex value
-   * @param <T> data type for {@code XlaRemoveDynamicDimensionSize} output and operands
-   * @return a new instance of XlaRemoveDynamicDimensionSize
-   */
-  public <T extends TType> XlaRemoveDynamicDimensionSize<T> xlaRemoveDynamicDimensionSize(
-      Operand<T> input, Operand<TInt32> dimIndex) {
-    return XlaRemoveDynamicDimensionSize.create(scope, input, dimIndex);
-  }
-
   /**
    * Creates a zeroed tensor given its type and shape.
    *

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

@@ -52,6 +52,7 @@
 import org.tensorflow.op.sparse.SparseSegmentSqrtNGrad;
 import org.tensorflow.op.sparse.SparseSegmentSqrtNWithNumSegments;
 import org.tensorflow.op.sparse.SparseSegmentSum;
+import org.tensorflow.op.sparse.SparseSegmentSumGrad;
 import org.tensorflow.op.sparse.SparseSegmentSumWithNumSegments;
 import org.tensorflow.op.sparse.SparseSlice;
 import org.tensorflow.op.sparse.SparseSliceGrad;
@@ -1053,6 +1054,25 @@ public <T extends TNumber> SparseSegmentSum<T> sparseSegmentSum(Operand<T> data,
     return SparseSegmentSum.create(scope, data, indices, segmentIds);
   }
 
+  /**
+   * Computes gradients for SparseSegmentSum.
+   *  Returns tensor &quot;output&quot; with same shape as grad, except for dimension 0 whose
+   *  value is output_dim0.
+   *
+   * @param <T> data type for {@code output} output
+   * @param grad gradient propagated to the SparseSegmentSum op.
+   * @param indices indices passed to the corresponding SparseSegmentSum op.
+   * @param segmentIds segment_ids passed to the corresponding SparseSegmentSum op.
+   * @param outputDim0 dimension 0 of &quot;data&quot; passed to SparseSegmentSum op.
+   * @param <T> data type for {@code SparseSegmentSumGrad} output and operands
+   * @return a new instance of SparseSegmentSumGrad
+   */
+  public <T extends TNumber> SparseSegmentSumGrad<T> sparseSegmentSumGrad(Operand<T> grad,
+      Operand<? extends TNumber> indices, Operand<? extends TNumber> segmentIds,
+      Operand<TInt32> outputDim0) {
+    return SparseSegmentSumGrad.create(scope, grad, indices, segmentIds, outputDim0);
+  }
+
   /**
    * Computes the sum along sparse segments of a tensor.
    *  Like {@code SparseSegmentSum}, but allows missing ids in {@code segment_ids}. If an id is

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

@@ -510,16 +510,17 @@ public <T extends TType> ApplyRmsProp<T> applyRmsProp(Operand<T> var, Operand<T>
    *  about broadcasting
    *   <a href="http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html">here</a> .
    *
-   * @param <T> data type for {@code output} output
+   * @param <V> data type for {@code output} output
    * @param x 2-D or higher with shape {@code [..., r_x, c_x]}.
    * @param y 2-D or higher with shape {@code [..., r_y, c_y]}.
+   * @param Tout If not spcified, Tout is the same type to input type.
    * @param options carries optional attribute values
-   * @param <T> data type for {@code BatchMatMulV2} output and operands
+   * @param <V> data type for {@code BatchMatMulV3} output and operands
    * @return a new instance of BatchMatMul
    */
-  public <T extends TType> BatchMatMul<T> batchMatMul(Operand<T> x, Operand<T> y,
-      BatchMatMul.Options... options) {
-    return BatchMatMul.create(scope, x, y, options);
+  public <V extends TType> BatchMatMul<V> batchMatMul(Operand<? extends TType> x,
+      Operand<? extends TType> y, Class<V> Tout, BatchMatMul.Options... options) {
+    return BatchMatMul.create(scope, x, y, Tout, options);
   }
 
   /**

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

@@ -36,6 +36,7 @@
 import org.tensorflow.op.xla.Recv;
 import org.tensorflow.op.xla.Reduce;
 import org.tensorflow.op.xla.ReduceWindow;
+import org.tensorflow.op.xla.RemoveDynamicDimensionSize;
 import org.tensorflow.op.xla.ReplicaId;
 import org.tensorflow.op.xla.Scatter;
 import org.tensorflow.op.xla.SelectAndScatter;
@@ -369,6 +370,24 @@ public <T extends TType, U extends TNumber> ReduceWindow<T> reduceWindow(Operand
     return ReduceWindow.create(scope, input, initValue, windowDimensions, windowStrides, baseDilations, windowDilations, padding, computation);
   }
 
+  /**
+   * Inverse of XlaSetDynamicDimensionSize. Make an xla bounded
+   *  <pre>
+   *      dynamic dimension into a static dimension. The bound of the size of
+   *      dimension `dim_index` becomes the static dimension size.
+   *  </pre>
+   *
+   * @param <T> data type for {@code output} output
+   * @param input the input value
+   * @param dimIndex the dimIndex value
+   * @param <T> data type for {@code XlaRemoveDynamicDimensionSize} output and operands
+   * @return a new instance of RemoveDynamicDimensionSize
+   */
+  public <T extends TType> RemoveDynamicDimensionSize<T> removeDynamicDimensionSize(
+      Operand<T> input, Operand<TInt32> dimIndex) {
+    return RemoveDynamicDimensionSize.create(scope, input, dimIndex);
+  }
+
   /**
    * Replica ID.
    *