Upgrading to TF 2.21.0

README.md

@@ -215,7 +215,7 @@ This table shows the mapping between TensorFlow, TensorFlow Java and minimum sup
 | 1.0.0-rc.2              | 2.16.2             | 11                   |
 | 1.0.0                   | 2.16.2             | 11                   |
 | 1.1.0                   | 2.18.0             | 11                   |
-| 1.2.0-SNAPSHOT          | 2.20.0             | 11                   |
+| 1.2.0-SNAPSHOT          | 2.21.0             | 11                   |
 
 ## How to Contribute?
 

tensorflow-core/pom.xml

@@ -42,7 +42,7 @@
     https://github.com/tensorflow/tensorflow/blob/master/tensorflow/workspace2.bzl
     (but for the currently used release, not master)
     -->
-    <protobuf.version>4.28.3</protobuf.version>
+    <protobuf.version>4.31.1</protobuf.version>
 
     <native.classifier>${javacpp.platform}${javacpp.platform.extension}</native.classifier>
     <javacpp.platform.extension></javacpp.platform.extension>

tensorflow-core/tensorflow-core-api/scripts/test_download.sh

@@ -5,13 +5,13 @@ DOWNLOAD_FOLDER="$1"
 
 case ${PLATFORM:-} in
   'linux-x86_64')
-    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/c3/e6/cfd784298ffb759a4235721cac2ac20f7ff758bf687069cfbaebb06c5804/tensorflow_text-2.20.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl'
+    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/a9/b9/02707723d44e5d0fe5f7d27ba4237528bc654bac1b0f638efe37988584b1/tensorflow_text-2.20.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl'
     ;;
   'linux-arm64')
-    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/f5/ca/796cfd97ae6693d3c84c37575a6d481be5f1ef36c920d1a73c884f31797b/tensorflow_text-2.20.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl'
+    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/17/99/b397038628a660a1446bb79c8be284443a28500d072227a77ebaeb5cd149/tensorflow_text-2.20.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl'
     ;;
   'macosx-arm64')
-    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/98/e4/e3c72d0a73caeba90cf5b31e69d44e9a08d614e0e829484d813f3b63e037/tensorflow_text-2.20.0-cp312-cp312-macosx_11_0_arm64.whl'
+    TEXT_WHEEL_URL='https://files.pythonhosted.org/packages/61/b8/ccb0c3b048f268860f4c92f6bb7ab55c6af6c2fe4e26746c8dc384063915/tensorflow_text-2.20.1-cp313-cp313-macosx_11_0_arm64.whl'
     ;;
   *)
     echo "TensorFlow Text distribution for ${PLATFORM} is not supported for download"

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

@@ -694,7 +694,6 @@ public ThreadPoolDataset threadPoolDataset(Operand<? extends TType> inputDataset
    * @param numThreads The number of threads in the thread pool.
    * @param displayName A human-readable name for the threads that may be visible in some
    *  visualizations.
-   *  threadpool.
    * @param options carries optional attribute values
    * @return a new instance of ThreadPoolHandle
    */

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

@@ -2197,7 +2197,6 @@ public ThreadPoolDataset threadPoolDataset(Operand<? extends TType> inputDataset
    * @param numThreads The number of threads in the thread pool.
    * @param displayName A human-readable name for the threads that may be visible in some
    *  visualizations.
-   *  threadpool.
    * @param options carries optional attribute values
    * @return a new instance of ThreadPoolHandle
    */

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

@@ -337,10 +337,10 @@ public DecodeGif decodeGif(Operand<TString> contents) {
   }
 
   /**
-   * Function for decode_bmp, decode_gif, decode_jpeg, decode_webp, and decode_png.
-   *  Detects whether an image is a BMP, GIF, JPEG, WebP, or PNG, and performs the
-   *  appropriate operation to convert the input bytes string into a Tensor of type
-   *  dtype.
+   * Function for decode_bmp, decode_gif, decode_jpeg, decode_jxl, decode_webp, and decode_png.
+   *  Detects whether an image is a BMP, GIF, JPEG, JPEG XL, WebP, or PNG, and
+   *  performs the appropriate operation to convert the input bytes string into a
+   *  Tensor of type dtype.
    *  <p><em>NOTE</em>: decode_gif and decode_webp return a 4-D
    *  array [num_frames, height, width, 3], as opposed to decode_bmp,
    *  decode_jpeg, and decode_png, which always return 3-D arrays [height,
@@ -365,10 +365,10 @@ public DecodeImage<TUint8> decodeImage(Operand<TString> contents,
   }
 
   /**
-   * Function for decode_bmp, decode_gif, decode_jpeg, decode_webp, and decode_png.
-   *  Detects whether an image is a BMP, GIF, JPEG, WebP, or PNG, and performs the
-   *  appropriate operation to convert the input bytes string into a Tensor of type
-   *  dtype.
+   * Function for decode_bmp, decode_gif, decode_jpeg, decode_jxl, decode_webp, and decode_png.
+   *  Detects whether an image is a BMP, GIF, JPEG, JPEG XL, WebP, or PNG, and
+   *  performs the appropriate operation to convert the input bytes string into a
+   *  Tensor of type dtype.
    *  <p><em>NOTE</em>: decode_gif and decode_webp return a 4-D
    *  array [num_frames, height, width, 3], as opposed to decode_bmp,
    *  decode_jpeg, and decode_png, which always return 3-D arrays [height,

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

@@ -668,7 +668,8 @@ public ComplexAbs<TFloat32> complexAbs(Operand<? extends TType> x) {
    *  </blockquote>
    *
    * @param x The x value
-   * @param Tout The value of the Tout attribute
+   * @param Tout Need to be {@code tf.float32} when the type of {@code x} is {@code tf.complex64}.
+   *  Need to be {@code tf.float64} when the type of {@code x} is {@code tf.complex128}.
    * @param <U> data type for {@code ComplexAbs} output and operands
    * @return a new instance of ComplexAbs
    */
@@ -850,9 +851,9 @@ public <T extends TNumber> CumulativeLogsumexp<T> cumulativeLogsumexp(Operand<T>
    *
    * @param input 1D or 2D int {@code Tensor}.
    * @param sizeOutput non-negative int scalar {@code Tensor}.
-   * @param weights is an int32, int64, float32, or float64 {@code Tensor} with the same
-   *  shape as {@code arr}, or a length-0 {@code Tensor}, in which case it acts as all weights
-   *  equal to 1.
+   * @param weights {@code Tensor} with the same shape as {@code arr}, or a length-0 {@code Tensor},
+   *  in which case it acts as all weights equal to 1.
+   *  Not supported by the GPU implementation of Bincount.
    * @param options carries optional attribute values
    * @param <U> data type for {@code DenseBincount} output and operands
    * @param <T> data type for {@code DenseBincount} output and operands
@@ -1196,7 +1197,8 @@ public Imag<TFloat32> imag(Operand<? extends TType> input) {
    *  </pre>
    *
    * @param input The input value
-   * @param Tout The value of the Tout attribute
+   * @param Tout Need to be {@code tf.float32} when the type of {@code x} is {@code tf.complex64}.
+   *  Need to be {@code tf.float64} when the type of {@code x} is {@code tf.complex128}.
    * @param <U> data type for {@code Imag} output and operands
    * @return a new instance of Imag
    */

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

@@ -1753,7 +1753,15 @@ public IsotonicRegression<TFloat32> isotonicRegression(Operand<? extends TNumber
    * Solves a batch of isotonic regression problems.
    *
    * @param input A (batch_size, dim)-tensor holding a batch of inputs.
-   * @param outputDtype Dtype of output.
+   * @param outputDtype Dtype of the output tensor.
+   *  <p>Note on supported input-output type combinations:
+   *  <ul>
+   *  <li>For floating-point types, the output has the same dtype as the input.</li>
+   *  <li>For 8-bit and 16-bit integer inputs, the output is a 32-bit float.</li>
+   *  <li>For 32-bit and 64-bit integer inputs, the output is a 64-bit float.</li>
+   *  </ul>
+   *  <p>Using unsupported dtype pairs (for example, input=float64 with output=float32)
+   *  will result in a &quot;Could not find device for node&quot; error.
    * @param <U> data type for {@code IsotonicRegression} output and operands
    * @return a new instance of IsotonicRegression
    */

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

@@ -7743,12 +7743,12 @@ public <T extends TType> TensorListGather<T> tensorListGather(
    * @param index The index value
    * @param elementShape The elementShape value
    * @param elementDtype The value of the elementDtype attribute
-   * @param <T> data type for {@code TensorListGetItem} output and operands
+   * @param <U> data type for {@code TensorListGetItem} output and operands
    * @return a new instance of TensorListGetItem
    */
-  public <T extends TType> TensorListGetItem<T> tensorListGetItem(
-      Operand<? extends TType> inputHandle, Operand<TInt32> index, Operand<TInt32> elementShape,
-      Class<T> elementDtype) {
+  public <U extends TType> TensorListGetItem<U> tensorListGetItem(
+      Operand<? extends TType> inputHandle, Operand<TInt32> index,
+      Operand<? extends TNumber> elementShape, Class<U> elementDtype) {
     return TensorListGetItem.create(scope, inputHandle, index, elementShape, elementDtype);
   }
 

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

@@ -562,8 +562,10 @@ public <T extends TNumber> QuantizeAndDequantizeV4<T> quantizeAndDequantizeV4(Op
    *
    * @param gradients The gradients value
    * @param input The input value
-   * @param inputMin The inputMin value
-   * @param inputMax The inputMax value
+   * @param inputMin If {@code axis} is specified, the shape of the minimum input tensor
+   *  must be rank 1.
+   * @param inputMax If {@code axis} is specified, the shape of the maximum input tensor
+   *  must be rank 1.
    * @param options carries optional attribute values
    * @param <T> data type for {@code QuantizeAndDequantizeV4Grad} output and operands
    * @return a new instance of QuantizeAndDequantizeV4Grad

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

@@ -444,9 +444,10 @@ public <U extends TNumber> IrfftNd<U> irfftNd(Operand<? extends TType> input,
    *  corresponding dimension of {@code input}, the dimension is cropped. If it is larger,
    *  the dimension is padded with zeros.
    *
-   * @param input A float32 tensor.
+   * @param input A float32 or float64 tensor.
    * @param fftLength An int32 tensor of shape [1]. The FFT length.
-   * @param Tcomplex The value of the Tcomplex attribute
+   * @param Tcomplex Should be {@code tf.complex64} when the type of {@code input} is {@code float32}.
+   *  Should be {@code tf.complex128} when the type of {@code input} is {@code float64}.
    * @param <U> data type for {@code RFFT} output and operands
    * @return a new instance of Rfft
    */
@@ -467,9 +468,10 @@ public <U extends TType> Rfft<U> rfft(Operand<? extends TNumber> input, Operand<
    *  corresponding dimension of {@code input}, the dimension is cropped. If it is larger,
    *  the dimension is padded with zeros.
    *
-   * @param input A float32 tensor.
+   * @param input A float32 or float64 tensor.
    * @param fftLength An int32 tensor of shape [2]. The FFT length for each dimension.
-   * @param Tcomplex The value of the Tcomplex attribute
+   * @param Tcomplex Should be {@code tf.complex64} when the type of {@code input} is {@code float32}.
+   *  Should be {@code tf.complex128} when the type of {@code input} is {@code float64}.
    * @param <U> data type for {@code RFFT2D} output and operands
    * @return a new instance of Rfft2d
    */
@@ -490,9 +492,10 @@ public <U extends TType> Rfft2d<U> rfft2d(Operand<? extends TNumber> input,
    *  corresponding dimension of {@code input}, the dimension is cropped. If it is larger,
    *  the dimension is padded with zeros.
    *
-   * @param input A float32 tensor.
+   * @param input A float32 or float64 tensor.
    * @param fftLength An int32 tensor of shape [3]. The FFT length for each dimension.
-   * @param Tcomplex The value of the Tcomplex attribute
+   * @param Tcomplex Should be {@code tf.complex64} when the type of {@code input} is {@code float32}.
+   *  Should be {@code tf.complex128} when the type of {@code input} is {@code float64}.
    * @param <U> data type for {@code RFFT3D} output and operands
    * @return a new instance of Rfft3d
    */
@@ -513,10 +516,11 @@ public <U extends TType> Rfft3d<U> rfft3d(Operand<? extends TNumber> input,
    *  <p>Axes mean the dimensions to perform the transform on. Default is to perform on
    *  all axes.
    *
-   * @param input A complex tensor.
+   * @param input A float32 or float64 tensor.
    * @param fftLength An int32 tensor. The FFT length for each dimension.
    * @param axes An int32 tensor with a same shape as fft_length. Axes to perform the transform.
-   * @param Tcomplex The value of the Tcomplex attribute
+   * @param Tcomplex Should be {@code tf.complex64} when the type of {@code input} is {@code float32}.
+   *  Should be {@code tf.complex128} when the type of {@code input} is {@code float64}.
    * @param <U> data type for {@code RFFTND} output and operands
    * @return a new instance of RfftNd
    */

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core/TensorListGetItem.java

@@ -33,6 +33,7 @@
 import org.tensorflow.op.annotation.Operator;
 import org.tensorflow.proto.DataType;
 import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TNumber;
 import org.tensorflow.types.family.TType;
 
 /**
@@ -46,13 +47,13 @@
     inputsClass = TensorListGetItem.Inputs.class
 )
 @Operator
-public final class TensorListGetItem<T extends TType> extends RawOp implements Operand<T> {
+public final class TensorListGetItem<U extends TType> extends RawOp implements Operand<U> {
   /**
    * The name of this op, as known by TensorFlow core engine
    */
   public static final String OP_NAME = "TensorListGetItem";
 
-  private Output<T> item;
+  private Output<U> item;
 
   public TensorListGetItem(Operation operation) {
     super(operation, OP_NAME);
@@ -68,15 +69,15 @@ public TensorListGetItem(Operation operation) {
    * @param index The index value
    * @param elementShape The elementShape value
    * @param elementDtype The value of the elementDtype attribute
-   * @param <T> data type for {@code TensorListGetItem} output and operands
+   * @param <U> data type for {@code TensorListGetItem} output and operands
    * @return a new instance of TensorListGetItem
    */
   @Endpoint(
       describeByClass = true
   )
-  public static <T extends TType> TensorListGetItem<T> create(Scope scope,
-      Operand<? extends TType> inputHandle, Operand<TInt32> index, Operand<TInt32> elementShape,
-      Class<T> elementDtype) {
+  public static <U extends TType> TensorListGetItem<U> create(Scope scope,
+      Operand<? extends TType> inputHandle, Operand<TInt32> index,
+      Operand<? extends TNumber> elementShape, Class<U> elementDtype) {
     OperationBuilder opBuilder = scope.opBuilder(OP_NAME, "TensorListGetItem");
     opBuilder.addInput(inputHandle.asOutput());
     opBuilder.addInput(index.asOutput());
@@ -90,12 +91,12 @@ public static <T extends TType> TensorListGetItem<T> create(Scope scope,
    *
    * @return item.
    */
-  public Output<T> item() {
+  public Output<U> item() {
     return item;
   }
 
   @Override
-  public Output<T> asOutput() {
+  public Output<U> asOutput() {
     return item;
   }
 
@@ -116,20 +117,26 @@ public static class Inputs extends RawOpInputs<TensorListGetItem<?>> {
     /**
      * The elementShape input
      */
-    public final Operand<TInt32> elementShape;
+    public final Operand<? extends TNumber> elementShape;
 
     /**
      * The elementDtype attribute
      */
     public final DataType elementDtype;
 
+    /**
+     * The Tshape attribute
+     */
+    public final DataType Tshape;
+
     public Inputs(GraphOperation op) {
-      super(new TensorListGetItem<>(op), op, Arrays.asList("element_dtype"));
+      super(new TensorListGetItem<>(op), op, Arrays.asList("element_dtype", "Tshape"));
       int inputIndex = 0;
       inputHandle = (Operand<? extends TType>) op.input(inputIndex++);
       index = (Operand<TInt32>) op.input(inputIndex++);
-      elementShape = (Operand<TInt32>) op.input(inputIndex++);
+      elementShape = (Operand<? extends TNumber>) op.input(inputIndex++);
       elementDtype = op.attributes().getAttrType("element_dtype");
+      Tshape = op.attributes().getAttrType("Tshape");
     }
   }
 }

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/data/StatsAggregatorHandle.java

@@ -85,7 +85,8 @@ public static StatsAggregatorHandle create(Scope scope, Options... options) {
   /**
    * Sets the container option.
    *
-   * @param container the container option
+   * @param container The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+   * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
    * @return this Options instance.
    */
   public static Options container(String container) {
@@ -131,7 +132,8 @@ private Options() {
     /**
      * Sets the container option.
      *
-     * @param container the container option
+     * @param container The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+     * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
      * @return this Options instance.
      */
     public Options container(String container) {
@@ -156,7 +158,8 @@ public Options sharedName(String sharedName) {
   )
   public static class Inputs extends RawOpInputs<StatsAggregatorHandle> {
     /**
-     * The container attribute
+     * The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+     * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
      */
     public final String container;
 

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/data/ThreadPoolHandle.java

@@ -64,7 +64,6 @@ public ThreadPoolHandle(Operation operation) {
    * @param numThreads The number of threads in the thread pool.
    * @param displayName A human-readable name for the threads that may be visible in some
    * visualizations.
-   * threadpool.
    * @param options carries optional attribute values
    * @return a new instance of ThreadPoolHandle
    */
@@ -106,7 +105,8 @@ public static Options maxIntraOpParallelism(Long maxIntraOpParallelism) {
   /**
    * Sets the container option.
    *
-   * @param container the container option
+   * @param container The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+   * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
    * @return this Options instance.
    */
   public static Options container(String container) {
@@ -125,8 +125,7 @@ public static Options sharedName(String sharedName) {
 
   /**
    * Gets handle.
-   * A resource that can be consumed by one or more ExperimentalThreadPoolDataset
-   * ops.
+   *
    * @return handle.
    */
   public Output<? extends TType> handle() {
@@ -167,7 +166,8 @@ public Options maxIntraOpParallelism(Long maxIntraOpParallelism) {
     /**
      * Sets the container option.
      *
-     * @param container the container option
+     * @param container The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+     * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
      * @return this Options instance.
      */
     public Options container(String container) {
@@ -205,12 +205,12 @@ public static class Inputs extends RawOpInputs<ThreadPoolHandle> {
     /**
      * A human-readable name for the threads that may be visible in some
      * visualizations.
-     * threadpool.
      */
     public final String displayName;
 
     /**
-     * The container attribute
+     * The name of {@code container} should start with {@code '.'} or {@code letter} or {@code digit},
+     * with ['-', '.', '/'] or {@code letter} or {@code digit} follows several times.
      */
     public final String container;