xref: /external/tensorflow/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Tensor.java

/* Copyright 2017 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.lite;

import java.lang.reflect.Array;
import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.nio.LongBuffer;
import java.nio.ShortBuffer;
import java.util.Arrays;

/**
 * A typed multi-dimensional array used in Tensorflow Lite.
 *
 * <p>The native handle of a {@code Tensor} is managed by {@code NativeInterpreterWrapper}, and does
 * not needed to be closed by the client. However, once the {@code NativeInterpreterWrapper} has
 * been closed, the tensor handle will be invalidated.
 */
// TODO(b/153882978): Add scalar getters similar to TF's Java API.
public final class Tensor {

  /**
   * Creates a Tensor wrapper from the provided interpreter instance and tensor index.
   *
   * <p>The caller is responsible for closing the created wrapper, and ensuring the provided native
   * interpreter is valid until the tensor is closed.
   */
  static Tensor fromIndex(long nativeInterpreterHandle, int tensorIndex) {
    return new Tensor(create(nativeInterpreterHandle, tensorIndex, /*subgraphIndex=*/ 0));
  }

  /**
   * Creates a Tensor wrapper from the provided interpreter instance, subgraph and tensor indices.
   *
   * <p>The caller is responsible for closing the created wrapper, and ensuring the provided native
   * interpreter is valid until the tensor is closed.
   */
  static Tensor fromSubgraphAndIndex(
      long nativeInterpreterHandle, int subgraphIndex, int tensorIndex) {
    return new Tensor(create(nativeInterpreterHandle, tensorIndex, subgraphIndex));
  }

  /**
   * Quantization parameters that corresponds to the table, {@code QuantizationParameters}, in the
   * <a
   * href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs">TFLite
   * Model schema file.</a>
   *
   * <p>Since per-channel quantization does not apply to input and output tensors, {@code scale} and
   * {@code zero_point} are both single values instead of arrays.
   *
   * <p>For tensor that are not quantized, the values of scale and zero_point are both 0.
   *
   * <p>Given a quantized value q, the corresponding float value f should be: <br>
   * f = scale * (q - zero_point) <br>
   */
  public static class QuantizationParams {
    /** The scale value used in quantization. */
    private final float scale;
    /** The zero point value used in quantization. */
    private final int zeroPoint;

    /**
     * Creates a {@link QuantizationParams} with {@code scale} and {@code zero_point}.
     *
     * @param scale The scale value used in quantization.
     * @param zeroPoint The zero point value used in quantization.
     */
    public QuantizationParams(final float scale, final int zeroPoint) {
      this.scale = scale;
      this.zeroPoint = zeroPoint;
    }

    /** Returns the scale value. */
    public float getScale() {
      return scale;
    }

    /** Returns the zero point value. */
    public int getZeroPoint() {
      return zeroPoint;
    }
  }

  /** Disposes of any resources used by the Tensor wrapper. */
  void close() {
    delete(nativeHandle);
    nativeHandle = 0;
  }

  /** Returns the {@link DataType} of elements stored in the Tensor. */
  public DataType dataType() {
    return dtype;
  }

  /**
   * Returns the number of dimensions (sometimes referred to as <a
   * href="https://www.tensorflow.org/resources/dims_types.html#rank">rank</a>) of the Tensor.
   *
   * <p>Will be 0 for a scalar, 1 for a vector, 2 for a matrix, 3 for a 3-dimensional tensor etc.
   */
  public int numDimensions() {
    return shapeCopy.length;
  }

  /** Returns the size, in bytes, of the tensor data. */
  public int numBytes() {
    return numBytes(nativeHandle);
  }

  /** Returns the number of elements in a flattened (1-D) view of the tensor. */
  public int numElements() {
    return computeNumElements(shapeCopy);
  }

  /**
   * Returns the <a href="https://www.tensorflow.org/resources/dims_types.html#shape">shape</a> of
   * the Tensor, i.e., the sizes of each dimension.
   *
   * @return an array where the i-th element is the size of the i-th dimension of the tensor.
   */
  public int[] shape() {
    return shapeCopy;
  }

  /**
   * Returns the original <a
   * href="https://www.tensorflow.org/resources/dims_types.html#shape">shape</a> of the Tensor,
   * i.e., the sizes of each dimension - before any resizing was performed. Unknown dimensions are
   * designated with a value of -1.
   *
   * @return an array where the i-th element is the size of the i-th dimension of the tensor.
   */
  public int[] shapeSignature() {
    return shapeSignatureCopy;
  }

  /**
   * Returns the (global) index of the tensor within the subgraph of the owning {@link Interpreter}.
   *
   * @hide
   */
  public int index() {
    return index(nativeHandle);
  }

  /**
   * Returns the subgraph index of the tensor within the owning {@link Interpreter}.
   *
   * @hide
   */
  public int subgraph() {
    return subgraph(nativeHandle);
  }

  /**
   * Returns the name of the tensor within the owning {@link Interpreter}.
   *
   * @hide
   */
  public String name() {
    return name(nativeHandle);
  }

  /**
   * Returns the quantization parameters of the tensor within the owning {@link Interpreter}.
   *
   * <p>Only quantized tensors have valid {@code QuantizationParameters}. For tensor that are not
   * quantized, the values of scale and zero_point are both 0.
   */
  public QuantizationParams quantizationParams() {
    return quantizationParamsCopy;
  }

  /**
   * Returns a read-only {@code ByteBuffer} view of the tensor data.
   *
   * <p>In general, this method is most useful for obtaining a read-only view of output tensor data,
   * *after* inference has been executed (e.g., via {@link Interpreter#run(Object,Object)}). In
   * particular, some graphs have dynamically shaped outputs, which can make feeding a predefined
   * output buffer to the {@link Interpreter} awkward. Example usage:
   *
   * <pre>{@code
   * interpreter.run(input, null);
   * ByteBuffer outputBuffer = interpreter.getOutputTensor(0).asReadOnlyBuffer();
   * // Copy or read from outputBuffer.
   * }</pre>
   *
   * <p>WARNING: If the tensor has not yet been allocated, e.g., before inference has been executed,
   * the result is undefined. Note that the underlying tensor pointer may also change when the
   * tensor is invalidated in any way (e.g., if inference is executed, or the graph is resized), so
   * it is *not* safe to hold a reference to the returned buffer beyond immediate use directly
   * following inference. Example *bad* usage:
   *
   * <pre>{@code
   * ByteBuffer outputBuffer = interpreter.getOutputTensor(0).asReadOnlyBuffer();
   * interpreter.run(input, null);
   * // Copy or read from outputBuffer (which may now be invalid).
   * }</pre>
   *
   * <p>WARNING: This is an experimental interface that is subject to change.
   *
   * @throws IllegalArgumentException if the tensor data has not been allocated.
   */
  public ByteBuffer asReadOnlyBuffer() {
    // Note that the ByteBuffer order is not preserved when duplicated or marked read only, so
    // we have to repeat the call.
    return buffer().asReadOnlyBuffer().order(ByteOrder.nativeOrder());
  }

  /**
   * Copies the contents of the provided {@code src} object to the Tensor.
   *
   * <p>The {@code src} should either be a (multi-dimensional) array with a shape matching that of
   * this tensor, a {@link ByteByffer} of compatible primitive type with a matching flat size, or
   * {@code null} iff the tensor has an underlying delegate buffer handle.
   *
   * @throws IllegalArgumentException if the tensor is a scalar or if {@code src} is not compatible
   *     with the tensor (for example, mismatched data types or shapes).
   */
  void setTo(Object src) {
    if (src == null) {
      if (hasDelegateBufferHandle(nativeHandle)) {
        return;
      }
      throw new IllegalArgumentException(
          "Null inputs are allowed only if the Tensor is bound to a buffer handle.");
    }
    throwIfTypeIsIncompatible(src);
    throwIfSrcShapeIsIncompatible(src);
    if (isBuffer(src)) {
      setTo((Buffer) src);
    } else if (dtype == DataType.STRING && shapeCopy.length == 0) {
      // Update scalar string input with 1-d byte array.
      writeScalar(nativeHandle, src);
    } else if (src.getClass().isArray()) {
      writeMultiDimensionalArray(nativeHandle, src);
    } else {
      writeScalar(nativeHandle, src);
    }
  }

  private void setTo(Buffer src) {
    // Note that we attempt to use a direct memcpy optimization for direct, native-ordered buffers.
    // There are no base Buffer#order() or Buffer#put() methods, so again we have to ugly cast.
    if (src instanceof ByteBuffer) {
      ByteBuffer srcBuffer = (ByteBuffer) src;
      if (srcBuffer.isDirect() && srcBuffer.order() == ByteOrder.nativeOrder()) {
        writeDirectBuffer(nativeHandle, src);
      } else {
        buffer().put(srcBuffer);
      }
    } else if (src instanceof LongBuffer) {
      LongBuffer srcBuffer = (LongBuffer) src;
      if (srcBuffer.isDirect() && srcBuffer.order() == ByteOrder.nativeOrder()) {
        writeDirectBuffer(nativeHandle, src);
      } else {
        buffer().asLongBuffer().put(srcBuffer);
      }
    } else if (src instanceof FloatBuffer) {
      FloatBuffer srcBuffer = (FloatBuffer) src;
      if (srcBuffer.isDirect() && srcBuffer.order() == ByteOrder.nativeOrder()) {
        writeDirectBuffer(nativeHandle, src);
      } else {
        buffer().asFloatBuffer().put(srcBuffer);
      }
    } else if (src instanceof IntBuffer) {
      IntBuffer srcBuffer = (IntBuffer) src;
      if (srcBuffer.isDirect() && srcBuffer.order() == ByteOrder.nativeOrder()) {
        writeDirectBuffer(nativeHandle, src);
      } else {
        buffer().asIntBuffer().put(srcBuffer);
      }
    } else if (src instanceof ShortBuffer) {
      ShortBuffer srcBuffer = (ShortBuffer) src;
      if (srcBuffer.isDirect() && srcBuffer.order() == ByteOrder.nativeOrder()) {
        writeDirectBuffer(nativeHandle, src);
      } else {
        buffer().asShortBuffer().put(srcBuffer);
      }
    } else {
      throw new IllegalArgumentException("Unexpected input buffer type: " + src);
    }
  }

  /**
   * Copies the contents of the tensor to {@code dst} and returns {@code dst}.
   *
   * @param dst the destination buffer, either an explicitly-typed array, a compatible {@link
   *     Buffer} or {@code null} iff the tensor has an underlying delegate buffer handle. If
   *     providing a (multi-dimensional) array, its shape must match the tensor shape *exactly*. If
   *     providing a {@link Buffer}, its capacity must be at least as large as the source tensor's
   *     capacity.
   * @throws IllegalArgumentException if {@code dst} is not compatible with the tensor (for example,
   *     mismatched data types or shapes).
   */
  Object copyTo(Object dst) {
    if (dst == null) {
      if (hasDelegateBufferHandle(nativeHandle)) {
        return dst;
      }
      throw new IllegalArgumentException(
          "Null outputs are allowed only if the Tensor is bound to a buffer handle.");
    }
    throwIfTypeIsIncompatible(dst);
    throwIfDstShapeIsIncompatible(dst);
    if (isBuffer(dst)) {
      copyTo((Buffer) dst);
    } else {
      readMultiDimensionalArray(nativeHandle, dst);
    }
    return dst;
  }

  private void copyTo(Buffer dst) {
    // There is no base Buffer#put() method, so we have to ugly cast.
    if (dst instanceof ByteBuffer) {
      ((ByteBuffer) dst).put(buffer());
    } else if (dst instanceof FloatBuffer) {
      ((FloatBuffer) dst).put(buffer().asFloatBuffer());
    } else if (dst instanceof LongBuffer) {
      ((LongBuffer) dst).put(buffer().asLongBuffer());
    } else if (dst instanceof IntBuffer) {
      ((IntBuffer) dst).put(buffer().asIntBuffer());
    } else if (dst instanceof ShortBuffer) {
      ((ShortBuffer) dst).put(buffer().asShortBuffer());
    } else {
      throw new IllegalArgumentException("Unexpected output buffer type: " + dst);
    }
  }

  /** Returns the provided buffer's shape if specified and different from this Tensor's shape. */
  // TODO(b/80431971): Remove this method after deprecating multi-dimensional array inputs.
  int[] getInputShapeIfDifferent(Object input) {
    if (input == null) {
      return null;
    }
    // Implicit resizes based on ByteBuffer capacity isn't supported, so short-circuit that path.
    // The Buffer's size will be validated against this Tensor's size in {@link #setTo(Object)}.
    if (isBuffer(input)) {
      return null;
    }
    throwIfTypeIsIncompatible(input);
    int[] inputShape = computeShapeOf(input);
    if (Arrays.equals(shapeCopy, inputShape)) {
      return null;
    }
    return inputShape;
  }

  /**
   * Forces a refresh of the tensor's cached shape.
   *
   * <p>This is useful if the tensor is resized or has a dynamic shape.
   */
  void refreshShape() {
    this.shapeCopy = shape(nativeHandle);
  }

  /** Returns the type of the data. */
  DataType dataTypeOf(Object o) {
    if (o != null) {
      Class<?> c = o.getClass();
      // For arrays, the data elements must be a *primitive* type, e.g., an
      // array of floats is fine, but not an array of Floats.
      if (c.isArray()) {
        while (c.isArray()) {
          c = c.getComponentType();
        }
        if (float.class.equals(c)) {
          return DataType.FLOAT32;
        } else if (int.class.equals(c)) {
          return DataType.INT32;
        } else if (short.class.equals(c)) {
          return DataType.INT16;
        } else if (byte.class.equals(c)) {
          // Byte array can be used for storing string tensors, especially for ParseExample op.
          if (dtype == DataType.STRING) {
            return DataType.STRING;
          }
          return DataType.UINT8;
        } else if (long.class.equals(c)) {
          return DataType.INT64;
        } else if (boolean.class.equals(c)) {
          return DataType.BOOL;
        } else if (String.class.equals(c)) {
          return DataType.STRING;
        }
      } else {
        // For scalars, the type will be boxed.
        if (Float.class.equals(c) || o instanceof FloatBuffer) {
          return DataType.FLOAT32;
        } else if (Integer.class.equals(c) || o instanceof IntBuffer) {
          return DataType.INT32;
        } else if (Short.class.equals(c) || o instanceof ShortBuffer) {
          return DataType.INT16;
        } else if (Byte.class.equals(c)) {
          // Note that we don't check for ByteBuffer here; ByteBuffer payloads
          // are allowed to map to any type, and should be handled earlier
          // in the input/output processing pipeline.
          return DataType.UINT8;
        } else if (Long.class.equals(c) || o instanceof LongBuffer) {
          return DataType.INT64;
        } else if (Boolean.class.equals(c)) {
          return DataType.BOOL;
        } else if (String.class.equals(c)) {
          return DataType.STRING;
        }
      }
    }
    throw new IllegalArgumentException(
        "DataType error: cannot resolve DataType of " + o.getClass().getName());
  }

  /** Returns the shape of an object as an int array. */
  int[] computeShapeOf(Object o) {
    int size = computeNumDimensions(o);
    if (dtype == DataType.STRING) {
      Class<?> c = o.getClass();
      if (c.isArray()) {
        while (c.isArray()) {
          c = c.getComponentType();
        }
        // If the given string data is stored in byte streams, the last array dimension should be
        // treated as a value.
        if (byte.class.equals(c)) {
          --size;
        }
      }
    }
    int[] dimensions = new int[size];
    fillShape(o, 0, dimensions);
    return dimensions;
  }

  /** Returns the number of elements in a flattened (1-D) view of the tensor's shape. */
  static int computeNumElements(int[] shape) {
    int n = 1;
    for (int i = 0; i < shape.length; ++i) {
      n *= shape[i];
    }
    return n;
  }

  /** Returns the number of dimensions of a multi-dimensional array, otherwise 0. */
  static int computeNumDimensions(Object o) {
    if (o == null || !o.getClass().isArray()) {
      return 0;
    }
    if (Array.getLength(o) == 0) {
      throw new IllegalArgumentException("Array lengths cannot be 0.");
    }
    return 1 + computeNumDimensions(Array.get(o, 0));
  }

  /** Recursively populates the shape dimensions for a given (multi-dimensional) array. */
  static void fillShape(Object o, int dim, int[] shape) {
    if (shape == null || dim == shape.length) {
      return;
    }
    final int len = Array.getLength(o);
    if (shape[dim] == 0) {
      shape[dim] = len;
    } else if (shape[dim] != len) {
      throw new IllegalArgumentException(
          String.format("Mismatched lengths (%d and %d) in dimension %d", shape[dim], len, dim));
    }
    for (int i = 0; i < len; ++i) {
      fillShape(Array.get(o, i), dim + 1, shape);
    }
  }

  private void throwIfTypeIsIncompatible(Object o) {
    // ByteBuffer payloads can map to any type, so exempt it from the check.
    if (isByteBuffer(o)) {
      return;
    }
    DataType oType = dataTypeOf(o);

    if (oType != dtype) {
      // INT8 and UINT8 have the same string name, "byte"
      if (oType.toStringName().equals(dtype.toStringName())) {
        return;
      }

      throw new IllegalArgumentException(
          String.format(
              "Cannot convert between a TensorFlowLite tensor with type %s and a Java "
                  + "object of type %s (which is compatible with the TensorFlowLite type %s).",
              dtype, o.getClass().getName(), oType));
    }
  }

  private void throwIfSrcShapeIsIncompatible(Object src) {
    if (isBuffer(src)) {
      Buffer srcBuffer = (Buffer) src;
      int bytes = numBytes();
      // Note that we allow the client to provide a ByteBuffer even for non-byte Tensors.
      // In such cases, we only care that the raw byte capacity matches the tensor byte capacity.
      int srcBytes =
          isByteBuffer(src) ? srcBuffer.capacity() : srcBuffer.capacity() * dtype.byteSize();
      if (bytes != srcBytes) {
        throw new IllegalArgumentException(
            String.format(
                "Cannot copy to a TensorFlowLite tensor (%s) with %d bytes from a "
                    + "Java Buffer with %d bytes.",
                name(), bytes, srcBytes));
      }
      return;
    }
    int[] srcShape = computeShapeOf(src);
    if (!Arrays.equals(srcShape, shapeCopy)) {
      throw new IllegalArgumentException(
          String.format(
              "Cannot copy to a TensorFlowLite tensor (%s) with shape %s from a Java object "
                  + "with shape %s.",
              name(), Arrays.toString(shapeCopy), Arrays.toString(srcShape)));
    }
  }

  private void throwIfDstShapeIsIncompatible(Object dst) {
    if (isBuffer(dst)) {
      Buffer dstBuffer = (Buffer) dst;
      int bytes = numBytes();
      // Note that we allow the client to provide a ByteBuffer even for non-byte Tensors.
      // In such cases, we only care that the raw byte capacity fits the tensor byte capacity.
      // This is subtly different than Buffer *inputs*, where the size should be exact.
      int dstBytes =
          isByteBuffer(dst) ? dstBuffer.capacity() : dstBuffer.capacity() * dtype.byteSize();
      if (bytes > dstBytes) {
        throw new IllegalArgumentException(
            String.format(
                "Cannot copy from a TensorFlowLite tensor (%s) with %d bytes to a "
                    + "Java Buffer with %d bytes.",
                name(), bytes, dstBytes));
      }
      return;
    }
    int[] dstShape = computeShapeOf(dst);
    if (!Arrays.equals(dstShape, shapeCopy)) {
      throw new IllegalArgumentException(
          String.format(
              "Cannot copy from a TensorFlowLite tensor (%s) with shape %s to a Java object "
                  + "with shape %s.",
              name(), Arrays.toString(shapeCopy), Arrays.toString(dstShape)));
    }
  }

  private static boolean isBuffer(Object o) {
    return o instanceof Buffer;
  }

  private static boolean isByteBuffer(Object o) {
    return o instanceof ByteBuffer;
  }

  private long nativeHandle;
  private final DataType dtype;
  private int[] shapeCopy;
  private final int[] shapeSignatureCopy;
  private final QuantizationParams quantizationParamsCopy;

  private Tensor(long nativeHandle) {
    this.nativeHandle = nativeHandle;
    this.dtype = DataType.fromC(dtype(nativeHandle));
    this.shapeCopy = shape(nativeHandle);
    this.shapeSignatureCopy = shapeSignature(nativeHandle);
    this.quantizationParamsCopy =
        new QuantizationParams(
            quantizationScale(nativeHandle), quantizationZeroPoint(nativeHandle));
  }

  private ByteBuffer buffer() {
    return buffer(nativeHandle).order(ByteOrder.nativeOrder());
  }

  private static native long create(long interpreterHandle, int tensorIndex, int subgraphIndex);

  private static native void delete(long handle);

  private static native ByteBuffer buffer(long handle);

  private static native void writeDirectBuffer(long handle, Buffer src);

  private static native int dtype(long handle);

  private static native int[] shape(long handle);

  private static native int[] shapeSignature(long handle);

  private static native int numBytes(long handle);

  private static native boolean hasDelegateBufferHandle(long handle);

  private static native void readMultiDimensionalArray(long handle, Object dst);

  private static native void writeMultiDimensionalArray(long handle, Object src);

  private static native void writeScalar(long handle, Object src);

  private static native int index(long handle);

  private static native int subgraph(long handle);

  private static native String name(long handle);

  private static native float quantizationScale(long handle);

  private static native int quantizationZeroPoint(long handle);
}