Repository: deepjavalibrary/djl
Pull request #1772 (2022-07-05 21:33:16)

[api] Supports Image output for ImageServingTranslator
Change-Id: I4646a9b411bbaaf81203b7531bf70f4bdfe9f9f8

## 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/1772?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1772](https://codecov.io/gh/deepjavalibrary/djl/pull/1772?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (d604521) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `1.46%`.
> The diff coverage is `70.30%`.

```diff
@@             Coverage Diff              @@
##             master    #1772      +/-   ##
============================================
- Coverage     72.08%   70.62%   -1.47%     
- Complexity     5126     5574     +448     
============================================
  Files           473      527      +54     
  Lines         21970    24745    +2775     
  Branches       2351     2682     +331     
============================================
+ Hits          15838    17477    +1639     
- Misses         4925     5939    +1014     
- Partials       1207     1329     +122     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1772?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <ø> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `8.33% <0.00%> (-0.50%)` | :arrow_down: |
| ... and [411 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1772/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/1772?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/1772?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [a1c7448...d604521](https://codecov.io/gh/deepjavalibrary/djl/pull/1772?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/modality/cv/translator/ImageServingTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/ImageServingTranslator.java
index a33c54ec6..4df9b7e2f 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/ImageServingTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/ImageServingTranslator.java
@@ -25,6 +25,7 @@ import ai.djl.translate.TranslatorContext;
 import ai.djl.util.JsonSerializable;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 
 /** A {@link Translator} that can handle generic CV {@link Input} and {@link Output}. */
@@ -56,10 +57,18 @@ public class ImageServingTranslator implements Translator<Input, Output> {
         Object obj = translator.processOutput(ctx, list);
         if (obj instanceof JsonSerializable) {
             output.add((JsonSerializable) obj);
+            output.addProperty("Content-Type", "application/json");
+        } else if (obj instanceof Image) {
+            Image img = (Image) obj;
+            try (ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
+                img.save(bos, "png");
+                output.add(bos.toByteArray());
+            }
+            output.addProperty("Content-Type", "image/png");
         } else {
             output.add(BytesSupplier.wrapAsJson(obj));
+            output.addProperty("Content-Type", "application/json");
         }
-        output.addProperty("Content-Type", "application/json");
         return output;
     }
 
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index c922d7e75..516e279a2 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -595,7 +595,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      *
      * @param index select the entries of an {@code NDArray}
      * @param data numbers to assign to the indexed entries
-     * @return The NDArray with updated values
+     * @return the NDArray with updated values
      */
     NDArray put(NDArray index, NDArray data);
 
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
index 47994b22c..3c7f67d4f 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
@@ -79,89 +79,80 @@ public abstract class NDArrayIndexer {
     }
 
     /**
-     * Sets the values of the array at the fullSlice with an array.
-     *
-     * @param array the array to set
-     * @param fullSlice the fullSlice of the index to set in the array
-     * @param value the value to set with
-     */
-    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value);
-
-    /**
-     * Sets the values of the array at the boolean locations with an array.
-     *
-     * @param array the array to set
-     * @param indices a boolean array where true indicates values to update
-     * @param value the value to set with when condition is true
-     */
-    public void set(NDArray array, NDIndexBooleans indices, NDArray value) {
-        array.intern(NDArrays.where(indices.getIndex(), value, array));
-    }
-
-    /**
-     * Sets the values of the array at the index locations with an array.
+     * Sets the entries of array at the indexed locations with the parameter value. The value can be
+     * only Number or NDArray.
      *
      * @param array the array to set
      * @param index the index to set at in the array
      * @param value the value to set with
      */
-    public void set(NDArray array, NDIndex index, NDArray value) {
+    public void set(NDArray array, NDIndex index, Object value) {
+        NDIndexFullSlice fullSlice =
+                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
+        if (fullSlice != null) {
+            if (value instanceof Number) {
+                set(array, fullSlice, (Number) value);
+            } else if (value instanceof NDArray) {
+                set(array, fullSlice, (NDArray) value);
+            } else {
+                throw new IllegalArgumentException(
+                        "The type of value to assign cannot be other than NDArray and Number.");
+            }
+            return;
+        }
+
         List<NDIndexElement> indices = index.getIndices();
         if (!indices.isEmpty() && indices.get(0) instanceof NDIndexBooleans) {
             if (indices.size() != 1) {
                 throw new IllegalArgumentException(
-                        "get() currently didn't support more that one boolean NDArray");
+                        "set() currently doesn't support more than one boolean NDArray");
+            }
+            if (value instanceof Number) {
+                set(
+                        array,
+                        (NDIndexBooleans) indices.get(0),
+                        array.getManager().create((Number) value));
+            } else if (value instanceof NDArray) {
+                set(array, (NDIndexBooleans) indices.get(0), (NDArray) value);
+            } else {
+                throw new IllegalArgumentException(
+                        "The type of value to assign cannot be other than NDArray and Number.");
             }
-            set(array, (NDIndexBooleans) indices.get(0), value);
         }
 
-        NDIndexFullSlice fullSlice =
-                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
-        if (fullSlice != null) {
-            set(array, fullSlice, value);
-            return;
-        }
         throw new UnsupportedOperationException(
                 "set() currently supports all, fixed, and slices indices");
     }
 
     /**
-     * Sets the values of the array at the fullSlice with a number.
+     * Sets the values of the array at the fullSlice with an array.
      *
      * @param array the array to set
      * @param fullSlice the fullSlice of the index to set in the array
      * @param value the value to set with
      */
-    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, Number value);
+    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value);
 
     /**
-     * Sets the values of the array at the index locations with a number.
+     * Sets the values of the array at the boolean locations with an array.
      *
      * @param array the array to set
-     * @param index the index to set at in the array
-     * @param value the value to set with
+     * @param indices a boolean array where true indicates values to update
+     * @param value the value to set with when condition is true
      */
-    public void set(NDArray array, NDIndex index, Number value) {
-        NDIndexFullSlice fullSlice =
-                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
-        if (fullSlice != null) {
-            set(array, fullSlice, value);
-            return;
-        }
-        // use booleanMask for NDIndexBooleans case
-        List<NDIndexElement> indices = index.getIndices();
-        if (!indices.isEmpty() && indices.get(0) instanceof NDIndexBooleans) {
-            if (indices.size() != 1) {
-                throw new IllegalArgumentException(
-                        "set() currently didn't support more that one boolean NDArray");
-            }
-            set(array, (NDIndexBooleans) indices.get(0), array.getManager().create(value));
-            return;
-        }
-        throw new UnsupportedOperationException(
-                "set() currently supports all, fixed, and slices indices");
+    public void set(NDArray array, NDIndexBooleans indices, NDArray value) {
+        array.intern(NDArrays.where(indices.getIndex(), value, array));
     }
 
+    /**
+     * Sets the values of the array at the fullSlice with a number.
+     *
+     * @param array the array to set
+     * @param fullSlice the fullSlice of the index to set in the array
+     * @param value the value to set with
+     */
+    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, Number value);
+
     /**
      * Sets a scalar value in the array at the indexed location.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDIndex.java b/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
index 8b8b40146..a5fb0a2f1 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
@@ -111,6 +111,11 @@ public class NDIndex {
      *     // Uses null to add an extra axis to the output array
      *     assertEquals(a.get(new NDIndex(":2, null, 0, :2")).getShape(), new Shape(2, 1, 2));
      *
+     *     // Gets entries of an NDArray with mixed index
+     *     index1 = manager.create(new long[] {0, 1, 1}, new Shape(2));
+     *     bool1 = manager.create(new boolean[] {true, false, true});
+     *     assertEquals(a.get(new NDIndex(":{}, {}, {}, {}" 2, index1, bool1, null).getShape(), new Shape(2, 2, 1));
+     *
      * </pre>
      *
      * @param indices a comma separated list of indices corresponding to either subsections,
diff --git a/docs/development/inference_performance_optimization.md b/docs/development/inference_performance_optimization.md
index e60b4728a..74d2aa78f 100644
--- a/docs/development/inference_performance_optimization.md
+++ b/docs/development/inference_performance_optimization.md
@@ -12,7 +12,7 @@ memory consumption compare to Python.
 DJL `Predictor` is not designed to be thread-safe (although some implementation is),
 we recommend creating a new [Predictor](https://javadoc.io/doc/ai.djl/api/latest/ai/djl/inference/Predictor.html) for each thread.
 
-For a reference implementation, see [Multi-threaded Benchmark](https://github.com/deepjavalibrary/djl/blob/master/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java).
+For a reference implementation, see [Multi-threaded Benchmark](https://github.com/deepjavalibrary/djl-serving/blob/master/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java).
 
 you need to set corresponding configuration based on the engine you want to use.
 
@@ -111,10 +111,11 @@ This should only be disabled when you do not have the time to "warmup" a model w
 #### Multithreading Inference
 You can follow the same steps as other engines for running multithreading inference using TensorFlow engine.
 It's recommended to use one `Predictor` for each thread and avoid using a new `Predictor` for each inference call.
-You can refer to our [Multithreading Benchmark](https://github.com/deepjavalibrary/djl/blob/master/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java) as an example,
+You can refer to our [Multithreading Benchmark](https://github.com/deepjavalibrary/djl-serving/blob/master/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java) as an example,
 here is how to run it using TensorFlow engine.
 
 ```bash
+cd djl-serving
 ./gradlew benchmark --args='-e TensorFlow -c 100 -t -1 -u djl://ai.djl.tensorflow/resnet/0.0.1/resnet50 -s 1,224,224,3'
 ```
 
diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml
index c54afb223..fdaedd333 100644
--- a/docs/mkdocs.yml
+++ b/docs/mkdocs.yml
@@ -81,7 +81,7 @@ nav:
         - 'docs/development/configure_logging.md'
         - 'docs/how_to_collect_metrics.md'
         - 'docs/development/inference_performance_optimization.md'
-        - 'extensions/benchmark/README.md'
+        - 'docs/serving/benchmark/README.md'
         - 'docs/development/profiler.md'
       - 'docs/development/cache_management.md'
       - 'docs/development/memory_management.md'
diff --git a/docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md b/docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md
index be1a0a346..b7091b9d1 100644
--- a/docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md
+++ b/docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md
@@ -52,6 +52,44 @@ Please refer to these two examples:
 1. [Object Detection with TensorFlow](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/ObjectDetection.java) for loading from TensorFlow Hub url.
 2. [BERT Classification](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/BertClassification.java) for loading from local downloaded model.
 
+## How to import TensorFlow models that use [TensorFlow Extensions](https://www.tensorflow.org/resources/libraries-extensions)
+
+You can use DJL to run TensorFlow models that require ops available in TensorFlow extensions (not in core), such as models that use ops in  [TensorFlow Text](https://www.tensorflow.org/text).
+
+To use such models in DJL you must know:
+* Which extension contains the additional required ops
+* The TensorFlow extension library version that is compatible with the version of Tensorflow currently being used in DJL
+
+Here we show an example that uses DJL to run the [Multilingual Universal Sentence Encoder](https://tfhub.dev/google/universal-sentence-encoder-multilingual/3). 
+This is a text encoder from tensorflow hub that uses ops from the TensorFlow Text extension.
+* You need to download the library files for the extension so that we can load them for use. [Here](https://github.com/tensorflow/text#install-using-pip) are instructions for downloading TensorFlow text. We'll use version `2.7.0` in this example.
+* After the library files are downloaded, you are ready to load them for use in DJL. You do this by using
+TensorFlow Java `loadLibrary(...)` API before loading the model in DJL. 
+* We can update the existing [UniversalSentenceEncoder](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/UniversalSentenceEncoder.java) example to use the Multilingual model:
+
+```java
+// Build Criteria for Multilingual Universal Sentence Encoder from TensorFlow Hub
+String modelUrl = "https://storage.googleapis.com/tfhub-modules/google/universal-sentence-encoder-multilingual/3.tar.gz"
+Criteria<String[], float[][]> criteria =
+    Criteria.builder()
+        .optApplication(Application.NLP.TEXT_EMBEDDING)
+        .setTypes(String[].class, float[][].class)
+        .optModelUrls(modelUrl)
+        .optTranslator(new MyTranslator())
+        .optEngine("TensorFlow")
+        .optProgress(new ProgressBar())
+        .build();
+
+// Load TensorFlow Text libraries for necessary ops
+// Example Path of installation on Linux may look like: /home/<user>/.local/python-3.8.3/lib/python3.8/site-packages/tensorflow_text/
+TensorFlow.loadLibrary("<path_to_library_installation>/python/ops/_sentencepiece_tokenizer.so");
+        
+// Run the model
+try (ZooModel<String[], float[][]> model = criteria.loadModel();
+    Predictor<String[], float[][]> predictor = model.newPredictor()) {
+    return predictor.predict(inputs.toArray(new String[0]));
+}
+```
 ## How to load TensorFlow Checkpoints
 
 To load an TensorFlow Estimator checkpoint, you need to convert it to SavedModel format in using Python.
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
index 84c1f38f2..6f2cf648b 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
@@ -179,6 +179,11 @@ public class OrtModel extends BaseModel {
             ortSession.setCPUArenaAllocator(true);
         }
 
+        String customOpLibrary = (String) options.get("customOpLibrary");
+        if (customOpLibrary != null) {
+            ortSession.registerCustomOpLibrary(customOpLibrary);
+        }
+
         Device device = manager.getDevice();
         if (options.containsKey("ortDevice")) {
             String ortDevice = (String) options.get("ortDevice");
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
index 07daec093..fcc8958d0 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
@@ -72,6 +72,25 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void set(NDArray array, NDIndex index, Object data) {
+        PtNDArray ptArray =
+                array instanceof PtNDArray
+                        ? (PtNDArray) array
+                        : manager.create(
+                                array.toByteBuffer(), array.getShape(), array.getDataType());
+
+        if (data instanceof Number) {
+            JniUtils.indexAdvPut(ptArray, index, (PtNDArray) manager.create((Number) data));
+        } else if (data instanceof NDArray) {
+            JniUtils.indexAdvPut(ptArray, index, (PtNDArray) data);
+        } else {
+            throw new IllegalArgumentException(
+                    "The type of value to assign cannot be other than NDArray and Number.");
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index a4f2a6100..3364cd45e 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -409,7 +409,69 @@ public final class JniUtils {
 
         return new PtNDArray(
                 ndArray.getManager(),
-                PyTorchLibrary.LIB.torchIndexReturn(ndArray.getHandle(), torchIndexHandle));
+                PyTorchLibrary.LIB.torchIndexAdvGet(ndArray.getHandle(), torchIndexHandle));
+    }
+
+    @SuppressWarnings("OptionalGetWithoutIsPresent")
+    public static void indexAdvPut(PtNDArray ndArray, NDIndex index, PtNDArray data) {
+        if (ndArray == null) {
+            return;
+        }
+
+        // Index aggregation
+        List<NDIndexElement> indices = index.getIndices();
+        long torchIndexHandle = PyTorchLibrary.LIB.torchIndexInit(indices.size());
+        ListIterator<NDIndexElement> it = indices.listIterator();
+        while (it.hasNext()) {
+            if (it.nextIndex() == index.getEllipsisIndex()) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+            }
+
+            NDIndexElement elem = it.next();
+            if (elem instanceof NDIndexNull) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
+            } else if (elem instanceof NDIndexSlice) {
+                Long min = ((NDIndexSlice) elem).getMin();
+                Long max = ((NDIndexSlice) elem).getMax();
+                Long step = ((NDIndexSlice) elem).getStep();
+                int nullSliceBin = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
+                // nullSliceBin encodes whether the slice (min, max) is null:
+                // is_null == 1, ! is_null == 0;
+                // 0b11 == 3, 0b10 = 2, ...
+                PyTorchLibrary.LIB.torchIndexAppendSlice(
+                        torchIndexHandle,
+                        min == null ? 0 : min,
+                        max == null ? 0 : max,
+                        step == null ? 1 : step,
+                        nullSliceBin);
+            } else if (elem instanceof NDIndexAll) {
+                PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, 0, 0, 1, 3);
+            } else if (elem instanceof NDIndexFixed) {
+                PyTorchLibrary.LIB.torchIndexAppendFixed(
+                        torchIndexHandle, ((NDIndexFixed) elem).getIndex());
+            } else if (elem instanceof NDIndexBooleans) {
+                PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
+                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
+            } else if (elem instanceof NDIndexTake) {
+                PtNDArray indexArr = (PtNDArray) ((NDIndexTake) elem).getIndex();
+                if (indexArr.getDataType() != DataType.INT64) {
+                    indexArr = indexArr.toType(DataType.INT64, true);
+                }
+                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
+            } else if (elem instanceof NDIndexPick) {
+                // Backward compatible
+                NDIndexFullPick fullPick =
+                        NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
+                pick(ndArray, ndArray.getManager().from(fullPick.getIndices()), fullPick.getAxis());
+                return;
+            }
+        }
+        if (indices.size() == index.getEllipsisIndex()) {
+            PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+        }
+
+        PyTorchLibrary.LIB.torchIndexAdvPut(
+                ndArray.getHandle(), torchIndexHandle, data.getHandle());
     }
 
     public static void indexSet(
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 36153510c..542ef3b54 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -192,6 +192,8 @@ final class PyTorchLibrary {
             long[] maxIndices,
             long[] stepIndices);
 
+    native void torchIndexAdvPut(long handle, long torchIndexHandle, long data);
+
     native void torchSet(long handle, ByteBuffer data);
 
     native long torchSlice(long handle, long dim, long start, long end, long step);
@@ -607,7 +609,7 @@ final class PyTorchLibrary {
 
     native long torchIndexInit(int size);
 
-    native long torchIndexReturn(long handle, long torchIndexHandle);
+    native long torchIndexAdvGet(long handle, long torchIndexHandle);
 
     native void torchIndexAppendNoneEllipsis(long torchIndexHandle, boolean isEllipsis);
 
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
index 21be55ae5..2b89fb98f 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
@@ -129,7 +129,7 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexInit(JN
   API_END_RETURN()
 }
 
-JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexReturn(
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAdvGet(
     JNIEnv* env, jobject jthis, jlong jhandle, jlong jtorch_index_handle) {
   API_BEGIN()
   const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
@@ -194,6 +194,16 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexPut(JNIE
   API_END()
 }
 
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAdvPut(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlong jtorch_index_handle, jlong jdata_handle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* data_ptr = reinterpret_cast<torch::Tensor*>(jdata_handle);
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  ((torch::Tensor) *tensor_ptr).index_put_(*index_ptr, *data_ptr);
+  API_END()
+}
+
 JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchSet(
     JNIEnv* env, jobject jthis, jlong jhandle, jobject jbuffer) {
   API_BEGIN()
diff --git a/extensions/benchmark/README.md b/extensions/benchmark/README.md
index ab3a8ab0f..9e32483c8 100644
--- a/extensions/benchmark/README.md
+++ b/extensions/benchmark/README.md
@@ -13,310 +13,5 @@ With djl-bench, you can easily compare your model's behavior in different use ca
 - running with different engines
 - running with different version of the engine
 
-djl-bench currently support benchmark the following type of models:
 
-- PyTorch TorchScript model
-- TensorFlow SavedModel bundle
-- Apache MXNet model
-- ONNX model
-- PaddlePaddle model
-- TFLite model
-- TensorRT model
-- XGBoost model
-- Python script model
-- Neo DLR (TVM) model
-
-You can build djl-bench from source if you need to benchmark fastText/BlazingText/Sentencepiece models.
-
-## Installation
-
-For macOS
-
-```
-brew install cask djl-bench
-```
-
-For Ubuntu
-
-- Install using snap
-
-```
-sudo snap install djlbench --classic
-sudo snap alias djlbench djl-bench
-```
-
-- Or download .deb package from S3
-
-```
-curl -O https://publish.djl.ai/djl-bench/0.17.0/djl-bench_0.17.0-1_all.deb
-sudo dpkg -i djl-bench_0.17.0-1_all.deb
-```
-
-For centOS or Amazon Linux 2
-
-You can download djl-bench zip file from [here](https://publish.djl.ai/djl-bench/0.17.0/benchmark-0.17.0.zip).
-
-```
-curl -O https://publish.djl.ai/djl-bench/0.17.0/benchmark-0.17.0.zip
-unzip benchmark-0.17.0.zip
-rm benchmark-0.17.0.zip
-sudo ln -s $PWD/benchmark-0.17.0/bin/benchmark /usr/bin/djl-bench
-```
-
-For Windows
-
-We are considering to create a `chocolatey` package for Windows. For the time being, you can
-download djl-bench zip file from [here](https://publish.djl.ai/djl-bench/0.17.0/benchmark-0.17.0.zip).
-
-Or you can run benchmark using gradle:
-
-```
-cd djl
-
-gradlew benchmark --args="--help"
-```
-
-## Prerequisite
-
-Please ensure Java 8+ is installed and you are using an OS that DJL supported with.
-
-After that, you need to clone the djl project and `cd` into the folder.
-
-DJL supported OS:
-
-- Ubuntu 18.04 and above
-- Amazon Linux 2 and above
-- MacOS latest version
-- Windows 10 (Windows Server 2016+)
-
-If you are trying to use GPU, please ensure the CUDA driver is installed. You can verify that through:
-
-```
-nvcc -V
-```
-
-to checkout the version. For different Deep Learning engine you are trying to run the benchmark,
-they have different CUDA version to support. Please check the individual Engine documentation to ensure your CUDA version is supported.
-
-## Sample benchmark script
-
-Here is a few sample benchmark script for you to refer. You can also skip this and directly follow
-the 4-step instructions for your own model.
-
-Benchmark on a Tensorflow model from [tfhub](https://tfhub.dev/) url with all-zeros NDArray input for 10 times:
-
-```
-djl-bench -e TensorFlow -u https://tfhub.dev/tensorflow/resnet_50/classification/1 -c 10 -s 1,224,224,3
-```
-
-Similarly, this is for PyTorch
-
-```
-djl-bench -e PyTorch -u https://alpha-djl-demos.s3.amazonaws.com/model/djl-blockrunner/pytorch_resnet18.zip -n traced_resnet18 -c 10 -s 1,3,224,224
-```
-
-Benchmark a model from [ONNX Model Zoo](https://github.com/onnx/models)
-
-```
-djl-bench -e OnnxRuntime -u https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet18v1/resnet18v1.tar.gz -s 1,3,224,224 -n resnet18v1/resnet18v1 -c 10
-```
-
-### Benchmark from ModelZoo
-
-#### MXNet
-
-Resnet50 image classification model:
-
-```
-djl-bench -c 2 -s 1,3,224,224 -u djl://ai.djl.mxnet/resnet/0.0.1/resnet50_v2
-```
-
-#### PyTorch
-
-SSD object detection model:
-
-```
-djl-bench -e PyTorch -c 2 -s 1,3,300,300 -u djl://ai.djl.pytorch/ssd/0.0.1/ssd_300_resnet50
-```
-
-## Configuration of Benchmark script
-
-To start your benchmarking, we need to make sure we provide the following information.
-
-- The Deep Learning Engine
-- The source of the model
-- How many runs you would like to make
-- Sample input for the model
-- (Optional) Multi-thread benchmark
-
-The benchmark script located [here](https://github.com/deepjavalibrary/djl/blob/master/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java).
-
-Just do the following:
-
-```
-djl-bench --help
-```
-
-This will print out the possible arguments to pass in:
-
-```
-usage: djl-bench [-p MODEL-PATH] -s INPUT-SHAPES [OPTIONS]
- -c,--iteration <ITERATION>               Number of total iterations.
- -d,--duration <DURATION>                 Duration of the test in minutes.
- -e,--engine <ENGINE-NAME>                Choose an Engine for the benchmark.
- -g,--gpus <NUMBER_GPUS>                  Number of GPUS to run multithreading inference.
- -h,--help                                Print this help.
- -l,--delay <DELAY>                       Delay of incremental threads.
-    --model-arguments <MODEL-ARGUMENTS>   Specify model loading arguments.
-    --model-options <MODEL-OPTIONS>       Specify model loading options.
- -n,--model-name <MODEL-NAME>             Specify model file name.
-    --neuron-cores <NEURON-CORES>         Number of neuron cores to run multithreading inference, See
-                                          https://awsdocs-neuron.readthedocs-hosted.com.
- -o,--output-dir <OUTPUT-DIR>             Directory for output logs.
- -p,--model-path <MODEL-PATH>             Model directory file path.
- -s,--input-shapes <INPUT-SHAPES>         Input data shapes for the model.
- -t,--threads <NUMBER_THREADS>            Number of inference threads.
- -u,--model-url <MODEL-URL>               Model archive file URL.
-```
-
-### Step 1: Pick your deep engine
-
-By default, the above script will use MXNet as the default Engine, but you can always change that by adding the followings:
-
-```
--e TensorFlow # TensorFlow
--e PyTorch # PyTorch
--e MXNet # Apache MXNet
--e PaddlePaddle # PaddlePaddle
--e OnnxRuntime # pytorch
--e TFLite # TFLite
--e TensorRT # TensorRT
--e DLR # Neo DLR
--e XGBoost # XGBoost
--e Python # Python script
-```
-
-### Step 2: Identify the source of your model
-
-DJL accept variety of models came from different places.
-
-#### Remote location
-
-Use `--model-url` option to load a model from a URL. The URL must point to an archive file.
-
-The following is a pytorch model
-
-```
--u https://alpha-djl-demos.s3.amazonaws.com/model/djl-blockrunner/pytorch_resnet18.zip
-```
-We would recommend to make model files in a zip for better file tracking.
-
-#### Local directory
-
-Use `--model-path` option to load model from a local directory or an archive file.
-
-Mac/Linux
-
-```
--p /home/ubuntu/models/pytorch_resnet18
-or
--p /home/ubuntu/models/pytorch_resnet18.zip
-```
-
-Windows
-
-```
--p C:\models\pytorch_resnet18
-or
--p C:\models\pytorch_resnet18.zip
-```
-
-If the model file name is different from the parent folder name (or the archive file name), you need
-to specify `--model-name` in the `--args`:
-
-```
--n traced_resnet18
-```
-
-### Step 3: Define how many runs you would like to make
-
-add `-c` inside with a number
-
-```
--c 1000
-```
-
-This will run 1000 times inference.
-
-### Step 4: Define your model inputs
-
-The benchmark script uses dummy NDArray inputs.
-It will make fake NDArrays (like `NDArray.ones`) to feed in the model for inference.
-
-If we would like to fake an image:
-
-```
--s 1,3,224,224
-```
-
-This will create a NDArray (DataType FLOAT32) of shape(1, 3, 224, 224).
-
-If your model requires multiple inputs like three NDArrays with shape 1, 384 and 384. You can do the followings:
-
-```
--s (1),(384),(384)
-```
-
-If you input `DataType` is not FLOAT32, you can specify the data type with suffix:
-
-- f: FLOAT32, this is default and is optional
-- s: FLOAT16 (short float)
-- d: FLOAT64 (double)
-- u: UINT8 (unsigned byte)
-- b: INT8 (byte)
-- i: INT32 (int)
-- l: INT64 (long)
-- B: BOOLEAN (boolean)
-
-For example:
-
-```
--s (1)i,(384)f,(384)
-```
-
-### Optional Step: multithreading inference
-
-You can also do multi-threading inference with DJL. For example, if you would like to run the inference with 10 threads:
-
-```
--t 10
-```
-
-Best thread number for your system: The same number of cores your system have or double of the total cores.
-
-You can also add `-l` to simulate the increment load for your inference server. It will add threads with the delay of time.
-
-```
--t 10 -l 100
-```
-
-The above code will create 10 threads with the wait time of 100ms.
-
-## Advanced use cases
-
-For different purposes, we designed different mode you can play with. Such as the following arg:
-
-```
--d 86400
-```
-
-This will ask the benchmark script repeatedly running the designed task for 86400 seconds (24 hour).
-If you would like to make sure DJL is stable in the long run, you can do that.
-
-You can also keep monitoring the DJL memory usages by enable the following flag:
-
-```
-export BENCHMARK_OPTS="-Dcollect-memory=true"
-```
-
-The memory report will be made available in `build/memory.log`.
+**This module has been moved to [deepjavalibrary/djl-serving/benchmark](https://github.com/deepjavalibrary/djl-serving/tree/master/benchmark).**
diff --git a/extensions/benchmark/build.gradle b/extensions/benchmark/build.gradle
deleted file mode 100644
index fe3a54aa3..000000000
--- a/extensions/benchmark/build.gradle
+++ /dev/null
@@ -1,151 +0,0 @@
-plugins {
-    id 'application'
-    id "nebula.ospackage" version "9.0.0"
-}
-
-boolean isRelease = project.hasProperty("release") || project.hasProperty("staging")
-
-dependencies {
-    implementation "commons-cli:commons-cli:${commons_cli_version}"
-    implementation "org.apache.logging.log4j:log4j-slf4j-impl:${log4j_slf4j_version}"
-    if (isRelease) {
-        implementation platform("ai.djl:bom:${djl_version}")
-
-        implementation "ai.djl:model-zoo"
-        runtimeOnly "ai.djl.pytorch:pytorch-model-zoo"
-        runtimeOnly "ai.djl.tensorflow:tensorflow-model-zoo"
-        runtimeOnly "ai.djl.mxnet:mxnet-model-zoo"
-        runtimeOnly "ai.djl.paddlepaddle:paddlepaddle-model-zoo"
-        runtimeOnly "ai.djl.onnxruntime:onnxruntime-engine"
-        runtimeOnly "ai.djl.tflite:tflite-engine"
-        runtimeOnly "ai.djl.dlr:dlr-engine"
-        runtimeOnly "ai.djl.ml.xgboost:xgboost"
-        runtimeOnly "ai.djl.python:python"
-        runtimeOnly "ai.djl.tensorrt:tensorrt"
-    } else {
-        implementation project(":model-zoo")
-
-        runtimeOnly project(":engines:pytorch:pytorch-model-zoo")
-        runtimeOnly project(":engines:tensorflow:tensorflow-model-zoo")
-        runtimeOnly project(":engines:mxnet:mxnet-model-zoo")
-        runtimeOnly project(":engines:paddlepaddle:paddlepaddle-model-zoo")
-
-        runtimeOnly project(":engines:tflite:tflite-engine")
-        runtimeOnly project(":engines:tensorrt")
-        ProcessBuilder pb = new ProcessBuilder("nvidia-smi", "-L")
-        def hasGPU = false;
-        try {
-            Process process = pb.start()
-            hasGPU = process.waitFor() == 0
-        } catch (IOException ignore) {
-        }
-
-        if (hasGPU) {
-            runtimeOnly(project(":engines:onnxruntime:onnxruntime-engine")) {
-                exclude group: "com.microsoft.onnxruntime", module: "onnxruntime"
-            }
-            runtimeOnly "com.microsoft.onnxruntime:onnxruntime_gpu:${onnxruntime_version}"
-        } else {
-            runtimeOnly project(":engines:onnxruntime:onnxruntime-engine")
-        }
-
-        runtimeOnly project(":engines:dlr:dlr-engine")
-        runtimeOnly(project(":engines:ml:xgboost")) {
-            exclude group: "ml.dmlc", module: "xgboost4j_2.12"
-        }
-    }
-
-    testImplementation("org.testng:testng:${testng_version}") {
-        exclude group: "junit", module: "junit"
-    }
-}
-
-application {
-    mainClass = System.getProperty("main", "ai.djl.benchmark.Benchmark")
-}
-
-run {
-    environment("TF_CPP_MIN_LOG_LEVEL", "1") // turn off TensorFlow print out
-    systemProperties System.getProperties()
-    systemProperties.remove("user.dir")
-    systemProperty("file.encoding", "UTF-8")
-}
-
-task benchmark(type: JavaExec) {
-    environment("TF_CPP_MIN_LOG_LEVEL", "1") // turn off TensorFlow print out
-    List<String> arguments = gradle.startParameter["taskRequests"]["args"].getAt(0)
-    for (String argument : arguments) {
-        if (argument.trim().startsWith("--args")) {
-            String[] line = argument.split("=", 2)
-            if (line.length == 2) {
-                line = line[1].split(" ")
-                if (line.contains("-t")) {
-                    if (System.getProperty("ai.djl.default_engine") == "TensorFlow") {
-                        environment("OMP_NUM_THREADS", "1")
-                        environment("TF_NUM_INTRAOP_THREADS", "1")
-                    } else {
-                        environment("MXNET_ENGINE_TYPE", "NaiveEngine")
-                        environment("OMP_NUM_THREADS", "1")
-                    }
-                }
-                break
-            }
-        }
-    }
-
-    systemProperties System.getProperties()
-    systemProperties.remove("user.dir")
-    systemProperty("file.encoding", "UTF-8")
-    classpath = sourceSets.main.runtimeClasspath
-    // restrict the jvm heap size for better monitoring benchmark
-    jvmArgs = ["-Xmx2g"]
-    if (Boolean.getBoolean("loggc")) {
-        if (JavaVersion.current() == JavaVersion.VERSION_1_8) {
-            jvmArgs += ["-XX:+PrintGCTimeStamps", "-Xloggc:build/gc.log"]
-        } else {
-            jvmArgs += ["-Xlog:gc*=debug:file=build/gc.log"]
-        }
-    }
-    mainClass = "ai.djl.benchmark.Benchmark"
-}
-
-task createDeb(type: Deb, dependsOn: distTar) {
-    doFirst {
-        exec {
-            commandLine "tar", "xvf", "${project.buildDir}/distributions/benchmark-${project.version}.tar", "-C", "${project.buildDir}"
-        }
-    }
-
-    packageName = "djl-bench"
-    archiveVersion = "${djl_version}"
-    release = 1
-    maintainer = "Deep Java Library <djl-dev@amazon.com>"
-    summary = "djl-bench is a command line tool that allows you to benchmark the\n" +
-            "  model on all different platforms for single-thread/multi-thread\n" +
-            "  inference performance."
-
-    from("${project.buildDir}/benchmark-${project.version}") {
-        into "/usr/local/djl-bench-${djl_version}"
-    }
-    link("/usr/bin/djl-bench", "/usr/local/djl-bench-${djl_version}/bin/benchmark")
-}
-
-startScripts {
-    defaultJvmOpts = []
-    doLast {
-        String replacement = 'CLASSPATH=\\$APP_HOME/lib/*\n\n' +
-                'if [[ "\\$*" == *-t* || "\\$*" == *--threads* ]]\n' +
-                'then\n' +
-                '    export TF_CPP_MIN_LOG_LEVEL=1\n' +
-                '    export MXNET_ENGINE_TYPE=NaiveEngine\n' +
-                '    export OMP_NUM_THREADS=1\n' +
-                '    export TF_NUM_INTRAOP_THREADS=1\n' +
-                'fi'
-
-        String text = unixScript.text.replaceAll('CLASSPATH=\\$APP_HOME/lib/.*', replacement)
-        text = text.replaceAll("/usr/bin/env sh", "/usr/bin/env bash")
-        text = text.replaceAll("#!/bin/sh", "#!/bin/bash")
-
-        unixScript.text = text
-    }
-}
diff --git a/extensions/benchmark/gradle b/extensions/benchmark/gradle
deleted file mode 120000
index 1ce6c4c1e..000000000
--- a/extensions/benchmark/gradle
+++ /dev/null
@@ -1,1 +0,0 @@
-../../gradle
\ No newline at end of file
diff --git a/extensions/benchmark/gradlew b/extensions/benchmark/gradlew
deleted file mode 120000
index 343e0d2ca..000000000
--- a/extensions/benchmark/gradlew
+++ /dev/null
@@ -1,1 +0,0 @@
-../../gradlew
\ No newline at end of file
diff --git a/extensions/benchmark/snapcraft/snapcraft.yaml b/extensions/benchmark/snapcraft/snapcraft.yaml
deleted file mode 100644
index afa85c747..000000000
--- a/extensions/benchmark/snapcraft/snapcraft.yaml
+++ /dev/null
@@ -1,44 +0,0 @@
-name: djlbench
-version: '0.17.0'
-title: DJL Benhmark
-license: Apache-2.0
-summary: A machine learning benchmarking toolkit
-description: |
-  djlbench is a command line tool that allows you to benchmark the
-  model on all different platforms for single-thread/multi-thread
-  inference performance.
-
-  Currently djlbench support the models from the following framework:
-  - PyTorch
-  - TensorFlow
-  - Apachmark MXNet
-  - PaddlePaddle
-  - ONNXRuntime
-  - TensorRT
-  - TensorFlow Lite
-  - Neo DLR
-  - XGBoost
-  - Python
-
-base: core18
-grade: stable
-confinement: classic
-
-apps:
-  djlbench:
-    command: benchmark-$SNAPCRAFT_PROJECT_VERSION/bin/benchmark
-    environment:
-      JAVA_HOME: "$SNAP/usr/lib/jvm/java-11-openjdk-amd64"
-      PATH: "$SNAP/bin:$PATH:$SNAP/usr/lib/jvm/java-11-openjdk-amd64/bin"
-
-parts:
-  djlbench:
-    plugin: gradle
-    source: https://github.com/deepjavalibrary/djl.git
-    source-tag: v$SNAPCRAFT_PROJECT_VERSION
-    gradle-output-dir: extensions/benchmark/build/libs
-    gradle-options: [ -Pstaging, ':extensions:benchmark:dT' ]
-    override-build: |
-      snapcraftctl build
-      tar xvf $SNAPCRAFT_PART_BUILD/extensions/benchmark/build/distributions/benchmark-*.tar -C $SNAPCRAFT_PART_INSTALL/
-      rm -rf $SNAPCRAFT_PART_INSTALL/jar
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java
deleted file mode 100644
index 8ef27cb1f..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/AbstractBenchmark.java
+++ /dev/null
@@ -1,299 +0,0 @@
-/*
- * Copyright 2019 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.benchmark;
-
-import ai.djl.Device;
-import ai.djl.ModelException;
-import ai.djl.engine.Engine;
-import ai.djl.metric.Metrics;
-import ai.djl.metric.Unit;
-import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.types.DataType;
-import ai.djl.ndarray.types.Shape;
-import ai.djl.repository.zoo.Criteria;
-import ai.djl.repository.zoo.ZooModel;
-import ai.djl.training.listener.MemoryTrainingListener;
-import ai.djl.training.util.ProgressBar;
-import ai.djl.translate.NoBatchifyTranslator;
-import ai.djl.translate.TranslateException;
-import ai.djl.translate.TranslatorContext;
-import ai.djl.util.Pair;
-import ai.djl.util.PairList;
-
-import org.apache.commons.cli.CommandLine;
-import org.apache.commons.cli.DefaultParser;
-import org.apache.commons.cli.Options;
-import org.apache.commons.cli.ParseException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.nio.FloatBuffer;
-import java.time.Duration;
-
-/** Abstract benchmark class. */
-public abstract class AbstractBenchmark {
-
-    private static final Logger logger = LoggerFactory.getLogger(AbstractBenchmark.class);
-
-    protected ProgressBar progressBar;
-
-    /**
-     * Abstract predict method that must be implemented by sub class.
-     *
-     * @param arguments command line arguments
-     * @param metrics {@link Metrics} to collect statistic information
-     * @param iteration number of prediction iteration to run
-     * @return prediction result
-     * @throws IOException if io error occurs when loading model.
-     * @throws ModelException if specified model not found or there is a parameter error
-     * @throws TranslateException if error occurs when processing input or output
-     * @throws ClassNotFoundException if input or output class cannot be loaded
-     */
-    protected abstract float[] predict(Arguments arguments, Metrics metrics, int iteration)
-            throws IOException, ModelException, TranslateException, ClassNotFoundException;
-
-    /**
-     * Execute benchmark.
-     *
-     * @param args input raw arguments
-     * @return if example execution complete successfully
-     */
-    public final boolean runBenchmark(String[] args) {
-        Options options = Arguments.getOptions();
-        try {
-            if (Arguments.hasHelp(args)) {
-                Arguments.printHelp(
-                        "usage: djl-bench [-p MODEL-PATH] -s INPUT-SHAPES [OPTIONS]", options);
-                return true;
-            }
-            DefaultParser parser = new DefaultParser();
-            CommandLine cmd = parser.parse(options, args, null, false);
-            Arguments arguments = new Arguments(cmd);
-            String engineName = arguments.getEngine();
-            Engine engine = Engine.getEngine(engineName);
-
-            long init = System.nanoTime();
-            String version = engine.getVersion();
-            long loaded = System.nanoTime();
-            logger.info(
-                    String.format(
-                            "Load %s (%s) in %.3f ms.",
-                            engineName, version, (loaded - init) / 1_000_000f));
-            Duration duration = Duration.ofSeconds(arguments.getDuration());
-            Object devices;
-            if (this instanceof MultithreadedBenchmark) {
-                devices = engine.getDevices(arguments.getMaxGpus());
-            } else {
-                devices = engine.defaultDevice();
-            }
-
-            if (arguments.getDuration() != 0) {
-                logger.info(
-                        "Running {} on: {}, duration: {} minutes.",
-                        getClass().getSimpleName(),
-                        devices,
-                        duration.toMinutes());
-            } else {
-                logger.info("Running {} on: {}.", getClass().getSimpleName(), devices);
-            }
-            int numOfThreads = arguments.getThreads();
-            int iteration = arguments.getIteration();
-            if (this instanceof MultithreadedBenchmark) {
-                int expected = 10 * numOfThreads;
-                if (iteration < expected) {
-                    iteration = expected;
-                    logger.info(
-                            "Iteration is too small for multi-threading benchmark. Adjust to: {}",
-                            iteration);
-                }
-            }
-            while (!duration.isNegative()) {
-                Metrics metrics = new Metrics(); // Reset Metrics for each test loop.
-                progressBar = new ProgressBar("Iteration", iteration);
-                float[] lastResult = predict(arguments, metrics, iteration);
-                if (lastResult == null) {
-                    return false;
-                }
-
-                long begin = metrics.getMetric("start").get(0).getValue().longValue();
-                long end = metrics.getMetric("end").get(0).getValue().longValue();
-                long totalTime = end - begin;
-
-                if (lastResult.length > 3) {
-                    logger.info(
-                            "Inference result: [{}, {}, {} ...]",
-                            lastResult[0],
-                            lastResult[1],
-                            lastResult[2]);
-                } else {
-                    logger.info("Inference result: {}", lastResult);
-                }
-
-                String throughput = String.format("%.2f", iteration * 1000d / totalTime);
-                logger.info(
-                        "Throughput: {}, completed {} iteration in {} ms.",
-                        throughput,
-                        iteration,
-                        totalTime);
-
-                if (metrics.hasMetric("LoadModel")) {
-                    long loadModelTime =
-                            metrics.getMetric("LoadModel").get(0).getValue().longValue();
-                    logger.info(
-                            "Model loading time: {} ms.",
-                            String.format("%.3f", loadModelTime / 1000f));
-                }
-
-                if (metrics.hasMetric("Inference") && iteration > 1) {
-                    float totalP50 = metrics.percentile("Total", 50).getValue().longValue() / 1000f;
-                    float totalP90 = metrics.percentile("Total", 90).getValue().longValue() / 1000f;
-                    float totalP99 = metrics.percentile("Total", 99).getValue().longValue() / 1000f;
-                    float p50 = metrics.percentile("Inference", 50).getValue().longValue() / 1000f;
-                    float p90 = metrics.percentile("Inference", 90).getValue().longValue() / 1000f;
-                    float p99 = metrics.percentile("Inference", 99).getValue().longValue() / 1000f;
-                    float preP50 =
-                            metrics.percentile("Preprocess", 50).getValue().longValue() / 1000f;
-                    float preP90 =
-                            metrics.percentile("Preprocess", 90).getValue().longValue() / 1000f;
-                    float preP99 =
-                            metrics.percentile("Preprocess", 99).getValue().longValue() / 1000f;
-                    float postP50 =
-                            metrics.percentile("Postprocess", 50).getValue().longValue() / 1000f;
-                    float postP90 =
-                            metrics.percentile("Postprocess", 90).getValue().longValue() / 1000f;
-                    float postP99 =
-                            metrics.percentile("Postprocess", 99).getValue().longValue() / 1000f;
-                    logger.info(
-                            String.format(
-                                    "total P50: %.3f ms, P90: %.3f ms, P99: %.3f ms",
-                                    totalP50, totalP90, totalP99));
-                    logger.info(
-                            String.format(
-                                    "inference P50: %.3f ms, P90: %.3f ms, P99: %.3f ms",
-                                    p50, p90, p99));
-                    logger.info(
-                            String.format(
-                                    "preprocess P50: %.3f ms, P90: %.3f ms, P99: %.3f ms",
-                                    preP50, preP90, preP99));
-                    logger.info(
-                            String.format(
-                                    "postprocess P50: %.3f ms, P90: %.3f ms, P99: %.3f ms",
-                                    postP50, postP90, postP99));
-
-                    if (Boolean.getBoolean("collect-memory")) {
-                        float heapBeforeModel =
-                                metrics.getMetric("Heap").get(0).getValue().longValue();
-                        float heapBeforeInference =
-                                metrics.getMetric("Heap").get(1).getValue().longValue();
-                        float heap = metrics.percentile("Heap", 90).getValue().longValue();
-                        float nonHeap = metrics.percentile("NonHeap", 90).getValue().longValue();
-                        int mb = 1024 * 1024;
-                        logger.info(String.format("heap (base): %.3f MB", heapBeforeModel / mb));
-                        logger.info(
-                                String.format("heap (model): %.3f MB", heapBeforeInference / mb));
-                        logger.info(String.format("heap P90: %.3f MB", heap / mb));
-                        logger.info(String.format("nonHeap P90: %.3f MB", nonHeap / mb));
-
-                        if (!System.getProperty("os.name").startsWith("Win")) {
-                            float rssBeforeModel =
-                                    metrics.getMetric("rss").get(0).getValue().longValue();
-                            float rssBeforeInference =
-                                    metrics.getMetric("rss").get(1).getValue().longValue();
-                            float rss = metrics.percentile("rss", 90).getValue().longValue();
-                            float cpu = metrics.percentile("cpu", 90).getValue().longValue();
-                            logger.info(String.format("cpu P90: %.3f %%", cpu));
-                            logger.info(String.format("rss (base): %.3f MB", rssBeforeModel / mb));
-                            logger.info(
-                                    String.format("rss (model): %.3f MB", rssBeforeInference / mb));
-                            logger.info(String.format("rss P90: %.3f MB", rss / mb));
-                        }
-                    }
-                }
-                MemoryTrainingListener.dumpMemoryInfo(metrics, arguments.getOutputDir());
-                long delta = System.currentTimeMillis() - begin;
-                duration = duration.minus(Duration.ofMillis(delta));
-                if (!duration.isNegative()) {
-                    logger.info(duration.toMinutes() + " minutes left");
-                }
-            }
-            return true;
-        } catch (ParseException e) {
-            Arguments.printHelp(e.getMessage(), options);
-        } catch (TranslateException | ModelException | IOException | ClassNotFoundException t) {
-            logger.error("Unexpected error", t);
-        }
-        return false;
-    }
-
-    protected ZooModel<Void, float[]> loadModel(Arguments arguments, Metrics metrics, Device device)
-            throws ModelException, IOException {
-        long begin = System.nanoTime();
-        PairList<DataType, Shape> shapes = arguments.getInputShapes();
-        BenchmarkTranslator translator = new BenchmarkTranslator(shapes);
-
-        Criteria<Void, float[]> criteria =
-                Criteria.builder()
-                        .setTypes(Void.class, float[].class)
-                        .optModelUrls(arguments.getModelUrl())
-                        .optModelName(arguments.getModelName())
-                        .optEngine(arguments.getEngine())
-                        .optOptions(arguments.getModelOptions())
-                        .optArguments(arguments.getModelArguments())
-                        .optDevice(device)
-                        .optTranslator(translator)
-                        .optProgress(new ProgressBar())
-                        .build();
-
-        ZooModel<Void, float[]> model = criteria.loadModel();
-        if (device == Device.cpu() || device == Device.gpu()) {
-            long delta = (System.nanoTime() - begin) / 1000;
-            logger.info(
-                    "Model {} loaded in: {} ms.",
-                    model.getName(),
-                    String.format("%.3f", delta / 1000f));
-            metrics.addMetric("LoadModel", delta, Unit.MICROSECONDS);
-        }
-        return model;
-    }
-
-    private static final class BenchmarkTranslator implements NoBatchifyTranslator<Void, float[]> {
-
-        private PairList<DataType, Shape> shapes;
-
-        public BenchmarkTranslator(PairList<DataType, Shape> shapes) {
-            this.shapes = shapes;
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public NDList processInput(TranslatorContext ctx, Void input) {
-            NDList list = new NDList();
-            for (Pair<DataType, Shape> pair : shapes) {
-                DataType dataType = pair.getKey();
-                Shape shape = pair.getValue();
-                list.add(ctx.getNDManager().zeros(shape, dataType));
-            }
-            return list;
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public float[] processOutput(TranslatorContext ctx, NDList list) {
-            FloatBuffer fb = list.get(0).toByteBuffer().asFloatBuffer();
-            float[] ret = new float[fb.remaining()];
-            fb.get(ret);
-            return ret;
-        }
-    }
-}
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/Arguments.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/Arguments.java
deleted file mode 100644
index 4a96e934f..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/Arguments.java
+++ /dev/null
@@ -1,340 +0,0 @@
-/*
- * Copyright 2019 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.benchmark;
-
-import ai.djl.Device;
-import ai.djl.engine.Engine;
-import ai.djl.ndarray.types.DataType;
-import ai.djl.ndarray.types.Shape;
-import ai.djl.util.PairList;
-
-import org.apache.commons.cli.CommandLine;
-import org.apache.commons.cli.HelpFormatter;
-import org.apache.commons.cli.Option;
-import org.apache.commons.cli.OptionGroup;
-import org.apache.commons.cli.Options;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
-/** A class represents parsed command line arguments. */
-public class Arguments {
-
-    private static final Logger logger = LoggerFactory.getLogger(Arguments.class);
-
-    private String modelUrl;
-    private String modelName;
-    private String engine;
-    private String modelOptions;
-    private String modelArguments;
-    private String outputDir;
-    private int duration;
-    private int iteration;
-    private int threads;
-    private int maxGpus;
-    private int neuronCores;
-    private int delay;
-    private PairList<DataType, Shape> inputShapes;
-
-    /**
-     * Constructs a {@code Arguments} instance.
-     *
-     * @param cmd command line options
-     */
-    Arguments(CommandLine cmd) {
-        if (cmd.hasOption("model-path")) {
-            String modelPath = cmd.getOptionValue("model-path");
-            Path path = Paths.get(modelPath);
-            try {
-                modelUrl = path.toUri().toURL().toExternalForm();
-            } catch (IOException e) {
-                throw new IllegalArgumentException("Invalid model-path: " + modelPath, e);
-            }
-        } else if (cmd.hasOption("model-url")) {
-            modelUrl = cmd.getOptionValue("model-url");
-        }
-
-        modelName = cmd.getOptionValue("model-name");
-        modelOptions = cmd.getOptionValue("model-options");
-        modelArguments = cmd.getOptionValue("model-arguments");
-        outputDir = cmd.getOptionValue("output-dir");
-
-        if (cmd.hasOption("engine")) {
-            engine = cmd.getOptionValue("engine");
-        } else {
-            engine = Engine.getDefaultEngineName();
-        }
-
-        if (cmd.hasOption("duration")) {
-            duration = Integer.parseInt(cmd.getOptionValue("duration"));
-        }
-        iteration = 1;
-        if (cmd.hasOption("iteration")) {
-            iteration = Integer.parseInt(cmd.getOptionValue("iteration"));
-        }
-        if (cmd.hasOption("gpus")) {
-            maxGpus = Integer.parseInt(cmd.getOptionValue("gpus"));
-            if (maxGpus < 0) {
-                maxGpus = Integer.MAX_VALUE;
-            }
-        } else {
-            maxGpus = Integer.MAX_VALUE;
-        }
-        if (cmd.hasOption("neuron-cores")) {
-            neuronCores = Integer.parseInt(cmd.getOptionValue("neuron-cores"));
-        }
-        if (cmd.hasOption("threads")) {
-            threads = Integer.parseInt(cmd.getOptionValue("threads"));
-            Engine eng = Engine.getEngine(engine);
-            Device[] devices = eng.getDevices(maxGpus);
-            if (devices[0].isGpu()) {
-                // one thread per GPU
-                if (threads <= 0) {
-                    threads = devices.length;
-                } else if (threads < devices.length) {
-                    threads = devices.length;
-                    logger.warn(
-                            "Number of threads is less than GPU count, adjust to: {}",
-                            devices.length);
-                } else if ("MXNet".equals(engine) && threads > devices.length) {
-                    threads = devices.length;
-                    logger.warn("MXNet inference can only have one worker per GPU.");
-                } else if (threads % devices.length != 0) {
-                    threads = threads / devices.length * devices.length;
-                    logger.warn("threads should be multiple of GPU count, change to: {}", threads);
-                }
-            } else if (threads <= 0) {
-                threads = Runtime.getRuntime().availableProcessors();
-            }
-        }
-        if (cmd.hasOption("delay")) {
-            delay = Integer.parseInt(cmd.getOptionValue("delay"));
-        }
-
-        String shape = cmd.getOptionValue("input-shapes");
-        inputShapes = NDListGenerator.parseShape(shape);
-    }
-
-    static Options getOptions() {
-        Options options = new Options();
-        options.addOption(
-                Option.builder("h").longOpt("help").hasArg(false).desc("Print this help.").build());
-        OptionGroup artifactGroup = new OptionGroup();
-        artifactGroup.setRequired(true);
-        artifactGroup.addOption(
-                Option.builder("p")
-                        .longOpt("model-path")
-                        .hasArg()
-                        .argName("MODEL-PATH")
-                        .desc("Model directory file path.")
-                        .build());
-        artifactGroup.addOption(
-                Option.builder("u")
-                        .longOpt("model-url")
-                        .hasArg()
-                        .argName("MODEL-URL")
-                        .desc("Model archive file URL.")
-                        .build());
-        options.addOptionGroup(artifactGroup);
-        options.addOption(
-                Option.builder("n")
-                        .longOpt("model-name")
-                        .hasArg()
-                        .argName("MODEL-NAME")
-                        .desc("Specify model file name.")
-                        .build());
-        options.addOption(
-                Option.builder()
-                        .longOpt("model-options")
-                        .hasArg()
-                        .argName("MODEL-OPTIONS")
-                        .desc("Specify model loading options.")
-                        .build());
-        options.addOption(
-                Option.builder()
-                        .longOpt("model-arguments")
-                        .hasArg()
-                        .argName("MODEL-ARGUMENTS")
-                        .desc("Specify model loading arguments.")
-                        .build());
-        options.addOption(
-                Option.builder("e")
-                        .longOpt("engine")
-                        .hasArg()
-                        .argName("ENGINE-NAME")
-                        .desc("Choose an Engine for the benchmark.")
-                        .build());
-        options.addOption(
-                Option.builder("s")
-                        .required()
-                        .longOpt("input-shapes")
-                        .hasArg()
-                        .argName("INPUT-SHAPES")
-                        .desc("Input data shapes for the model.")
-                        .build());
-        options.addOption(
-                Option.builder("d")
-                        .longOpt("duration")
-                        .hasArg()
-                        .argName("DURATION")
-                        .desc("Duration of the test in minutes.")
-                        .build());
-        options.addOption(
-                Option.builder("c")
-                        .longOpt("iteration")
-                        .hasArg()
-                        .argName("ITERATION")
-                        .desc("Number of total iterations.")
-                        .build());
-        options.addOption(
-                Option.builder("t")
-                        .longOpt("threads")
-                        .hasArg()
-                        .argName("NUMBER_THREADS")
-                        .desc("Number of inference threads.")
-                        .build());
-        OptionGroup deviceGroup = new OptionGroup();
-        deviceGroup.addOption(
-                Option.builder("g")
-                        .longOpt("gpus")
-                        .hasArg()
-                        .argName("NUMBER_GPUS")
-                        .desc("Number of GPUS to run multithreading inference.")
-                        .build());
-        deviceGroup.addOption(
-                Option.builder()
-                        .longOpt("neuron-cores")
-                        .hasArg()
-                        .argName("NEURON-CORES")
-                        .desc(
-                                "Number of neuron cores to run multithreading inference, See"
-                                        + " https://awsdocs-neuron.readthedocs-hosted.com.")
-                        .build());
-        options.addOptionGroup(deviceGroup);
-        options.addOption(
-                Option.builder("l")
-                        .longOpt("delay")
-                        .hasArg()
-                        .argName("DELAY")
-                        .desc("Delay of incremental threads.")
-                        .build());
-        options.addOption(
-                Option.builder("o")
-                        .longOpt("output-dir")
-                        .hasArg()
-                        .argName("OUTPUT-DIR")
-                        .desc("Directory for output logs.")
-                        .build());
-        return options;
-    }
-
-    static boolean hasHelp(String[] args) {
-        List<String> list = Arrays.asList(args);
-        return list.contains("-h") || list.contains("--help");
-    }
-
-    static void printHelp(String msg, Options options) {
-        HelpFormatter formatter = new HelpFormatter();
-        formatter.setSyntaxPrefix("");
-        formatter.setLeftPadding(1);
-        formatter.setWidth(120);
-        formatter.printHelp(msg, options);
-    }
-
-    int getDuration() {
-        return duration;
-    }
-
-    String getEngine() {
-        return engine;
-    }
-
-    String getModelUrl() {
-        return modelUrl;
-    }
-
-    String getModelName() {
-        return modelName;
-    }
-
-    Map<String, String> getModelOptions() {
-        if (modelOptions == null) {
-            return null;
-        }
-        Map<String, String> map = new ConcurrentHashMap<>();
-        for (String option : modelOptions.split(",")) {
-            String[] tokens = option.split("=", 2);
-            if (tokens.length == 2) {
-                map.put(tokens[0].trim(), tokens[1].trim());
-            } else {
-                map.put(tokens[0].trim(), "");
-            }
-        }
-        return map;
-    }
-
-    Map<String, Object> getModelArguments() {
-        if (modelArguments == null) {
-            return null;
-        }
-
-        Map<String, Object> map = new ConcurrentHashMap<>();
-        for (String option : modelArguments.split(",")) {
-            String[] tokens = option.split("=", 2);
-            if (tokens.length == 2) {
-                map.put(tokens[0].trim(), tokens[1].trim());
-            } else {
-                map.put(tokens[0].trim(), "");
-            }
-        }
-        return map;
-    }
-
-    int getIteration() {
-        return iteration;
-    }
-
-    int getThreads() {
-        return threads;
-    }
-
-    int getMaxGpus() {
-        return maxGpus;
-    }
-
-    int getNeuronCores() {
-        return neuronCores;
-    }
-
-    String getOutputDir() {
-        if (outputDir == null) {
-            outputDir = "build";
-        }
-        return outputDir;
-    }
-
-    int getDelay() {
-        return delay;
-    }
-
-    PairList<DataType, Shape> getInputShapes() {
-        return inputShapes;
-    }
-}
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java
deleted file mode 100644
index 3fad6c377..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/Benchmark.java
+++ /dev/null
@@ -1,119 +0,0 @@
-/*
- * Copyright 2019 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.benchmark;
-
-import ai.djl.Device;
-import ai.djl.ModelException;
-import ai.djl.engine.Engine;
-import ai.djl.engine.EngineException;
-import ai.djl.inference.Predictor;
-import ai.djl.metric.Metrics;
-import ai.djl.metric.Unit;
-import ai.djl.repository.zoo.ZooModel;
-import ai.djl.training.listener.MemoryTrainingListener;
-import ai.djl.translate.TranslateException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.List;
-
-/** A class runs single threaded benchmark. */
-public final class Benchmark extends AbstractBenchmark {
-
-    private static final Logger logger = LoggerFactory.getLogger(Benchmark.class);
-
-    /**
-     * Main entry point.
-     *
-     * @param args command line arguments
-     */
-    public static void main(String[] args) {
-        List<String> list = Arrays.asList(args);
-        try {
-            boolean success;
-            if (!list.isEmpty() && "ndlist-gen".equals(list.get(0))) {
-                success = NDListGenerator.generate(Arrays.copyOfRange(args, 1, args.length));
-            } else {
-                boolean multithreading = list.contains("-t") || list.contains("--threads");
-                configEngines(multithreading);
-                if (multithreading) {
-                    success = new MultithreadedBenchmark().runBenchmark(args);
-                } else {
-                    success = new Benchmark().runBenchmark(args);
-                }
-            }
-            if (!success) {
-                System.exit(-1); // NOPMD
-            }
-        } catch (EngineException e) {
-            String osName = System.getProperty("os.name");
-            String arch = System.getProperty("os.arch");
-            logger.warn("Engine is not supported on {}:{}.", osName, arch);
-            logger.debug("Failed to load engine", e);
-        }
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public float[] predict(Arguments arguments, Metrics metrics, int iteration)
-            throws IOException, ModelException, TranslateException {
-        Device device = Engine.getEngine(arguments.getEngine()).defaultDevice();
-        try (ZooModel<Void, float[]> model = loadModel(arguments, metrics, device)) {
-            float[] predictResult = null;
-
-            try (Predictor<Void, float[]> predictor = model.newPredictor()) {
-                predictor.predict(null); // warmup
-
-                predictor.setMetrics(metrics); // Let predictor collect metrics
-                metrics.addMetric("start", System.currentTimeMillis(), Unit.MILLISECONDS);
-                for (int i = 0; i < iteration; ++i) {
-                    predictResult = predictor.predict(null);
-
-                    progressBar.update(i);
-                    MemoryTrainingListener.collectMemoryInfo(metrics);
-                }
-                metrics.addMetric("end", System.currentTimeMillis(), Unit.MILLISECONDS);
-            }
-            return predictResult;
-        }
-    }
-
-    private static void configEngines(boolean multithreading) {
-        if (multithreading) {
-            if (System.getProperty("ai.djl.pytorch.num_interop_threads") == null) {
-                System.setProperty("ai.djl.pytorch.num_interop_threads", "1");
-            }
-            if (System.getProperty("ai.djl.pytorch.num_threads") == null) {
-                System.setProperty("ai.djl.pytorch.num_threads", "1");
-            }
-        }
-        if (System.getProperty("ai.djl.tflite.disable_alternative") == null) {
-            System.setProperty("ai.djl.tflite.disable_alternative", "true");
-        }
-        if (System.getProperty("ai.djl.dlr.disable_alternative") == null) {
-            System.setProperty("ai.djl.dlr.disable_alternative", "true");
-        }
-        if (System.getProperty("ai.djl.paddlepaddle.disable_alternative") == null) {
-            System.setProperty("ai.djl.paddlepaddle.disable_alternative", "true");
-        }
-        if (System.getProperty("ai.djl.onnx.disable_alternative") == null) {
-            System.setProperty("ai.djl.onnx.disable_alternative", "true");
-        }
-        if (System.getProperty("ai.djl.tensorrt.disable_alternative") == null) {
-            System.setProperty("ai.djl.tensorrt.disable_alternative", "true");
-        }
-    }
-}
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java
deleted file mode 100644
index 438b176d4..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/MultithreadedBenchmark.java
+++ /dev/null
@@ -1,194 +0,0 @@
-/*
- * Copyright 2019 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.benchmark;
-
-import ai.djl.Device;
-import ai.djl.ModelException;
-import ai.djl.engine.Engine;
-import ai.djl.inference.Predictor;
-import ai.djl.metric.Metrics;
-import ai.djl.metric.Unit;
-import ai.djl.repository.zoo.ZooModel;
-import ai.djl.training.listener.MemoryTrainingListener;
-import ai.djl.translate.TranslateException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-import java.util.concurrent.Callable;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.atomic.AtomicInteger;
-
-/** A class runs single threaded benchmark. */
-public class MultithreadedBenchmark extends AbstractBenchmark {
-
-    private static final Logger logger = LoggerFactory.getLogger(MultithreadedBenchmark.class);
-
-    /** {@inheritDoc} */
-    @Override
-    public float[] predict(Arguments arguments, Metrics metrics, int iteration)
-            throws IOException, ModelException, TranslateException {
-
-        MemoryTrainingListener.collectMemoryInfo(metrics); // Measure memory before loading model
-
-        Engine engine = Engine.getEngine(arguments.getEngine());
-        Device[] devices = engine.getDevices(arguments.getMaxGpus());
-        int numOfThreads = arguments.getThreads();
-        int neuronCores = arguments.getNeuronCores();
-        if (neuronCores > 0) {
-            devices = new Device[neuronCores];
-            Arrays.fill(devices, Device.cpu());
-            if (numOfThreads > 1) {
-                numOfThreads = 2 * neuronCores;
-            }
-        }
-
-        int delay = arguments.getDelay();
-        AtomicInteger counter = new AtomicInteger(iteration);
-        logger.info("Multithreading inference with {} threads.", numOfThreads);
-
-        List<ZooModel<Void, float[]>> models = new ArrayList<>(devices.length);
-        List<PredictorCallable> callables = new ArrayList<>(numOfThreads);
-        for (Device device : devices) {
-            ZooModel<Void, float[]> model = loadModel(arguments, metrics, device);
-            models.add(model);
-
-            for (int i = 0; i < numOfThreads / devices.length; ++i) {
-                callables.add(new PredictorCallable(model, metrics, counter, i, i == 0));
-            }
-        }
-
-        float[] result = null;
-        ExecutorService executorService = Executors.newFixedThreadPool(numOfThreads);
-
-        MemoryTrainingListener.collectMemoryInfo(metrics); // Measure memory before worker kickoff
-
-        int successThreads = 0;
-        try {
-            for (PredictorCallable callable : callables) {
-                callable.warmup();
-            }
-
-            metrics.addMetric("start", System.currentTimeMillis(), Unit.MILLISECONDS);
-            try {
-                List<Future<float[]>> futures;
-                if (delay > 0) {
-                    futures = new ArrayList<>();
-                    for (PredictorCallable callable : callables) {
-                        futures.add(executorService.submit(callable));
-                        Thread.sleep(delay);
-                    }
-                } else {
-                    futures = executorService.invokeAll(callables);
-                }
-
-                for (Future<float[]> future : futures) {
-                    result = future.get();
-                    if (result != null) {
-                        ++successThreads;
-                    }
-                }
-            } catch (InterruptedException | ExecutionException e) {
-                logger.error("", e);
-            }
-            metrics.addMetric("end", System.currentTimeMillis(), Unit.MILLISECONDS);
-            for (PredictorCallable callable : callables) {
-                callable.close();
-            }
-        } finally {
-            executorService.shutdown();
-        }
-
-        models.forEach(ZooModel::close);
-        if (successThreads != numOfThreads) {
-            logger.error("Only {}/{} threads finished.", successThreads, numOfThreads);
-            return null;
-        }
-
-        return result;
-    }
-
-    private static class PredictorCallable implements Callable<float[]> {
-
-        private Predictor<Void, float[]> predictor;
-
-        private Metrics metrics;
-        private String workerId;
-        private boolean collectMemory;
-        private AtomicInteger counter;
-        private int total;
-        private int steps;
-
-        public PredictorCallable(
-                ZooModel<Void, float[]> model,
-                Metrics metrics,
-                AtomicInteger counter,
-                int workerId,
-                boolean collectMemory) {
-            this.predictor = model.newPredictor();
-            this.metrics = metrics;
-            this.counter = counter;
-            this.workerId = String.format("%02d", workerId);
-            this.collectMemory = collectMemory;
-            predictor.setMetrics(metrics);
-            total = counter.get();
-            if (total < 10) {
-                steps = 1;
-            } else {
-                steps = (int) Math.pow(10, (int) Math.log10(total));
-            }
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public float[] call() throws Exception {
-            float[] result = null;
-            int count = 0;
-            int remaining;
-            while ((remaining = counter.decrementAndGet()) > 0 || result == null) {
-                try {
-                    result = predictor.predict(null);
-                } catch (Exception e) {
-                    // stop immediately when we find any exception
-                    counter.set(0);
-                    throw e;
-                }
-                if (collectMemory) {
-                    MemoryTrainingListener.collectMemoryInfo(metrics);
-                }
-                int processed = total - remaining + 1;
-                logger.trace("Worker-{}: {} iteration finished.", workerId, ++count);
-                if (processed % steps == 0 || processed == total) {
-                    logger.info("Completed {} requests", processed);
-                }
-            }
-            logger.debug("Worker-{}: finished.", workerId);
-            return result;
-        }
-
-        public void warmup() throws TranslateException {
-            predictor.predict(null);
-        }
-
-        public void close() {
-            predictor.close();
-        }
-    }
-}
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/NDListGenerator.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/NDListGenerator.java
deleted file mode 100644
index eb52178b6..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/NDListGenerator.java
+++ /dev/null
@@ -1,171 +0,0 @@
-/*
- * 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.benchmark;
-
-import ai.djl.Device;
-import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.DataType;
-import ai.djl.ndarray.types.Shape;
-import ai.djl.util.Pair;
-import ai.djl.util.PairList;
-
-import org.apache.commons.cli.CommandLine;
-import org.apache.commons.cli.DefaultParser;
-import org.apache.commons.cli.Option;
-import org.apache.commons.cli.Options;
-import org.apache.commons.cli.ParseException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.BufferedOutputStream;
-import java.io.OutputStream;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.Paths;
-import java.util.Arrays;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-/** A class generates NDList files. */
-final class NDListGenerator {
-
-    private static final Logger logger = LoggerFactory.getLogger(NDListGenerator.class);
-
-    private NDListGenerator() {}
-
-    static boolean generate(String[] args) {
-        Options options = getOptions();
-        try {
-            if (Arguments.hasHelp(args)) {
-                Arguments.printHelp(
-                        "usage: djl-bench ndlist-gen -s INPUT-SHAPES -o OUTPUT_FILE", options);
-                return true;
-            }
-            DefaultParser parser = new DefaultParser();
-            CommandLine cmd = parser.parse(options, args, null, false);
-            String inputShapes = cmd.getOptionValue("input-shapes");
-            String output = cmd.getOptionValue("output-file");
-            boolean ones = cmd.hasOption("ones");
-            Path path = Paths.get(output);
-
-            try (NDManager manager = NDManager.newBaseManager(Device.cpu(), "PyTorch")) {
-                NDList list = new NDList();
-                for (Pair<DataType, Shape> pair : parseShape(inputShapes)) {
-                    DataType dataType = pair.getKey();
-                    Shape shape = pair.getValue();
-                    if (ones) {
-                        list.add(manager.ones(shape, dataType));
-                    } else {
-                        list.add(manager.zeros(shape, dataType));
-                    }
-                }
-                try (OutputStream os = new BufferedOutputStream(Files.newOutputStream(path))) {
-                    list.encode(os);
-                }
-            }
-            logger.info("NDList file created: {}", path.toAbsolutePath());
-            return true;
-        } catch (ParseException e) {
-            Arguments.printHelp(e.getMessage(), options);
-        } catch (Throwable t) {
-            logger.error("Unexpected error", t);
-        }
-        return false;
-    }
-
-    static PairList<DataType, Shape> parseShape(String shape) {
-        PairList<DataType, Shape> inputShapes = new PairList<>();
-        if (shape != null) {
-            if (shape.contains("(")) {
-                Pattern pattern =
-                        Pattern.compile("\\((\\s*(\\d+)([,\\s]+\\d+)*\\s*)\\)([sdubilBfS]?)");
-                Matcher matcher = pattern.matcher(shape);
-                while (matcher.find()) {
-                    String[] tokens = matcher.group(1).split(",");
-                    long[] array = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
-                    DataType dataType;
-                    String dataTypeStr = matcher.group(4);
-                    if (dataTypeStr == null || dataTypeStr.isEmpty()) {
-                        dataType = DataType.FLOAT32;
-                    } else {
-                        switch (dataTypeStr) {
-                            case "s":
-                                dataType = DataType.FLOAT16;
-                                break;
-                            case "d":
-                                dataType = DataType.FLOAT64;
-                                break;
-                            case "u":
-                                dataType = DataType.UINT8;
-                                break;
-                            case "b":
-                                dataType = DataType.INT8;
-                                break;
-                            case "i":
-                                dataType = DataType.INT32;
-                                break;
-                            case "l":
-                                dataType = DataType.INT64;
-                                break;
-                            case "B":
-                                dataType = DataType.BOOLEAN;
-                                break;
-                            case "f":
-                                dataType = DataType.FLOAT32;
-                                break;
-                            default:
-                                throw new IllegalArgumentException("Invalid input-shape: " + shape);
-                        }
-                    }
-                    inputShapes.add(dataType, new Shape(array));
-                }
-            } else {
-                String[] tokens = shape.split(",");
-                long[] shapes = Arrays.stream(tokens).mapToLong(Long::parseLong).toArray();
-                inputShapes.add(DataType.FLOAT32, new Shape(shapes));
-            }
-        }
-        return inputShapes;
-    }
-
-    private static Options getOptions() {
-        Options options = new Options();
-        options.addOption(
-                Option.builder("h").longOpt("help").hasArg(false).desc("Print this help.").build());
-        options.addOption(
-                Option.builder("s")
-                        .required()
-                        .longOpt("input-shapes")
-                        .hasArg()
-                        .argName("INPUT-SHAPES")
-                        .desc("Input data shapes for the model.")
-                        .build());
-        options.addOption(
-                Option.builder("o")
-                        .required()
-                        .longOpt("output-file")
-                        .hasArg()
-                        .argName("OUTPUT-FILE")
-                        .desc("Write output NDList to file.")
-                        .build());
-        options.addOption(
-                Option.builder("1")
-                        .longOpt("ones")
-                        .hasArg(false)
-                        .argName("ones")
-                        .desc("Use all ones instead of zeros.")
-                        .build());
-        return options;
-    }
-}
diff --git a/extensions/benchmark/src/main/java/ai/djl/benchmark/package-info.java b/extensions/benchmark/src/main/java/ai/djl/benchmark/package-info.java
deleted file mode 100644
index 6436a24fe..000000000
--- a/extensions/benchmark/src/main/java/ai/djl/benchmark/package-info.java
+++ /dev/null
@@ -1,15 +0,0 @@
-/*
- * Copyright 2019 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.
- */
-
-/** Contains benchmarking utility classes. */
-package ai.djl.benchmark;
diff --git a/extensions/benchmark/src/main/resources/log4j2.xml b/extensions/benchmark/src/main/resources/log4j2.xml
deleted file mode 100644
index cb6d3c6fb..000000000
--- a/extensions/benchmark/src/main/resources/log4j2.xml
+++ /dev/null
@@ -1,20 +0,0 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<Configuration status="INFO">
-  <Appenders>
-    <Console name="console" target="SYSTEM_OUT">
-      <PatternLayout
-          pattern="[%-5level] - %msg%n"/>
-    </Console>
-  </Appenders>
-  <Loggers>
-    <Root level="info" additivity="false">
-      <AppenderRef ref="console"/>
-    </Root>
-    <Logger name="ai.djl" level="${sys:ai.djl.logging.level:-info}" additivity="false">
-      <AppenderRef ref="console"/>
-    </Logger>
-    <Logger name="ai.djl.repository.zoo" level="${sys:ai.djl.modelzoo.logging.level:-info}" additivity="false">
-      <AppenderRef ref="console"/>
-    </Logger>
-  </Loggers>
-</Configuration>
diff --git a/settings.gradle b/settings.gradle
index 859c5362e..412daf4fe 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -27,7 +27,6 @@ include ':engines:tflite:tflite-native'
 include ':examples'
 include 'extensions:audio'
 include ':extensions:aws-ai'
-include ':extensions:benchmark'
 include ':extensions:fasttext'
 include ':extensions:hadoop'
 include ':extensions:opencv'
diff --git a/tools/gradle/release.gradle b/tools/gradle/release.gradle
index 5645ad5fe..85662d8b1 100644
--- a/tools/gradle/release.gradle
+++ b/tools/gradle/release.gradle
@@ -59,7 +59,6 @@ task increaseFinalVersion {
         collection += fileTree(".").filter {
             it.name.endsWith(".md") || it.name.endsWith("overview.html")
         }
-        collection += file("extensions/benchmark/snapcraft/snapcraft.yaml")
 
         collection.each { File file ->
             file.text = file.text.replaceAll("/${previousVersion}/", "/${djl_version}/")


Test patch:
diff --git a/extensions/benchmark/src/test/java/ai/djl/benchmark/BenchmarkTest.java b/extensions/benchmark/src/test/java/ai/djl/benchmark/BenchmarkTest.java
deleted file mode 100644
index 3b53e6fc8..000000000
--- a/extensions/benchmark/src/test/java/ai/djl/benchmark/BenchmarkTest.java
+++ /dev/null
@@ -1,123 +0,0 @@
-/*
- * Copyright 2020 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.benchmark;
-
-import ai.djl.ndarray.types.DataType;
-
-import org.apache.commons.cli.CommandLine;
-import org.apache.commons.cli.DefaultParser;
-import org.apache.commons.cli.Options;
-import org.apache.commons.cli.ParseException;
-import org.testng.Assert;
-import org.testng.annotations.Test;
-
-import java.net.MalformedURLException;
-import java.nio.file.Paths;
-import java.util.Map;
-
-public class BenchmarkTest {
-
-    @Test
-    public void testHelp() {
-        String[] args = {"-h"};
-        Benchmark.main(args);
-    }
-
-    @Test
-    public void testArguments() throws ParseException, MalformedURLException {
-        Options options = Arguments.getOptions();
-        DefaultParser parser = new DefaultParser();
-
-        String[] args = {
-            "-p",
-            "/opt/ml/resnet18_v1",
-            "-s",
-            "(1)s,(1)d,(1)u,(1)b,(1)i,(1)l,(1)B,(1)",
-            "--model-options",
-            "fp16,dlaCore=1",
-            "--model-arguments",
-            "width=28"
-        };
-        CommandLine cmd = parser.parse(options, args, null, false);
-        Arguments arguments = new Arguments(cmd);
-        String expected = Paths.get("/opt/ml/resnet18_v1").toUri().toURL().toString();
-        Assert.assertEquals(arguments.getModelUrl(), expected);
-        DataType[] types = arguments.getInputShapes().keyArray(new DataType[0]);
-        Assert.assertEquals(types[0], DataType.FLOAT16);
-        Assert.assertEquals(types[1], DataType.FLOAT64);
-        Assert.assertEquals(types[2], DataType.UINT8);
-        Assert.assertEquals(types[3], DataType.INT8);
-        Assert.assertEquals(types[4], DataType.INT32);
-        Assert.assertEquals(types[5], DataType.INT64);
-        Assert.assertEquals(types[6], DataType.BOOLEAN);
-        Assert.assertEquals(types[7], DataType.FLOAT32);
-
-        Assert.assertThrows(
-                IllegalArgumentException.class,
-                () -> {
-                    String[] arg = {"-p", "/opt/ml/resnet18_v1", "-s", "(1)S"};
-                    CommandLine commandLine = parser.parse(options, arg, null, false);
-                    new Arguments(commandLine);
-                });
-
-        Map<String, String> map = arguments.getModelOptions();
-        Assert.assertEquals(map.get("dlaCore"), "1");
-        Assert.assertTrue(map.containsKey("fp16"));
-
-        Map<String, Object> modelArguments = arguments.getModelArguments();
-        Assert.assertEquals(modelArguments.get("width"), "28");
-    }
-
-    @Test
-    public void testBenchmark() {
-        String[] args = {
-            "-e",
-            "PyTorch",
-            "-u",
-            "djl://ai.djl.pytorch/resnet/0.0.1/traced_resnet18",
-            "-s",
-            "1,3,224,224",
-            "-c",
-            "2"
-        };
-        new Benchmark().runBenchmark(args);
-    }
-
-    @Test
-    public void testMultithreadedBenchmark() {
-        System.setProperty("collect-memory", "true");
-        try {
-            String[] args = {
-                "-e",
-                "PyTorch",
-                "-u",
-                "djl://ai.djl.pytorch/resnet/0.0.1/traced_resnet18",
-                "-s",
-                "(1,3,224,224)f",
-                "-d",
-                "1",
-                "-l",
-                "1",
-                "-c",
-                "2",
-                "-t",
-                "-1",
-                "-g",
-                "-1"
-            };
-            Benchmark.main(args);
-        } finally {
-            System.clearProperty("collect-memory");
-        }
-    }
-}
diff --git a/extensions/benchmark/src/test/java/ai/djl/benchmark/NDListGeneratorTest.java b/extensions/benchmark/src/test/java/ai/djl/benchmark/NDListGeneratorTest.java
deleted file mode 100644
index 50964a9a1..000000000
--- a/extensions/benchmark/src/test/java/ai/djl/benchmark/NDListGeneratorTest.java
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * 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.benchmark;
-
-import org.testng.Assert;
-import org.testng.annotations.Test;
-
-public class NDListGeneratorTest {
-
-    @Test
-    public void testHelp() {
-        String[] args = {"ndlist-gen", "-h"};
-        Benchmark.main(args);
-    }
-
-    @Test
-    public void testMissingOptions() {
-        String[] args = {"ndlist-gen", "-s"};
-        boolean success = NDListGenerator.generate(args);
-        Assert.assertFalse(success);
-    }
-
-    @Test
-    public void testOnes() {
-        String[] args = {"ndlist-gen", "-s", "1", "-o", "build/ones.ndlist", "-1"};
-        boolean success = NDListGenerator.generate(args);
-        Assert.assertTrue(success);
-    }
-
-    @Test
-    public void testZeros() {
-        String[] args = {"ndlist-gen", "-s", "1", "-o", "build/ones.ndlist"};
-        boolean success = NDListGenerator.generate(args);
-        Assert.assertTrue(success);
-    }
-}
diff --git a/extensions/benchmark/src/test/java/ai/djl/benchmark/package-info.java b/extensions/benchmark/src/test/java/ai/djl/benchmark/package-info.java
deleted file mode 100644
index fd842219c..000000000
--- a/extensions/benchmark/src/test/java/ai/djl/benchmark/package-info.java
+++ /dev/null
@@ -1,15 +0,0 @@
-/*
- * 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.
- */
-
-/** Contains tests for the benchmark module. */
-package ai.djl.benchmark;
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
index 49c165fd1..c9c6744bd 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
@@ -157,6 +157,15 @@ public class NDIndexTest {
             expected = manager.create(new int[] {0, 1, 9, 10, 4, 5, 11, 12}, new Shape(2, 4));
             original.set(new NDIndex(":, 2:"), manager.arange(9, 13).reshape(2, 2));
             Assert.assertEquals(original, expected);
+
+            // set by index array
+            original = manager.arange(1, 10).reshape(3, 3);
+            NDArray index = manager.create(new long[] {0, 1}, new Shape(2));
+            value = manager.create(new int[] {666, 777, 888, 999}, new Shape(2, 2));
+            original.set(new NDIndex("{}, :{}", index, 2), value);
+            expected =
+                    manager.create(new int[] {666, 777, 3, 888, 999, 6, 7, 8, 9}, new Shape(3, 3));
+            Assert.assertEquals(original, expected);
         }
     }
 
@@ -188,6 +197,14 @@ public class NDIndexTest {
             expected = manager.create(new float[] {1, 1, 1, 3}).reshape(2, 2);
             original.set(new NDIndex("..., 0"), 1);
             Assert.assertEquals(original, expected);
+
+            // set by index array
+            original = manager.arange(1, 10).reshape(3, 3);
+            NDArray index = manager.create(new long[] {0, 1}, new Shape(2));
+            original.set(new NDIndex("{}, :{}", index, 2), 666);
+            expected =
+                    manager.create(new int[] {666, 666, 3, 666, 666, 6, 7, 8, 9}, new Shape(3, 3));
+            Assert.assertEquals(original, expected);
         }
     }