Java Tensor and EValue serialization (#6620)

Add serialization and deserialization for EValue (except string) and
Tensor.

RFC: #6569

Author: kirklandsign

extension/android/src/main/java/org/pytorch/executorch/EValue.java

@@ -9,6 +9,8 @@
 package org.pytorch.executorch;
 
 import com.facebook.jni.annotations.DoNotStrip;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Locale;
 import java.util.Optional;
 import org.pytorch.executorch.annotations.Experimental;
@@ -287,4 +289,75 @@ private void preconditionType(int typeCodeExpected, int typeCode) {
   private String getTypeName(int typeCode) {
     return typeCode >= 0 && typeCode < TYPE_NAMES.length ? TYPE_NAMES[typeCode] : "Unknown";
   }
+
+  /**
+   * Serializes an {@code EValue} into a byte array.
+   *
+   * @return The serialized byte array.
+   * @apiNote This method is experimental and subject to change without notice. This does NOT
+   *     supoprt list type.
+   */
+  public byte[] toByteArray() {
+    if (isNone()) {
+      return ByteBuffer.allocate(1).put((byte) TYPE_CODE_NONE).array();
+    } else if (isTensor()) {
+      Tensor t = toTensor();
+      byte[] tByteArray = t.toByteArray();
+      return ByteBuffer.allocate(1 + tByteArray.length)
+          .put((byte) TYPE_CODE_TENSOR)
+          .put(tByteArray)
+          .array();
+    } else if (isBool()) {
+      return ByteBuffer.allocate(2)
+          .put((byte) TYPE_CODE_BOOL)
+          .put((byte) (toBool() ? 1 : 0))
+          .array();
+    } else if (isInt()) {
+      return ByteBuffer.allocate(9).put((byte) TYPE_CODE_INT).putLong(toInt()).array();
+    } else if (isDouble()) {
+      return ByteBuffer.allocate(9).put((byte) TYPE_CODE_DOUBLE).putDouble(toDouble()).array();
+    } else if (isString()) {
+      return ByteBuffer.allocate(1 + toString().length())
+          .put((byte) TYPE_CODE_STRING)
+          .put(toString().getBytes())
+          .array();
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+  }
+
+  /**
+   * Deserializes an {@code EValue} from a byte[].
+   *
+   * @param bytes The byte array to deserialize from.
+   * @return The deserialized {@code EValue}.
+   * @apiNote This method is experimental and subject to change without notice. This does NOT list
+   *     type.
+   */
+  public static EValue fromByteArray(byte[] bytes) {
+    ByteBuffer buffer = ByteBuffer.wrap(bytes);
+    if (buffer == null) {
+      throw new IllegalArgumentException("buffer cannot be null");
+    }
+    if (!buffer.hasRemaining()) {
+      throw new IllegalArgumentException("invalid buffer");
+    }
+    int typeCode = buffer.get();
+    switch (typeCode) {
+      case TYPE_CODE_NONE:
+        return new EValue(TYPE_CODE_NONE);
+      case TYPE_CODE_TENSOR:
+        byte[] bufferArray = buffer.array();
+        return from(Tensor.fromByteArray(Arrays.copyOfRange(bufferArray, 1, bufferArray.length)));
+      case TYPE_CODE_STRING:
+        throw new IllegalArgumentException("TYPE_CODE_STRING is not supported");
+      case TYPE_CODE_DOUBLE:
+        return from(buffer.getDouble());
+      case TYPE_CODE_INT:
+        return from(buffer.getLong());
+      case TYPE_CODE_BOOL:
+        return from(buffer.get() != 0);
+    }
+    throw new IllegalArgumentException("invalid type code: " + typeCode);
+  }
 }

extension/android/src/main/java/org/pytorch/executorch/Tensor.java

@@ -53,9 +53,10 @@ public abstract class Tensor {
 
   @DoNotStrip final long[] shape;
 
+  private static final int BYTE_SIZE_BYTES = 1;
   private static final int INT_SIZE_BYTES = 4;
-  private static final int FLOAT_SIZE_BYTES = 4;
   private static final int LONG_SIZE_BYTES = 8;
+  private static final int FLOAT_SIZE_BYTES = 4;
   private static final int DOUBLE_SIZE_BYTES = 8;
 
   /**
@@ -679,4 +680,105 @@ private static Tensor nativeNewTensor(
     tensor.mHybridData = hybridData;
     return tensor;
   }
+
+  /**
+   * Serializes a {@code Tensor} into a byte array.
+   *
+   * @return The serialized byte array.
+   * @apiNote This method is experimental and subject to change without notice. This does NOT
+   *     supoprt list type.
+   */
+  public byte[] toByteArray() {
+    int dtypeSize = 0;
+    byte[] tensorAsByteArray = null;
+    if (dtype() == DType.UINT8) {
+      dtypeSize = BYTE_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel()];
+      Tensor_uint8 thiz = (Tensor_uint8) this;
+      ByteBuffer.wrap(tensorAsByteArray).put(thiz.getDataAsUnsignedByteArray());
+    } else if (dtype() == DType.INT8) {
+      dtypeSize = BYTE_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel()];
+      Tensor_int8 thiz = (Tensor_int8) this;
+      ByteBuffer.wrap(tensorAsByteArray).put(thiz.getDataAsByteArray());
+    } else if (dtype() == DType.INT16) {
+      throw new IllegalArgumentException("DType.INT16 is not supported in Java so far");
+    } else if (dtype() == DType.INT32) {
+      dtypeSize = INT_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel() * dtypeSize];
+      Tensor_int32 thiz = (Tensor_int32) this;
+      ByteBuffer.wrap(tensorAsByteArray).asIntBuffer().put(thiz.getDataAsIntArray());
+    } else if (dtype() == DType.INT64) {
+      dtypeSize = LONG_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel() * dtypeSize];
+      Tensor_int64 thiz = (Tensor_int64) this;
+      ByteBuffer.wrap(tensorAsByteArray).asLongBuffer().put(thiz.getDataAsLongArray());
+    } else if (dtype() == DType.FLOAT) {
+      dtypeSize = FLOAT_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel() * dtypeSize];
+      Tensor_float32 thiz = (Tensor_float32) this;
+      ByteBuffer.wrap(tensorAsByteArray).asFloatBuffer().put(thiz.getDataAsFloatArray());
+    } else if (dtype() == DType.DOUBLE) {
+      dtypeSize = DOUBLE_SIZE_BYTES;
+      tensorAsByteArray = new byte[(int) numel() * dtypeSize];
+      Tensor_float64 thiz = (Tensor_float64) this;
+      ByteBuffer.wrap(tensorAsByteArray).asDoubleBuffer().put(thiz.getDataAsDoubleArray());
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+    ByteBuffer byteBuffer =
+        ByteBuffer.allocate(1 + 1 + 4 * shape.length + dtypeSize * (int) numel());
+    byteBuffer.put((byte) dtype().jniCode);
+    byteBuffer.put((byte) shape.length);
+    for (long s : shape) {
+      byteBuffer.putInt((int) s);
+    }
+    byteBuffer.put(tensorAsByteArray);
+    return byteBuffer.array();
+  }
+
+  /**
+   * Deserializes a {@code Tensor} from a byte[].
+   *
+   * @param buffer The byte array to deserialize from.
+   * @return The deserialized {@code Tensor}.
+   * @apiNote This method is experimental and subject to change without notice. This does NOT
+   *     supoprt list type.
+   */
+  public static Tensor fromByteArray(byte[] bytes) {
+    if (bytes == null) {
+      throw new IllegalArgumentException("bytes cannot be null");
+    }
+    ByteBuffer buffer = ByteBuffer.wrap(bytes);
+    if (!buffer.hasRemaining()) {
+      throw new IllegalArgumentException("invalid buffer");
+    }
+    byte dtype = buffer.get();
+    byte shapeLength = buffer.get();
+    long[] shape = new long[(int) shapeLength];
+    long numel = 1;
+    for (int i = 0; i < shapeLength; i++) {
+      int dim = buffer.getInt();
+      if (dim < 0) {
+        throw new IllegalArgumentException("invalid shape");
+      }
+      shape[i] = dim;
+      numel *= dim;
+    }
+    if (dtype == DType.UINT8.jniCode) {
+      return new Tensor_uint8(buffer, shape);
+    } else if (dtype == DType.INT8.jniCode) {
+      return new Tensor_int8(buffer, shape);
+    } else if (dtype == DType.INT32.jniCode) {
+      return new Tensor_int32(buffer.asIntBuffer(), shape);
+    } else if (dtype == DType.INT64.jniCode) {
+      return new Tensor_int64(buffer.asLongBuffer(), shape);
+    } else if (dtype == DType.FLOAT.jniCode) {
+      return new Tensor_float32(buffer.asFloatBuffer(), shape);
+    } else if (dtype == DType.DOUBLE.jniCode) {
+      return new Tensor_float64(buffer.asDoubleBuffer(), shape);
+    } else {
+      throw new IllegalArgumentException("Unknown Tensor dtype");
+    }
+  }
 }