Repository: deepjavalibrary/djl
Pull request #1251 (2021-09-25 23:51:09)

[tensorflow] Fixes string tensor implemenation
1. Add create(String[], Shape) to NDArray interface
2. Implemented full string tensor for TF
3. Fixed Onnx NDManager.create(String) implementation

Change-Id: If51e55373ac800bd2df0ed1deb937fc60fca8b70

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1251](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (d452391) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/5907d505a55351394ca26babcd0d4dd7411a500a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (5907d50) will **increase** coverage by `0.18%`.
> The diff coverage is `82.50%`.

[![Impacted file tree graph](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/graphs/tree.svg?width=650&height=150&src=pr&token=mDKVNwFXRV&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

```diff
@@             Coverage Diff              @@
##             master    #1251      +/-   ##
============================================
+ Coverage     70.23%   70.41%   +0.18%     
- Complexity     5063     5079      +16     
============================================
  Files           485      485              
  Lines         22262    22292      +30     
  Branches       2375     2380       +5     
============================================
+ Hits          15635    15698      +63     
+ Misses         5414     5383      -31     
+ Partials       1213     1211       -2     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [...pi/src/main/java/ai/djl/ndarray/BaseNDManager.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvQmFzZU5ETWFuYWdlci5qYXZh) | `67.72% <ø> (+0.42%)` | :arrow_up: |
| [...n/java/ai/djl/onnxruntime/engine/OrtNDManager.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9vbm54cnVudGltZS9vbm54cnVudGltZS1lbmdpbmUvc3JjL21haW4vamF2YS9haS9kamwvb25ueHJ1bnRpbWUvZW5naW5lL09ydE5ETWFuYWdlci5qYXZh) | `67.56% <0.00%> (-5.41%)` | :arrow_down: |
| [...ain/java/ai/djl/tensorflow/engine/TfNDManager.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JmbG93L3RlbnNvcmZsb3ctZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3RlbnNvcmZsb3cvZW5naW5lL1RmTkRNYW5hZ2VyLmphdmE=) | `84.37% <50.00%> (-3.02%)` | :arrow_down: |
| [api/src/main/java/ai/djl/ndarray/NDManager.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRNYW5hZ2VyLmphdmE=) | `58.00% <100.00%> (+0.56%)` | :arrow_up: |
| [...ain/java/ai/djl/onnxruntime/engine/OrtNDArray.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9vbm54cnVudGltZS9vbm54cnVudGltZS1lbmdpbmUvc3JjL21haW4vamF2YS9haS9kamwvb25ueHJ1bnRpbWUvZW5naW5lL09ydE5EQXJyYXkuamF2YQ==) | `76.00% <100.00%> (+3.27%)` | :arrow_up: |
| [.../main/java/ai/djl/tensorflow/engine/TfNDArray.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JmbG93L3RlbnNvcmZsb3ctZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3RlbnNvcmZsb3cvZW5naW5lL1RmTkRBcnJheS5qYXZh) | `85.71% <100.00%> (+1.42%)` | :arrow_up: |
| [...ai/djl/tensorflow/engine/javacpp/JavacppUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JmbG93L3RlbnNvcmZsb3ctZW5naW5lL3NyYy9tYWluL2phdmEvYWkvZGpsL3RlbnNvcmZsb3cvZW5naW5lL2phdmFjcHAvSmF2YWNwcFV0aWxzLmphdmE=) | `81.28% <100.00%> (+16.33%)` | :arrow_up: |
| [...gine/src/main/java/ai/djl/dlr/engine/DlrModel.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy9kbHIvZGxyLWVuZ2luZS9zcmMvbWFpbi9qYXZhL2FpL2RqbC9kbHIvZW5naW5lL0Rsck1vZGVsLmphdmE=) | `51.61% <0.00%> (-3.23%)` | :arrow_down: |
| [...rt/src/main/java/ai/djl/tensorrt/jni/LibUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-ZW5naW5lcy90ZW5zb3JydC9zcmMvbWFpbi9qYXZhL2FpL2RqbC90ZW5zb3JydC9qbmkvTGliVXRpbHMuamF2YQ==) | `71.05% <0.00%> (-2.64%)` | :arrow_down: |
| [api/src/main/java/ai/djl/util/Platform.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL3V0aWwvUGxhdGZvcm0uamF2YQ==) | `69.81% <0.00%> (-1.89%)` | :arrow_down: |
| ... and [10 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1251/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [5907d50...d452391](https://codecov.io/gh/deepjavalibrary/djl/pull/1251?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
index 0e026e68a..018852a41 100644
--- a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -26,6 +26,7 @@ import java.nio.FloatBuffer;
 import java.nio.IntBuffer;
 import java.nio.LongBuffer;
 import java.nio.ShortBuffer;
+import java.nio.charset.Charset;
 import java.nio.file.Path;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
@@ -62,13 +63,7 @@ public abstract class BaseNDManager implements NDManager {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray create(String data) {
-        throw new UnsupportedOperationException("Not supported!");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public NDArray create(String[] data) {
+    public NDArray create(String[] data, Charset charset, Shape shape) {
         throw new UnsupportedOperationException("Not supported!");
     }
 
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index eca7ca041..00b90ff12 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -26,6 +26,8 @@ import java.nio.DoubleBuffer;
 import java.nio.FloatBuffer;
 import java.nio.IntBuffer;
 import java.nio.LongBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.function.Function;
 import java.util.stream.IntStream;
@@ -346,7 +348,19 @@ public interface NDArray extends NDResource {
      *
      * @return Array of Strings
      */
-    String[] toStringArray();
+    default String[] toStringArray() {
+        return toStringArray(StandardCharsets.UTF_8);
+    }
+
+    /**
+     * Converts this {@code NDArray} to a String array with the specified charset.
+     *
+     * <p>This method is only applicable to the String typed NDArray and not for printing purpose
+     *
+     * @param charset to charset for the string
+     * @return Array of Strings
+     */
+    String[] toStringArray(Charset charset);
 
     /**
      * Converts this {@code NDArray} to a Number array based on its {@link DataType}.
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 88fd54a70..3d456c699 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -18,6 +18,7 @@ import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.ndarray.types.SparseFormat;
 import java.nio.Buffer;
+import java.nio.charset.Charset;
 import java.util.Arrays;
 
 /**
@@ -165,7 +166,7 @@ public abstract class NDArrayAdapter implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public String[] toStringArray() {
+    public String[] toStringArray(Charset charset) {
         throw new UnsupportedOperationException(UNSUPPORTED_MSG);
     }
 
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index e1253766f..935f820fe 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -30,6 +30,8 @@ import java.nio.DoubleBuffer;
 import java.nio.FloatBuffer;
 import java.nio.IntBuffer;
 import java.nio.LongBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 
 /**
@@ -268,15 +270,51 @@ public interface NDManager extends AutoCloseable {
      * @param data the String data that needs to be set
      * @return a new instance of {@link NDArray}
      */
-    NDArray create(String data);
+    default NDArray create(String data) {
+        return create(new String[] {data}, StandardCharsets.UTF_8, new Shape());
+    }
+
+    /**
+     * Creates and initializes 1D {@link NDArray}.
+     *
+     * @param data the String data that needs to be set
+     * @return a new instance of {@link NDArray}
+     */
+    default NDArray create(String[] data) {
+        return create(data, StandardCharsets.UTF_8);
+    }
 
     /**
      * Creates and initializes 1D {@link NDArray}.
      *
      * @param data the String data that needs to be set
+     * @param charset the charset to decode the string
      * @return a new instance of {@link NDArray}
      */
-    NDArray create(String[] data);
+    default NDArray create(String[] data, Charset charset) {
+        return create(data, charset, new Shape(data.length));
+    }
+
+    /**
+     * Creates a String {@link NDArray} based on the provided shape.
+     *
+     * @param data the flattened String array
+     * @param shape the shape of the String NDArray
+     * @return a new instance of {@code NDArray}
+     */
+    default NDArray create(String[] data, Shape shape) {
+        return create(data, StandardCharsets.UTF_8, shape);
+    }
+
+    /**
+     * Creates a String {@link NDArray} based on the provided shape.
+     *
+     * @param data the flattened String array
+     * @param charset the charset to decode the string
+     * @param shape the shape of the String NDArray
+     * @return a new instance of {@code NDArray}
+     */
+    NDArray create(String[] data, Charset charset, Shape shape);
 
     /**
      * Creates and initializes a 1D {@link NDArray}.
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index 51f92db76..4c7718b31 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -29,6 +29,7 @@ import com.sun.jna.Pointer;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.IntBuffer;
+import java.nio.charset.Charset;
 import java.util.Arrays;
 import java.util.stream.IntStream;
 
@@ -269,7 +270,7 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
 
     /** {@inheritDoc} */
     @Override
-    public String[] toStringArray() {
+    public String[] toStringArray(Charset charset) {
         throw new UnsupportedOperationException("String NDArray is not supported!");
     }
 
@@ -293,14 +294,14 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
     @Override
     public void set(Buffer data) {
         int size = Math.toIntExact(size());
-        BaseNDManager.validateBufferSize(data, getDataType(), size);
+        DataType type = getDataType();
+        BaseNDManager.validateBufferSize(data, type, size);
         if (data.isDirect()) {
             JnaUtils.syncCopyFromCPU(getHandle(), data, size);
             return;
         }
 
-        DataType inputType = DataType.fromBuffer(data);
-        ByteBuffer buf = manager.allocateDirect(size * inputType.getNumOfBytes());
+        ByteBuffer buf = manager.allocateDirect(size * type.getNumOfBytes());
         BaseNDManager.copyBuffer(data, buf);
         JnaUtils.syncCopyFromCPU(getHandle(), buf, size);
     }
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
index 0fa6e1567..af4b05853 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
@@ -22,6 +22,7 @@ import ai.onnxruntime.OnnxTensor;
 import ai.onnxruntime.OrtException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.Charset;
 import java.util.UUID;
 
 /** {@code OrtNDArray} is the ONNX Runtime implementation of {@link NDArray}. */
@@ -74,9 +75,14 @@ public class OrtNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public String[] toStringArray() {
+    public String[] toStringArray(Charset charset) {
         try {
-            return (String[]) tensor.getValue();
+            Object obj = tensor.getValue();
+            if (obj instanceof String) {
+                // Scalar type;
+                return new String[] {(String) obj};
+            }
+            return (String[]) obj;
         } catch (OrtException e) {
             throw new EngineException(e);
         }
@@ -85,6 +91,9 @@ public class OrtNDArray extends NDArrayAdapter {
     /** {@inheritDoc} */
     @Override
     public ByteBuffer toByteBuffer() {
+        if (getDataType() == DataType.STRING) {
+            throw new IllegalArgumentException("Please use toStringArray() for String NDArray.");
+        }
         return tensor.getByteBuffer().order(ByteOrder.nativeOrder());
     }
 
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDManager.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDManager.java
index 6da4c8ba2..56235e806 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDManager.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDManager.java
@@ -26,6 +26,7 @@ import ai.onnxruntime.OrtException;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.Charset;
 
 /** {@code OrtNDManager} is the ONNX Runtime implementation of {@link NDManager}. */
 public class OrtNDManager extends BaseNDManager {
@@ -67,6 +68,10 @@ public class OrtNDManager extends BaseNDManager {
     /** {@inheritDoc} */
     @Override
     public OrtNDArray create(Buffer data, Shape shape, DataType dataType) {
+        if (dataType == DataType.STRING) {
+            throw new IllegalArgumentException(
+                    "Use NDManager.create(String[], Shape) to create String NDArray.");
+        }
         int size = Math.toIntExact(shape.size());
         BaseNDManager.validateBufferSize(data, dataType, size);
         OnnxTensor tensor = OrtUtils.toTensor(env, data, shape, dataType);
@@ -85,14 +90,9 @@ public class OrtNDManager extends BaseNDManager {
         return create(data, new Shape(data.length));
     }
 
-    /**
-     * Create A String tensor based on the provided shape.
-     *
-     * @param data the flattened String array
-     * @param shape the shape of the String NDArray
-     * @return a new instance of {@link NDArray}
-     */
-    public NDArray create(String[] data, Shape shape) {
+    /** {@inheritDoc} */
+    @Override
+    public NDArray create(String[] data, Charset charset, Shape shape) {
         try {
             return new OrtNDArray(this, alternativeManager, OrtUtils.toTensor(env, data, shape));
         } catch (OrtException e) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index 4c6ef56ab..a317642f7 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -24,6 +24,7 @@ import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.util.NativeResource;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -200,7 +201,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public String[] toStringArray() {
+    public String[] toStringArray(Charset charset) {
         throw new UnsupportedOperationException("String NDArray is not supported!");
     }
 
diff --git a/engines/tensorflow/tensorflow-engine/build.gradle b/engines/tensorflow/tensorflow-engine/build.gradle
index 78a8e846a..6123664f4 100644
--- a/engines/tensorflow/tensorflow-engine/build.gradle
+++ b/engines/tensorflow/tensorflow-engine/build.gradle
@@ -9,6 +9,7 @@ dependencies {
         exclude group: "junit", module: "junit"
     }
     testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
+    testRuntimeOnly "ai.djl.tensorflow:tensorflow-native-auto:${tensorflow_version}"
 }
 
 publishing {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 5ef8ced12..d315ff5af 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -13,6 +13,7 @@
 package ai.djl.tensorflow.engine;
 
 import ai.djl.Device;
+import ai.djl.ndarray.BaseNDManager;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDArrays;
 import ai.djl.ndarray.NDList;
@@ -26,7 +27,7 @@ import ai.djl.util.NativeResource;
 import ai.djl.util.Preconditions;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
+import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -173,21 +174,32 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
 
     /** {@inheritDoc} */
     @Override
-    public String[] toStringArray() {
+    public String[] toStringArray(Charset charset) {
         int size = Math.toIntExact(getShape().size());
-        return JavacppUtils.getString(getHandle(), size, StandardCharsets.UTF_8);
+        return JavacppUtils.getString(getHandle(), size, charset);
     }
 
     /** {@inheritDoc} */
     @Override
     public ByteBuffer toByteBuffer() {
+        if (getDataType() == DataType.STRING) {
+            throw new IllegalArgumentException("Please use toStringArray() for String NDArray.");
+        }
         return JavacppUtils.getByteBuffer(getHandle());
     }
 
     /** {@inheritDoc} */
     @Override
     public void set(Buffer data) {
-        throw new UnsupportedOperationException("Tensor cannot be modified after creation");
+        int size = Math.toIntExact(getShape().size());
+        BaseNDManager.validateBufferSize(data, getDataType(), size);
+        if (data instanceof ByteBuffer) {
+            JavacppUtils.setByteBuffer(getHandle(), (ByteBuffer) data);
+            return;
+        }
+        ByteBuffer buf = getManager().allocateDirect(size * getDataType().getNumOfBytes());
+        BaseNDManager.copyBuffer(data, buf);
+        JavacppUtils.setByteBuffer(getHandle(), buf);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDManager.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDManager.java
index 12ca09ab8..12c03e739 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDManager.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDManager.java
@@ -24,6 +24,7 @@ import ai.djl.util.Pair;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.Charset;
 import org.tensorflow.internal.c_api.TFE_Context;
 import org.tensorflow.internal.c_api.TFE_TensorHandle;
 import org.tensorflow.internal.c_api.TF_Tensor;
@@ -59,6 +60,10 @@ public class TfNDManager extends BaseNDManager {
     /** {@inheritDoc} */
     @Override
     public NDArray create(Shape shape, DataType dataType) {
+        if (dataType == DataType.STRING) {
+            throw new IllegalArgumentException(
+                    "Use NDManager.create(String[], Charset, Shape) to create String NDArray.");
+        }
         if (shape.dimension() == 0) {
             // TensorFlow does not support empty scalar(emtpy NDArray with 0 dimension)
             // initialize with scalar 0
@@ -71,6 +76,10 @@ public class TfNDManager extends BaseNDManager {
     /** {@inheritDoc} */
     @Override
     public TfNDArray create(Buffer data, Shape shape, DataType dataType) {
+        if (dataType == DataType.STRING) {
+            throw new IllegalArgumentException(
+                    "Use NDManager.create(String[], Charset, Shape) to create String NDArray.");
+        }
         int size = Math.toIntExact(shape.size());
         BaseNDManager.validateBufferSize(data, dataType, size);
         if (data.isDirect() && data instanceof ByteBuffer) {
@@ -86,8 +95,24 @@ public class TfNDManager extends BaseNDManager {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray create(String data) {
-        Pair<TF_Tensor, TFE_TensorHandle> pair = JavacppUtils.createStringTensor(data);
+    public NDArray create(String[] data, Charset charset, Shape shape) {
+        ByteBuffer[] buf = new ByteBuffer[data.length];
+        for (int i = 0; i < data.length; ++i) {
+            buf[i] = ByteBuffer.wrap(data[i].getBytes(charset));
+        }
+        return createStringTensor(shape, buf);
+    }
+
+    /**
+     * Creates a String {@link NDArray} based on the provided shape.
+     *
+     * @param shape the shape of the String NDArray
+     * @param data the flattened String array
+     * @return a new instance of {@code NDArray}
+     */
+    public NDArray createStringTensor(Shape shape, ByteBuffer... data) {
+        Pair<TF_Tensor, TFE_TensorHandle> pair =
+                JavacppUtils.createStringTensor(shape.getShape(), data);
         return new TfNDArray(this, pair.getValue(), pair.getKey());
     }
 
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/javacpp/JavacppUtils.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/javacpp/JavacppUtils.java
index d3cf5701a..cba1d1c27 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/javacpp/JavacppUtils.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/javacpp/JavacppUtils.java
@@ -271,10 +271,10 @@ public final class JavacppUtils {
     }
 
     @SuppressWarnings({"unchecked", "try"})
-    public static Pair<TF_Tensor, TFE_TensorHandle> createStringTensor(String src) {
+    public static Pair<TF_Tensor, TFE_TensorHandle> createStringTensor(
+            long[] dims, ByteBuffer[] src) {
         int dType = TfDataType.toTf(DataType.STRING);
-        long[] dims = {};
-        long numBytes = Loader.sizeof(TF_TString.class);
+        long numBytes = (long) Loader.sizeof(TF_TString.class) * src.length;
         try (PointerScope ignored = new PointerScope()) {
             /*
              * String tensor allocates a separate TF_TString memory. The TF_TString will
@@ -285,9 +285,12 @@ public final class JavacppUtils {
              */
             TF_Tensor tensor = AbstractTF_Tensor.allocateTensor(dType, dims, numBytes);
             Pointer pointer = tensorflow.TF_TensorData(tensor).capacity(numBytes);
-            TF_TString data = new TF_TString(pointer).capacity(pointer.position() + 1);
-            byte[] buf = src.getBytes(StandardCharsets.UTF_8);
-            tensorflow.TF_TString_Copy(data, new BytePointer(buf), buf.length);
+            TF_TString data = new TF_TString(pointer).capacity(pointer.position() + src.length);
+            for (int i = 0; i < src.length; ++i) {
+                TF_TString tstring = data.getPointer(i);
+                tensorflow.TF_TString_Copy(tstring, new BytePointer(src[i]), src[i].remaining());
+            }
+
             TF_Status status = TF_Status.newStatus();
             TFE_TensorHandle handle = AbstractTFE_TensorHandle.newTensor(tensor, status);
             status.throwExceptionIfNotOK();
@@ -367,6 +370,19 @@ public final class JavacppUtils {
         }
     }
 
+    @SuppressWarnings({"unchecked", "try"})
+    public static void setByteBuffer(TFE_TensorHandle handle, ByteBuffer data) {
+        try (PointerScope ignored = new PointerScope()) {
+            // convert to TF_Tensor
+            TF_Status status = TF_Status.newStatus();
+            TF_Tensor tensor = tensorflow.TFE_TensorHandleResolve(handle, status).withDeallocator();
+            status.throwExceptionIfNotOK();
+            Pointer pointer =
+                    tensorflow.TF_TensorData(tensor).capacity(tensorflow.TF_TensorByteSize(tensor));
+            pointer.asByteBuffer().put(data);
+        }
+    }
+
     @SuppressWarnings({"unchecked", "try"})
     public static ByteBuffer getByteBuffer(TFE_TensorHandle handle) {
         try (PointerScope ignored = new PointerScope()) {


Test patch:
diff --git a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
index 7057f52c9..cc421bb33 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/test/java/ai/djl/onnxruntime/engine/OrtTest.java
@@ -54,8 +54,8 @@ public class OrtTest {
                 Assert.assertThrows(() -> m.load(path, null));
                 Assert.assertThrows(() -> m.load(path, "invalid.onnx"));
 
-                Path modleFile = path.resolve(model.getName() + ".onnx");
-                m.load(modleFile);
+                Path modelFile = path.resolve(model.getName() + ".onnx");
+                m.load(modelFile);
 
                 m.close();
             }
@@ -86,6 +86,11 @@ public class OrtTest {
 
             array = manager.create("string");
             Assert.assertEquals(array.toStringArray()[0], "string");
+            final NDArray a = array;
+            Assert.assertThrows(IllegalArgumentException.class, a::toByteBuffer);
+
+            array = manager.create(new String[] {"string1", "string2"});
+            Assert.assertEquals(array.toStringArray()[1], "string2");
         }
     }
 
diff --git a/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/TfNDManagerTest.java b/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/TfNDManagerTest.java
new file mode 100644
index 000000000..4ea91004f
--- /dev/null
+++ b/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/TfNDManagerTest.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.tensorflow.engine;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class TfNDManagerTest {
+
+    @Test
+    public void testNDArray() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.create("string");
+            Assert.assertEquals(array.toStringArray()[0], "string");
+
+            array = manager.create(new String[] {"string1", "string2"});
+            Assert.assertEquals(array.toStringArray()[1], "string2");
+            final NDArray a = array;
+            Assert.assertThrows(IllegalArgumentException.class, a::toByteBuffer);
+
+            ByteBuffer buf1 = ByteBuffer.wrap("string1".getBytes(StandardCharsets.UTF_8));
+            ByteBuffer buf2 = ByteBuffer.wrap(new byte[] {1, 0});
+
+            array = ((TfNDManager) manager).createStringTensor(new Shape(2, 1), buf1, buf2);
+            Assert.assertEquals(
+                    array.toStringArray()[1].getBytes(StandardCharsets.UTF_8), buf2.array());
+
+            array = manager.zeros(new Shape(2));
+            float[] expected = {2, 3};
+            array.set(expected);
+            Assert.assertEquals(array.toFloatArray(), expected);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> manager.create(new Shape(1), DataType.STRING));
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> manager.create(buf1, new Shape(1), DataType.STRING));
+        }
+    }
+}
diff --git a/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/package-info.java b/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/package-info.java
new file mode 100644
index 000000000..27819d8a3
--- /dev/null
+++ b/engines/tensorflow/tensorflow-engine/src/test/java/ai/djl/tensorflow/engine/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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.
+ */
+/** The test classes for TensorFlow engine features. */
+package ai.djl.tensorflow.engine;