[Type Refactor] Move type implementations and mappers to internal package

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/DataType.java

@@ -15,7 +15,6 @@
 
 package org.tensorflow;
 
-import org.tensorflow.internal.c_api.TF_Tensor;
 import org.tensorflow.types.TBfloat16;
 import org.tensorflow.types.TBool;
 import org.tensorflow.types.TFloat16;
@@ -30,23 +29,6 @@
 /** Represents a type of elements in a {@link Tensor} */
 public final class DataType<T extends TType> {
 
-  @FunctionalInterface
-  public interface TensorMapper<T> {
-
-    /**
-     * Maps the tensor memory to a n-dimensional typed data space.
-     *
-     * <p>This method is designed to be invoked internally by this library only, in order to pass the
-     * native handle of {@code tensor} as {@code nativeHandle} (and since only classes from the
-     * {@code org.tensorflow} package can retrieve such handle).
-     *
-     * @param tensor the tensor to map in its raw nature
-     * @param nativeHandle native handle of the tensor
-     * @return a typed tensor of type {@code T}
-     */
-    T apply(RawTensor tensor, TF_Tensor nativeHandle);
-  }
-
   /**
    * Creates a new datatype
    *
@@ -167,7 +149,7 @@ int nativeCode() {
    * @return data structure of elements of this type
    */
   T map(RawTensor tensor) {
-    return tensorMapper.apply(tensor, tensor.nativeHandle());
+    return tensorMapper.mapDense(tensor);
   }
 
   private final int nativeCode;

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/TensorMapper.java

@@ -0,0 +1,35 @@
+package org.tensorflow;
+
+import org.tensorflow.internal.c_api.TF_Tensor;
+import org.tensorflow.types.family.TType;
+
+/**
+ * Maps the native memory of a {@link RawTensor} to a n-dimensional typed data space
+ * accessible from the JVM.
+ *
+ * <p>Usage of this class is reserved for internal purposes only.
+ *
+ * @param <T> tensor type mapped by this object
+ * @see {@link TType}
+ */
+public abstract class TensorMapper<T extends TType> {
+
+  /**
+   * Maps the provided dense raw {@code tensor} as a tensor of type {@code T}.
+   *
+   * @param tensor the dense tensor to map, in its raw nature
+   * @return an instance of {@code T}
+   */
+  protected abstract T mapDense(RawTensor tensor);
+
+  /**
+   * Helper for retrieving the native handle of a raw tensor
+   *
+   * @param tensor a raw tensor
+   * @return the native handle of that tensor
+   * @throws IllegalStateException if the tensor has been released
+   */
+  protected static TF_Tensor nativeHandle(RawTensor tensor) {
+    return tensor.nativeHandle();
+  }
+}

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/buffer/ByteSequenceProvider.java

@@ -0,0 +1,53 @@
+package org.tensorflow.internal.buffer;
+
+import java.util.Iterator;
+import java.util.function.Function;
+import org.tensorflow.ndarray.NdArray;
+import org.tensorflow.ndarray.NdArraySequence;
+
+/**
+ * Produces sequence of bytes to be stored in a {@link ByteSequenceTensorBuffer}.
+ *
+ * @param <T> source of bytes (byte arrays or strings)
+ */
+public class ByteSequenceProvider<T> implements Iterable<byte[]> {
+
+  /**
+   * Constructor
+   *
+   * @param source source of data
+   * @param byteExtractor method that converts one value of the source into a sequence of bytes
+   */
+  public ByteSequenceProvider(NdArray<T> source, Function<T, byte[]> byteExtractor) {
+    this.source = source;
+    this.byteExtractor = byteExtractor;
+  }
+
+  @Override
+  public Iterator<byte[]> iterator() {
+    return new Iterator<byte[]>() {
+
+      @Override
+      public boolean hasNext() {
+        return scalarIterator.hasNext();
+      }
+
+      @Override
+      public byte[] next() {
+        return byteExtractor.apply(scalarIterator.next().getObject());
+      }
+
+      private final Iterator<? extends NdArray<T>> scalarIterator = source.scalars().iterator();
+    };
+  }
+
+  /**
+   * @return total number of byte sequences that can be produced by this sequencer
+   */
+  long numSequences() {
+    return source.size();
+  }
+
+  private final NdArray<T> source;
+  private final Function<T, byte[]> byteExtractor;
+}

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/buffer/StringTensorBuffer.java → tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/buffer/ByteSequenceTensorBuffer.java

@@ -50,23 +50,21 @@
  * <p>After its data has been initialized, the buffer is read-only as it is not possible to change
  * safely a value without reinitializing the whole data.
  */
-public class StringTensorBuffer extends AbstractDataBuffer<byte[]> {
+public class ByteSequenceTensorBuffer extends AbstractDataBuffer<byte[]> {
 
   /**
    * Computes how many bytes are required to store the given data in a string buffer.
    *
-   * @param data data to store eventually by calling {@link #init(NdArray, Function)}
-   * @param getBytes method that converts one value of the data into a sequence of bytes
+   * @param byteSequenceProvider produces sequences of bytes
    * @return number of bytes required to store the data.
    */
-  public static <T> long computeSize(NdArray<T> data, Function<T, byte[]> getBytes) {
+  public static <T> long computeSize(ByteSequenceProvider<?> byteSequenceProvider) {
     // reserve space to store 64-bit offsets
-    long size = data.size() * Long.BYTES;
+    long size = byteSequenceProvider.numSequences() * Long.BYTES;
 
     // reserve space to store length and data of each values
-    for (NdArray<T> scalar : data.scalars()) {
-      byte[] elementBytes = getBytes.apply(scalar.getObject());
-      size += elementBytes.length + StringTensorBuffer.varintLength(elementBytes.length);
+    for (byte[] elementBytes : byteSequenceProvider) {
+      size += elementBytes.length + ByteSequenceTensorBuffer.varintLength(elementBytes.length);
     }
     return size;
   }
@@ -79,14 +77,11 @@ public static <T> long computeSize(NdArray<T> data, Function<T, byte[]> getBytes
    * same set of data, calling {@link #computeSize(NdArray, Function)} priory to make sure there is
    * enough space to store it.
    *
-   * @param data data to store
-   * @param getBytes method that converts one value of the data into a sequence of bytes
+   * @param byteSequenceProvider produces sequences of bytes to use as the tensor data
    */
-  public <T> void init(NdArray<T> data, Function<T, byte[]> getBytes) {
+  public <T> void init(ByteSequenceProvider<T> byteSequenceProvider) {
     InitDataWriter writer = new InitDataWriter();
-    for (NdArray<T> scalar : data.scalars()) {
-      writer.writeNext(getBytes.apply(scalar.getObject()));
-    }
+    byteSequenceProvider.forEach(writer::writeNext);
   }
 
   @Override
@@ -129,8 +124,8 @@ public boolean isReadOnly() {
 
   @Override
   public DataBuffer<byte[]> copyTo(DataBuffer<byte[]> dst, long size) {
-    if (size == size() && dst instanceof StringTensorBuffer) {
-      StringTensorBuffer tensorDst = (StringTensorBuffer) dst;
+    if (size == size() && dst instanceof ByteSequenceTensorBuffer) {
+      ByteSequenceTensorBuffer tensorDst = (ByteSequenceTensorBuffer) dst;
       if (offsets.size() != size || data.size() != size) {
         throw new IllegalArgumentException(
             "Cannot copy string tensor data to another tensor of a different size");
@@ -145,20 +140,20 @@ public DataBuffer<byte[]> copyTo(DataBuffer<byte[]> dst, long size) {
 
   @Override
   public DataBuffer<byte[]> offset(long index) {
-    return new StringTensorBuffer(offsets.offset(index), data);
+    return new ByteSequenceTensorBuffer(offsets.offset(index), data);
   }
 
   @Override
   public DataBuffer<byte[]> narrow(long size) {
-    return new StringTensorBuffer(offsets.narrow(size), data);
+    return new ByteSequenceTensorBuffer(offsets.narrow(size), data);
   }
 
   @Override
   public DataBuffer<byte[]> slice(long index, long size) {
-    return new StringTensorBuffer(offsets.slice(index, size), data);
+    return new ByteSequenceTensorBuffer(offsets.slice(index, size), data);
   }
 
-  StringTensorBuffer(LongDataBuffer offsets, ByteDataBuffer data) {
+  ByteSequenceTensorBuffer(LongDataBuffer offsets, ByteDataBuffer data) {
     this.offsets = offsets;
     this.data = data;
   }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/buffer/TensorBuffers.java

@@ -156,7 +156,7 @@ public static BooleanDataBuffer toBooleans(TF_Tensor nativeTensor) {
    * @param nativeTensor native reference to the tensor
    * @return a string buffer
    */
-  public static StringTensorBuffer toStrings(TF_Tensor nativeTensor, long numElements) {
+  public static ByteSequenceTensorBuffer toStrings(TF_Tensor nativeTensor, long numElements) {
     Pointer tensorMemory = tensorMemory(nativeTensor);
     if (TensorRawDataBufferFactory.canBeUsed()) {
       return TensorRawDataBufferFactory.mapTensorToStrings(tensorMemory, numElements);
@@ -173,7 +173,7 @@ public static StringTensorBuffer toStrings(TF_Tensor nativeTensor, long numEleme
     dataBuffer.position((int)numElements * Long.BYTES);
     ByteDataBuffer data = DataBuffers.of(dataBuffer.slice());
 
-    return new StringTensorBuffer(offsets, data);
+    return new ByteSequenceTensorBuffer(offsets, data);
   }
 
   private static Pointer tensorMemory(TF_Tensor nativeTensor) {

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/buffer/TensorRawDataBufferFactory.java

@@ -57,13 +57,13 @@ static BooleanDataBuffer mapTensorToBooleans(Pointer tensorMemory) {
     return mapNativeBooleans(tensorMemory.address(), tensorMemory.capacity(), false);
   }
 
-  static StringTensorBuffer mapTensorToStrings(Pointer tensorMemory, long numElements) {
+  static ByteSequenceTensorBuffer mapTensorToStrings(Pointer tensorMemory, long numElements) {
     long offsetByteSize = numElements * Long.BYTES;
     LongDataBuffer offsets = mapNativeLongs(tensorMemory.address(), offsetByteSize, false);
     ByteDataBuffer data = mapNativeBytes(
         tensorMemory.address() + offsetByteSize,
         tensorMemory.capacity() - offsetByteSize,
         false);
-    return new StringTensorBuffer(offsets, data);
+    return new ByteSequenceTensorBuffer(offsets, data);
   }
 }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/types/TBfloat16Mapper.java

@@ -0,0 +1,43 @@
+package org.tensorflow.internal.types;
+
+import org.tensorflow.TensorMapper;
+import org.tensorflow.DataType;
+import org.tensorflow.RawTensor;
+import org.tensorflow.internal.buffer.TensorBuffers;
+import org.tensorflow.ndarray.buffer.FloatDataBuffer;
+import org.tensorflow.ndarray.buffer.layout.DataLayouts;
+import org.tensorflow.ndarray.impl.dense.FloatDenseNdArray;
+import org.tensorflow.types.TBfloat16;
+
+/**
+ * Maps memory of {@link org.tensorflow.proto.framework.DataType#DT_BFLOAT16} tensors
+ * to a n-dimensional data space.
+ */
+public final class TBfloat16Mapper extends TensorMapper<TBfloat16> {
+
+  @Override
+  protected TBfloat16 mapDense(RawTensor tensor) {
+    FloatDataBuffer buffer = DataLayouts.BFLOAT16.applyTo(TensorBuffers.toShorts(nativeHandle(tensor)));
+    return new DenseTBfloat16(tensor, buffer);
+  }
+
+  private static final class DenseTBfloat16 extends FloatDenseNdArray implements TBfloat16 {
+
+    @Override
+    public DataType<TBfloat16> dataType() {
+      return TBfloat16.DTYPE;
+    }
+
+    @Override
+    public RawTensor asRawTensor() {
+      return rawTensor;
+    }
+
+    final RawTensor rawTensor;
+
+    DenseTBfloat16(RawTensor rawTensor, FloatDataBuffer buffer) {
+      super(buffer, rawTensor.shape());
+      this.rawTensor = rawTensor;
+    }
+  }
+}

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/types/TBoolMapper.java

@@ -0,0 +1,42 @@
+package org.tensorflow.internal.types;
+
+import org.tensorflow.DataType;
+import org.tensorflow.RawTensor;
+import org.tensorflow.TensorMapper;
+import org.tensorflow.internal.buffer.TensorBuffers;
+import org.tensorflow.ndarray.buffer.BooleanDataBuffer;
+import org.tensorflow.ndarray.impl.dense.BooleanDenseNdArray;
+import org.tensorflow.types.TBool;
+
+/**
+ * Maps memory of {@link org.tensorflow.proto.framework.DataType#DT_BOOL} tensors
+ * to a n-dimensional data space.
+ */
+public final class TBoolMapper extends TensorMapper<TBool> {
+
+  @Override
+  protected TBool mapDense(RawTensor tensor) {
+    BooleanDataBuffer buffer = TensorBuffers.toBooleans(nativeHandle(tensor));
+    return new DenseTBool(tensor, buffer);
+  }
+
+  private static final class DenseTBool extends BooleanDenseNdArray implements TBool {
+
+    @Override
+    public DataType<TBool> dataType() {
+      return TBool.DTYPE;
+    }
+
+    @Override
+    public RawTensor asRawTensor() {
+      return rawTensor;
+    }
+
+    final RawTensor rawTensor;
+
+    DenseTBool(RawTensor rawTensor, BooleanDataBuffer buffer) {
+      super(buffer, rawTensor.shape());
+      this.rawTensor = rawTensor;
+    }
+  }
+}

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/types/TFloat16Mapper.java

@@ -0,0 +1,43 @@
+package org.tensorflow.internal.types;
+
+import org.tensorflow.DataType;
+import org.tensorflow.RawTensor;
+import org.tensorflow.TensorMapper;
+import org.tensorflow.internal.buffer.TensorBuffers;
+import org.tensorflow.ndarray.buffer.FloatDataBuffer;
+import org.tensorflow.ndarray.buffer.layout.DataLayouts;
+import org.tensorflow.ndarray.impl.dense.FloatDenseNdArray;
+import org.tensorflow.types.TFloat16;
+
+/**
+ * Maps memory of {@link org.tensorflow.proto.framework.DataType#DT_HALF} tensors
+ * to a n-dimensional data space.
+ */
+public final class TFloat16Mapper extends TensorMapper<TFloat16> {
+
+  @Override
+  protected TFloat16 mapDense(RawTensor tensor) {
+    FloatDataBuffer buffer = DataLayouts.FLOAT16.applyTo(TensorBuffers.toShorts(nativeHandle(tensor)));
+    return new DenseTFloat16(tensor, buffer);
+  }
+
+  private static final class DenseTFloat16 extends FloatDenseNdArray implements TFloat16 {
+
+    @Override
+    public DataType<?> dataType() {
+      return TFloat16.DTYPE;
+    }
+
+    @Override
+    public RawTensor asRawTensor() {
+      return rawTensor;
+    }
+
+    final RawTensor rawTensor;
+
+    DenseTFloat16(RawTensor rawTensor, FloatDataBuffer buffer) {
+      super(buffer, rawTensor.shape());
+      this.rawTensor = rawTensor;
+    }
+  }
+}