Repository: deepjavalibrary/djl
Pull request #2216 (2022-12-08 03:19:51)

[api][example] Implement gatherNd, partial flatten and enable BertOnCode training
## Description ##

* Implement gatherNd and write the document to explain its defination. Pytorch and MXNet. (Implemetation on PyTorch contains the explaination of the algorithm)
* Implement partial flatten on PyTorch
* Enable BertOnCode. `TrainBertOnCode.java` can now run.
* Delete unused file: MissingOps

## Irrelavant document updates to ignore ##
 Shape.java, Block.java, Adam.java

## Other issue ##
By the way, the `gradle :integration:test` on `Continuous / build (ubuntu-latest)`  seems more flaky than before. The error in the following runs are flaky, i.e. they ran into failures on integration:test, then just by re-running them, they finally pass.

https://github.com/deepjavalibrary/djl/actions/runs/3645090324/jobs/6154918749
https://github.com/deepjavalibrary/djl/actions/runs/3645118004/jobs/6154969134
https://github.com/deepjavalibrary/djl/actions/runs/3645120287/jobs/6154972716

https://github.com/deepjavalibrary/djl/actions/runs/3641764881/jobs/6148092580

The failed tests are Adam and AdamW. It's possible that the JNI is not updated when running the PR tests.


Patch:
diff --git a/android/README.md b/android/README.md
index ef167cb3f..f158c5825 100644
--- a/android/README.md
+++ b/android/README.md
@@ -16,9 +16,12 @@ In gradle, you can add the 5 modules in your dependencies:
 
 ```
 dependencies {
-    implementation "ai.djl:api:0.20.0"
+    implementation platform("ai.djl:bom:0.20.0")
+
+    implementation "ai.djl:api"
+    // TODO: remove djl android package version in 0.21.0
     implementation "ai.djl.android:core:0.20.0"
-    runtimeOnly "ai.djl.pytorch:pytorch-engine:0.20.0"
+    runtimeOnly "ai.djl.pytorch:pytorch-engine"
     runtimeOnly "ai.djl.android:pytorch-native:0.20.0"
     runtimeOnly "ai.djl.android:onnxruntime:0.20.0"
 }
diff --git a/android/core/build.gradle b/android/core/build.gradle
index 0a31876cf..8f964f6b3 100644
--- a/android/core/build.gradle
+++ b/android/core/build.gradle
@@ -22,8 +22,8 @@ android {
         }
     }
     compileOptions {
-        sourceCompatibility = 1.8
-        targetCompatibility = 1.8
+        sourceCompatibility = 11
+        targetCompatibility = 11
     }
 }
 
@@ -98,11 +98,15 @@ afterEvaluate {
 }
 
 dependencies {
-    implementation "ai.djl:api:${version}"
+    implementation platform("ai.djl:bom:${version}")
+
+    implementation "ai.djl:api"
     testImplementation 'junit:junit:4.13.2'
-    androidTestRuntimeOnly "ai.djl.pytorch:pytorch-engine:${version}"
+    androidTestRuntimeOnly "ai.djl.pytorch:pytorch-engine"
+
+    // TODO: remove version in 0.21.0
     androidTestRuntimeOnly "ai.djl.android:pytorch-native:${version}"
     androidTestImplementation "ai.djl.android:onnxruntime:${version}"
-    androidTestImplementation 'androidx.test.ext:junit:1.1.3'
-    androidTestImplementation 'androidx.test.espresso:espresso-core:3.4.0'
+    androidTestImplementation 'androidx.test.ext:junit:1.1.4'
+    androidTestImplementation 'androidx.test.espresso:espresso-core:3.5.0'
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslatorFactory.java
new file mode 100644
index 000000000..d48176aa4
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/translator/BaseImageTranslatorFactory.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright 2022 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.modality.cv.translator;
+
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.translator.wrapper.FileImagePreProcesor;
+import ai.djl.modality.cv.translator.wrapper.InputStreamImagePreProcessor;
+import ai.djl.modality.cv.translator.wrapper.UrlImagePreProcessor;
+import ai.djl.translate.ExpansionTranslatorFactory;
+import ai.djl.translate.PreProcessor;
+import ai.djl.util.Pair;
+
+import java.io.InputStream;
+import java.lang.reflect.Type;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+/**
+ * A helper to create a {@link ai.djl.translate.TranslatorFactory} with the {@link
+ * BaseImageTranslator}.
+ *
+ * @param <O> the output type for the {@link ai.djl.translate.TranslatorFactory}.
+ */
+public abstract class BaseImageTranslatorFactory<O> extends ExpansionTranslatorFactory<Image, O> {
+
+    /** {@inheritDoc} */
+    @Override
+    protected Map<Type, Function<PreProcessor<Image>, PreProcessor<?>>>
+            getPreprocessorExpansions() {
+        Map<Type, Function<PreProcessor<Image>, PreProcessor<?>>> expansions =
+                new ConcurrentHashMap<>();
+        expansions.put(Path.class, FileImagePreProcesor::new);
+        expansions.put(URL.class, UrlImagePreProcessor::new);
+        expansions.put(InputStream.class, InputStreamImagePreProcessor::new);
+        return expansions;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected Map<Pair<Type, Type>, TranslatorExpansion<Image, O>> getExpansions() {
+        Map<Pair<Type, Type>, TranslatorExpansion<Image, O>> expansions = new ConcurrentHashMap<>();
+        expansions.put(new Pair<>(Input.class, Output.class), ImageServingTranslator::new);
+        return expansions;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Class<Image> getBaseInputType() {
+        return Image.class;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
index 9ee08ceb5..24a63135b 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactory.java
@@ -14,42 +14,15 @@ package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
 import ai.djl.modality.Classifications;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
-import ai.djl.translate.ExpansionTranslatorFactory;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
-import java.io.InputStream;
-import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.function.Function;
 
 /** A {@link TranslatorFactory} that creates an {@link ImageClassificationTranslator}. */
 public class ImageClassificationTranslatorFactory
-        extends ExpansionTranslatorFactory<Image, Classifications> {
-
-    private static final Map<
-                    Pair<Type, Type>,
-                    Function<Translator<Image, Classifications>, Translator<?, ?>>>
-            EXPANSIONS = new ConcurrentHashMap<>();
-
-    static {
-        EXPANSIONS.put(new Pair<>(Image.class, Classifications.class), t -> t);
-        EXPANSIONS.put(new Pair<>(Path.class, Classifications.class), FileTranslator::new);
-        EXPANSIONS.put(new Pair<>(URL.class, Classifications.class), UrlTranslator::new);
-        EXPANSIONS.put(
-                new Pair<>(InputStream.class, Classifications.class), InputStreamTranslator::new);
-        EXPANSIONS.put(new Pair<>(Input.class, Output.class), ImageServingTranslator::new);
-    }
+        extends BaseImageTranslatorFactory<Classifications> {
 
     /** {@inheritDoc} */
     @Override
@@ -60,8 +33,7 @@ public class ImageClassificationTranslatorFactory
 
     /** {@inheritDoc} */
     @Override
-    protected Map<Pair<Type, Type>, Function<Translator<Image, Classifications>, Translator<?, ?>>>
-            getExpansions() {
-        return EXPANSIONS;
+    public Class<Classifications> getBaseOutputType() {
+        return Classifications.class;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/ImageFeatureExtractorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/ImageFeatureExtractorFactory.java
index b89c7a155..12e41b648 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/ImageFeatureExtractorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/ImageFeatureExtractorFactory.java
@@ -13,60 +13,24 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
-import java.io.InputStream;
-import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 
 /** A {@link TranslatorFactory} that creates an {@link ImageClassificationTranslator}. */
-public class ImageFeatureExtractorFactory implements TranslatorFactory {
-
-    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
-
-    static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, byte[].class));
-        SUPPORTED_TYPES.add(new Pair<>(Path.class, byte[].class));
-        SUPPORTED_TYPES.add(new Pair<>(URL.class, byte[].class));
-        SUPPORTED_TYPES.add(new Pair<>(InputStream.class, byte[].class));
-        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
-    }
+public class ImageFeatureExtractorFactory extends BaseImageTranslatorFactory<byte[]> {
 
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return SUPPORTED_TYPES;
+    protected Translator<Image, byte[]> buildBaseTranslator(Model model, Map<String, ?> arguments) {
+        return ImageFeatureExtractor.builder(arguments).build();
     }
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        ImageFeatureExtractor translator = ImageFeatureExtractor.builder(arguments).build();
-        if (input == Image.class && output == byte[].class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == byte[].class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == byte[].class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == byte[].class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    public Class<byte[]> getBaseOutputType() {
+        return byte[].class;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
index 7394356cb..d2216779c 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactory.java
@@ -13,19 +13,11 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** A {@link TranslatorFactory} that creates a {@link InstanceSegmentationTranslator} instance. */
@@ -33,22 +25,8 @@ public class InstanceSegmentationTranslatorFactory extends ObjectDetectionTransl
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        InstanceSegmentationTranslator translator =
-                InstanceSegmentationTranslator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return InstanceSegmentationTranslator.builder(arguments).build();
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/ObjectDetectionTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/ObjectDetectionTranslatorFactory.java
index 3cd2dac1a..ee89e4002 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/ObjectDetectionTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/ObjectDetectionTranslatorFactory.java
@@ -12,38 +12,18 @@
  */
 package ai.djl.modality.cv.translator;
 
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
-import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
-
-import java.io.InputStream;
-import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
-import java.util.HashSet;
-import java.util.Set;
 
 /**
  * An abstract {@link TranslatorFactory} that creates a {@link ObjectDetectionTranslator} instance.
  */
-public abstract class ObjectDetectionTranslatorFactory implements TranslatorFactory {
-
-    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
-
-    static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, DetectedObjects.class));
-        SUPPORTED_TYPES.add(new Pair<>(Path.class, DetectedObjects.class));
-        SUPPORTED_TYPES.add(new Pair<>(URL.class, DetectedObjects.class));
-        SUPPORTED_TYPES.add(new Pair<>(InputStream.class, DetectedObjects.class));
-        SUPPORTED_TYPES.add(new Pair<>(Input.class, Output.class));
-    }
+public abstract class ObjectDetectionTranslatorFactory
+        extends BaseImageTranslatorFactory<DetectedObjects> {
 
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return SUPPORTED_TYPES;
+    public Class<DetectedObjects> getBaseOutputType() {
+        return DetectedObjects.class;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
index d4f794ad7..e2c3f11a0 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactory.java
@@ -13,19 +13,11 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link SingleShotDetectionTranslator} instance. */
@@ -33,22 +25,8 @@ public class SingleShotDetectionTranslatorFactory extends ObjectDetectionTransla
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        SingleShotDetectionTranslator translator =
-                SingleShotDetectionTranslator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return SingleShotDetectionTranslator.builder(arguments).build();
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
index b474e15fc..bd799ec15 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloTranslatorFactory.java
@@ -13,19 +13,11 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link YoloTranslator} instance. */
@@ -33,21 +25,8 @@ public class YoloTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        YoloTranslator translator = YoloTranslator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return YoloTranslator.builder(arguments).build();
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
index c18051aa0..c5b0a971f 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactory.java
@@ -13,19 +13,11 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link YoloV5Translator} instance. */
@@ -33,21 +25,8 @@ public class YoloV5TranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        YoloV5Translator translator = YoloV5Translator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return YoloV5Translator.builder(arguments).build();
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileImagePreProcesor.java b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileImagePreProcesor.java
new file mode 100644
index 000000000..ad91e8f52
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileImagePreProcesor.java
@@ -0,0 +1,43 @@
+/*
+ * 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.modality.cv.translator.wrapper;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.PreProcessor;
+import ai.djl.translate.TranslatorContext;
+
+import java.nio.file.Path;
+
+/** Built-in {@code PreProcessor} that provides image pre-processing from file path. */
+public class FileImagePreProcesor implements PreProcessor<Path> {
+
+    private PreProcessor<Image> preProcessor;
+
+    /**
+     * Creates a {@code FileImagePreProcesor} instance.
+     *
+     * @param preProcessor a {@code {@link PreProcessor}} that can process image
+     */
+    public FileImagePreProcesor(PreProcessor<Image> preProcessor) {
+        this.preProcessor = preProcessor;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Path input) throws Exception {
+        Image image = ImageFactory.getInstance().fromFile(input);
+        return preProcessor.processInput(ctx, image);
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileTranslator.java
deleted file mode 100644
index b3c9b9a42..000000000
--- a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/FileTranslator.java
+++ /dev/null
@@ -1,66 +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.modality.cv.translator.wrapper;
-
-import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.ImageFactory;
-import ai.djl.ndarray.NDList;
-import ai.djl.translate.Batchifier;
-import ai.djl.translate.Translator;
-import ai.djl.translate.TranslatorContext;
-
-import java.nio.file.Path;
-
-/**
- * Built-in {@code Translator} that provides image pre-processing from file path.
- *
- * @param <T> the output object type
- */
-public class FileTranslator<T> implements Translator<Path, T> {
-
-    private Translator<Image, T> translator;
-
-    /**
-     * Creates a {@code FileTranslator} instance.
-     *
-     * @param translator a {@code Translator} that can process image
-     */
-    public FileTranslator(Translator<Image, T> translator) {
-        this.translator = translator;
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public NDList processInput(TranslatorContext ctx, Path input) throws Exception {
-        Image image = ImageFactory.getInstance().fromFile(input);
-        return translator.processInput(ctx, image);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public T processOutput(TranslatorContext ctx, NDList list) throws Exception {
-        return translator.processOutput(ctx, list);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Batchifier getBatchifier() {
-        return translator.getBatchifier();
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void prepare(TranslatorContext ctx) throws Exception {
-        translator.prepare(ctx);
-    }
-}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamImagePreProcessor.java b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamImagePreProcessor.java
new file mode 100644
index 000000000..786623e3b
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamImagePreProcessor.java
@@ -0,0 +1,43 @@
+/*
+ * 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.modality.cv.translator.wrapper;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.PreProcessor;
+import ai.djl.translate.TranslatorContext;
+
+import java.io.InputStream;
+
+/** Built-in {@code PreProcessor} that provides image pre-processing from {@code InputStream}. */
+public class InputStreamImagePreProcessor<T> implements PreProcessor<InputStream> {
+
+    private PreProcessor<Image> preProcessor;
+
+    /**
+     * Creates a {@code InputStreamImagePreProcessor} instance.
+     *
+     * @param preProcessor a {@code PreProcessor} that can process image
+     */
+    public InputStreamImagePreProcessor(PreProcessor<Image> preProcessor) {
+        this.preProcessor = preProcessor;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, InputStream input) throws Exception {
+        Image image = ImageFactory.getInstance().fromInputStream(input);
+        return preProcessor.processInput(ctx, image);
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamTranslator.java
deleted file mode 100644
index 1ed0698c0..000000000
--- a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/InputStreamTranslator.java
+++ /dev/null
@@ -1,66 +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.modality.cv.translator.wrapper;
-
-import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.ImageFactory;
-import ai.djl.ndarray.NDList;
-import ai.djl.translate.Batchifier;
-import ai.djl.translate.Translator;
-import ai.djl.translate.TranslatorContext;
-
-import java.io.InputStream;
-
-/**
- * Built-in {@code Translator} that provides image pre-processing from {@code InputStream}.
- *
- * @param <T> the output object type
- */
-public class InputStreamTranslator<T> implements Translator<InputStream, T> {
-
-    private Translator<Image, T> translator;
-
-    /**
-     * Creates a {@code UrlTranslator} instance.
-     *
-     * @param translator a {@code Translator} that can process image
-     */
-    public InputStreamTranslator(Translator<Image, T> translator) {
-        this.translator = translator;
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public NDList processInput(TranslatorContext ctx, InputStream input) throws Exception {
-        Image image = ImageFactory.getInstance().fromInputStream(input);
-        return translator.processInput(ctx, image);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public T processOutput(TranslatorContext ctx, NDList list) throws Exception {
-        return translator.processOutput(ctx, list);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Batchifier getBatchifier() {
-        return translator.getBatchifier();
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void prepare(TranslatorContext ctx) throws Exception {
-        translator.prepare(ctx);
-    }
-}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlImagePreProcessor.java
similarity index 52%
rename from api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlTranslator.java
rename to api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlImagePreProcessor.java
index bf6b06a76..c12a32da8 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/wrapper/UrlImagePreProcessor.java
@@ -15,52 +15,33 @@ package ai.djl.modality.cv.translator.wrapper;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.ndarray.NDList;
-import ai.djl.translate.Batchifier;
-import ai.djl.translate.Translator;
+import ai.djl.translate.PreProcessor;
 import ai.djl.translate.TranslatorContext;
 
 import java.net.URL;
 
 /**
- * Built-in {@code Translator} that provides image pre-processing from URL.
+ * Built-in {@code PreProcessor} that provides image pre-processing from URL.
  *
  * @param <T> the output object type
  */
-public class UrlTranslator<T> implements Translator<URL, T> {
+public class UrlImagePreProcessor<T> implements PreProcessor<URL> {
 
-    private Translator<Image, T> translator;
+    private PreProcessor<Image> preProcessor;
 
     /**
-     * Creates a {@code UrlTranslator} instance.
+     * Creates a {@code UrlImagePreProcessor} instance.
      *
-     * @param translator a {@code Translator} that can process image
+     * @param preProcessor a {@code PreProcessor} that can process image
      */
-    public UrlTranslator(Translator<Image, T> translator) {
-        this.translator = translator;
+    public UrlImagePreProcessor(PreProcessor<Image> preProcessor) {
+        this.preProcessor = preProcessor;
     }
 
     /** {@inheritDoc} */
     @Override
     public NDList processInput(TranslatorContext ctx, URL input) throws Exception {
         Image image = ImageFactory.getInstance().fromUrl(input);
-        return translator.processInput(ctx, image);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public T processOutput(TranslatorContext ctx, NDList list) throws Exception {
-        return translator.processOutput(ctx, list);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Batchifier getBatchifier() {
-        return translator.getBatchifier();
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void prepare(TranslatorContext ctx) throws Exception {
-        translator.prepare(ctx);
+        return preProcessor.processInput(ctx, image);
     }
 }
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 434e62542..59ec56534 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -572,9 +572,11 @@ public interface NDArray extends NDResource, BytesSupplier {
     }
 
     /**
-     * Returns a partial {@code NDArray} pointed by the indexed array. Given NDArray arr, NDArray
-     * idx, and long axis, the output is out_{ijk} = arr_{idx_{ijk}, j, k} if axis=0 or arr_{i,
-     * idx_{ijk}, k} if axis=1 or arr_{i, j, idx_{ijk}} if axis=2
+     * Returns a partial {@code NDArray} pointed by the indexed array. \( Given NDArray arr, NDArray
+     * idx, and long axis, \)<br>
+     * \( out_{ijk} = arr_{idx_{ijk}, j, k} if axis=0 \)<br>
+     * \( arr_{i, idx_{ijk}, k} if axis=1 \)<br>
+     * \( arr_{i, j, idx_{ijk}} if axis=2 \)<br>
      *
      * @param index picks the elements of an NDArray to the same position as index
      * @param axis the entries of index are indices of axis
@@ -582,6 +584,24 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray gather(NDArray index, int axis);
 
+    /**
+     * Returns a partial {@code NDArray} pointed by the indexed array. Given NDArray arr and NDArray
+     * idx. idx is the following structure: <br>
+     * \( idx = [ idx[0, ...], idx[1, ...],..., idx[indexingDepth,...] ] \)<br>
+     * corresponding to x, y, z index, i.e. [idx_x, idx_y, idx_z, ...]. <br>
+     * So indexingDepth smaller than or equal to data.shape[0] <br>
+     * If indexingDepth is smaller than data.shape[0], for instance, data.shape[0] = 3, i.e. x, y,
+     * z, <br>
+     * but indexingDepth = 2, i.e. [idx_x, idx_y], <br>
+     * then the tail co-rank = data.shape[0] - indexingDepth will be kept. <br>
+     * With it, the output shape = idx_y.shape appended by data.shape[indexingDepth:]
+     * https://mxnet.apache.org/versions/1.6/api/r/docs/api/mx.symbol.gather_nd.html?highlight=gather_nd
+     *
+     * @param index picks the elements of an NDArray to the same position as index
+     * @return the partial {@code NDArray} of the same shape as index
+     */
+    NDArray gatherNd(NDArray index);
+
     /**
      * Returns a partial {@code NDArray} pointed by index according to linear indexing, and the of
      * output is of the same shape as index.
@@ -3157,6 +3177,24 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray flatten();
 
+    /**
+     * Flattens this {@code NDArray} into a partially flatten {@code NDArray}.
+     *
+     * <p>Examples
+     *
+     * <pre>
+     * jshell&gt; NDArray array = manager.create(new float[]{1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f}, new Shape(2, 2, 2));
+     * jshell&gt; array.flatten(0, 1);
+     * ND: (4) cpu() float32
+     * [[1., 2], [3., 4.], [5., 6.], [7., 8.]]
+     * </pre>
+     *
+     * @param startDim the first dim to flatten, inclusive
+     * @param endDim the last dim to flatten, inclusive
+     * @return a partially fallen {@code NDArray}
+     */
+    NDArray flatten(int startDim, int endDim);
+
     /**
      * Reshapes this {@code NDArray} to the given {@link Shape}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 1e9a4a486..fbe2d7b05 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -191,6 +191,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException(UNSUPPORTED_MSG);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gatherNd(NDArray index) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray take(NDManager manager, NDArray index) {
@@ -822,6 +828,12 @@ public abstract class NDArrayAdapter implements NDArray {
         return getAlternativeArray().flatten();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray flatten(int startDim, int endDim) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/api/src/main/java/ai/djl/ndarray/types/Shape.java b/api/src/main/java/ai/djl/ndarray/types/Shape.java
index 08eaf0960..4c3262974 100644
--- a/api/src/main/java/ai/djl/ndarray/types/Shape.java
+++ b/api/src/main/java/ai/djl/ndarray/types/Shape.java
@@ -227,6 +227,9 @@ public class Shape {
      * @return a new {@code Shape} whose content is a slice of this shape
      */
     public Shape slice(int beginIndex, int endIndex) {
+        beginIndex = beginIndex + (beginIndex < 0 ? shape.length : 0);
+        endIndex = endIndex + (endIndex < 0 ? shape.length : 0);
+
         int size = endIndex - beginIndex;
         long[] out = new long[size];
         System.arraycopy(shape, beginIndex, out, 0, size);
diff --git a/api/src/main/java/ai/djl/nn/Block.java b/api/src/main/java/ai/djl/nn/Block.java
index 6ca7bc257..0adf2e457 100644
--- a/api/src/main/java/ai/djl/nn/Block.java
+++ b/api/src/main/java/ai/djl/nn/Block.java
@@ -185,7 +185,8 @@ public interface Block {
     void setInitializer(Initializer initializer, Predicate<Parameter> predicate);
 
     /**
-     * Initializes the parameters of the block. This method must be called before calling `forward`.
+     * Initializes the parameters of the block, set require gradient if required and infer the block
+     * inputShape. This method must be called before calling `forward`.
      *
      * @param manager the NDManager to initialize the parameters
      * @param dataType the datatype of the parameters
diff --git a/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java b/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java
index d8002d6e5..a0b49b943 100644
--- a/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java
+++ b/api/src/main/java/ai/djl/nn/transformer/BertMaskedLanguageModelBlock.java
@@ -96,7 +96,7 @@ public class BertMaskedLanguageModelBlock extends AbstractBlock {
         // Now we create one long sequence by appending all sequences
         NDArray flattenedSequences = sequences.reshape((long) batchSize * sequenceLength, width);
         // We use the absolute indices to gather the elements of the flattened sequences
-        return MissingOps.gatherNd(flattenedSequences, absoluteIndices);
+        return flattenedSequences.gatherNd(absoluteIndices);
     }
 
     /** {@inheritDoc} */
diff --git a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
index cfb1e08ae..d135d5c8b 100644
--- a/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
+++ b/api/src/main/java/ai/djl/nn/transformer/IdEmbedding.java
@@ -68,7 +68,7 @@ public final class IdEmbedding extends AbstractBlock {
         // create the embedding Table
         NDArray embeddingTable = ps.getValue(embedding, ids.getDevice(), training);
         // We do not perform a sparse lookup, instead we just project into the table
-        NDArray result = MissingOps.gatherNd(embeddingTable, ids);
+        NDArray result = embeddingTable.gatherNd(ids);
         // we want the original shape of the input + the last dimension of the embedding
         Shape targetShape = input.getShape().addAll(new Shape(embeddingTable.getShape().get(1)));
         return new NDList(result.reshape(targetShape));
diff --git a/api/src/main/java/ai/djl/nn/transformer/MissingOps.java b/api/src/main/java/ai/djl/nn/transformer/MissingOps.java
deleted file mode 100644
index 0412a6797..000000000
--- a/api/src/main/java/ai/djl/nn/transformer/MissingOps.java
+++ /dev/null
@@ -1,33 +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.nn.transformer;
-
-import ai.djl.ndarray.NDArray;
-import ai.djl.ndarray.NDList;
-
-/** Operators missing from NDArray that are necessary to implement Bert pretraining. */
-public final class MissingOps {
-
-    private MissingOps() {}
-
-    /**
-     * Applies the mxnet gather_nd operator.
-     *
-     * @param lookup array with the data to look up (e.g. an embedding table)
-     * @param indices indices to use for lookup
-     * @return the lookup result
-     */
-    public static NDArray gatherNd(NDArray lookup, NDArray indices) {
-        return indices.getManager().invoke("gather_nd", new NDList(lookup, indices), null).head();
-    }
-}
diff --git a/api/src/main/java/ai/djl/training/optimizer/Adam.java b/api/src/main/java/ai/djl/training/optimizer/Adam.java
index 22ee0b390..1d7ce86a3 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Adam.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Adam.java
@@ -72,9 +72,9 @@ public class Adam extends Optimizer {
         float weightDecay = getWeightDecay();
 
         Preconditions.checkArgument(
-                !Float.isNaN(newLearningRate)
+                !Float.isNaN(learningRateBiasCorrection)
                         && !Float.isNaN(weightDecay)
-                        && !Float.isInfinite(newLearningRate)
+                        && !Float.isInfinite(learningRateBiasCorrection)
                         && !Float.isInfinite(weightDecay),
                 "learning rate or weight decay is nan or infinite");
         NDList inputs =
diff --git a/api/src/main/java/ai/djl/translate/BasicTranslator.java b/api/src/main/java/ai/djl/translate/BasicTranslator.java
new file mode 100644
index 000000000..70ad2a749
--- /dev/null
+++ b/api/src/main/java/ai/djl/translate/BasicTranslator.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright 2022 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.translate;
+
+import ai.djl.ndarray.NDList;
+
+/**
+ * A {@link Translator} made by combining a {@link PreProcessor} and a {@link PostProcessor}.
+ *
+ * @param <I> the input class
+ * @param <O> the output class
+ */
+public class BasicTranslator<I, O> implements Translator<I, O> {
+
+    private PreProcessor<I> preProcessor;
+    private PostProcessor<O> postProcessor;
+    private Batchifier batchifier;
+
+    /**
+     * Constructs a {@link BasicTranslator} with the default {@link Batchifier}.
+     *
+     * @param preProcessor the preProcessor to use for pre-processing
+     * @param postProcessor the postProcessor to use for post-processing
+     */
+    public BasicTranslator(PreProcessor<I> preProcessor, PostProcessor<O> postProcessor) {
+        this.preProcessor = preProcessor;
+        this.postProcessor = postProcessor;
+    }
+
+    /**
+     * Constructs a {@link BasicTranslator}.
+     *
+     * @param preProcessor the preProcessor to use for pre-processing
+     * @param postProcessor the postProcessor to use for post-processing
+     * @param batchifier the batchifier to use
+     */
+    public BasicTranslator(
+            PreProcessor<I> preProcessor, PostProcessor<O> postProcessor, Batchifier batchifier) {
+        this.preProcessor = preProcessor;
+        this.postProcessor = postProcessor;
+        this.batchifier = batchifier;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public O processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        return postProcessor.processOutput(ctx, list);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, I input) throws Exception {
+        return preProcessor.processInput(ctx, input);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Batchifier getBatchifier() {
+        if (batchifier != null) {
+            return batchifier;
+        }
+        return Translator.super.getBatchifier();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws Exception {
+        if (preProcessor instanceof Translator) {
+            ((Translator<?, ?>) preProcessor).prepare(ctx);
+        }
+        if (postProcessor instanceof Translator && postProcessor != preProcessor) {
+            ((Translator<?, ?>) postProcessor).prepare(ctx);
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/translate/ExpansionTranslatorFactory.java b/api/src/main/java/ai/djl/translate/ExpansionTranslatorFactory.java
index 952f2901a..6a58bd55d 100644
--- a/api/src/main/java/ai/djl/translate/ExpansionTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/translate/ExpansionTranslatorFactory.java
@@ -16,6 +16,8 @@ import ai.djl.Model;
 import ai.djl.util.Pair;
 
 import java.lang.reflect.Type;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
@@ -26,13 +28,29 @@ import java.util.function.Function;
  * @param <IbaseT> the input type for the base translator
  * @param <ObaseT> the output type for the base translator
  */
-@SuppressWarnings("PMD.GenericsNaming")
+@SuppressWarnings({"PMD.GenericsNaming", "InterfaceTypeParameterName"})
 public abstract class ExpansionTranslatorFactory<IbaseT, ObaseT> implements TranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
     public Set<Pair<Type, Type>> getSupportedTypes() {
-        return getExpansions().keySet();
+        Set<Pair<Type, Type>> results = new HashSet<>();
+        results.addAll(getExpansions().keySet());
+
+        Set<Type> preProcessorTypes = new HashSet<>();
+        preProcessorTypes.addAll(getPreprocessorExpansions().keySet());
+        preProcessorTypes.add(getBaseInputType());
+
+        Set<Type> postProcessorTypes = new HashSet<>();
+        postProcessorTypes.addAll(getPostprocessorExpansions().keySet());
+        postProcessorTypes.add(getBaseOutputType());
+
+        for (Type i : preProcessorTypes) {
+            for (Type o : postProcessorTypes) {
+                results.add(new Pair<>(i, o));
+            }
+        }
+        return results;
     }
 
     /** {@inheritDoc} */
@@ -43,15 +61,58 @@ public abstract class ExpansionTranslatorFactory<IbaseT, ObaseT> implements Tran
         return newInstance(input, output, baseTranslator);
     }
 
+    /**
+     * Returns a new instance of the {@link Translator} class.
+     *
+     * @param <I> the input data type
+     * @param <O> the output data type
+     * @param input the input class
+     * @param output the output class
+     * @param translator the base translator to expand from
+     * @return a new instance of the {@code Translator} class
+     */
     @SuppressWarnings("unchecked")
     <I, O> Translator<I, O> newInstance(
             Class<I> input, Class<O> output, Translator<IbaseT, ObaseT> translator) {
-        Function<Translator<IbaseT, ObaseT>, Translator<?, ?>> expansion =
+
+        if (input.equals(getBaseInputType()) && output.equals(getBaseOutputType())) {
+            return (Translator<I, O>) translator;
+        }
+
+        TranslatorExpansion<IbaseT, ObaseT> expansion =
                 getExpansions().get(new Pair<>(input, output));
-        if (expansion == null) {
-            throw new IllegalArgumentException("Unsupported expansion input/output types.");
+        if (expansion != null) {
+            return (Translator<I, O>) expansion.apply(translator);
+        }
+
+        // Note that regular expansions take precedence over pre-processor+post-processor expansions
+        PreProcessor<I> preProcessor = null;
+        if (input.equals(getBaseInputType())) {
+            preProcessor = (PreProcessor<I>) translator;
+        } else {
+            Function<PreProcessor<IbaseT>, PreProcessor<?>> expander =
+                    getPreprocessorExpansions().get(input);
+            if (expander != null) {
+                preProcessor = (PreProcessor<I>) expander.apply(translator);
+            }
+        }
+
+        PostProcessor<O> postProcessor = null;
+        if (output.equals(getBaseOutputType())) {
+            postProcessor = (PostProcessor<O>) translator;
+        } else {
+            Function<PostProcessor<ObaseT>, PostProcessor<?>> expander =
+                    getPostprocessorExpansions().get(output);
+            if (expander != null) {
+                postProcessor = (PostProcessor<O>) expander.apply(translator);
+            }
         }
-        return (Translator<I, O>) expansion.apply(translator);
+
+        if (preProcessor != null && postProcessor != null) {
+            return new BasicTranslator<>(preProcessor, postProcessor, translator.getBatchifier());
+        }
+
+        throw new IllegalArgumentException("Unsupported expansion input/output types.");
     }
 
     /**
@@ -74,14 +135,50 @@ public abstract class ExpansionTranslatorFactory<IbaseT, ObaseT> implements Tran
     protected abstract Translator<IbaseT, ObaseT> buildBaseTranslator(
             Model model, Map<String, ?> arguments);
 
+    /**
+     * Returns the input type for the base translator.
+     *
+     * @return the input type for the base translator
+     */
+    public abstract Class<IbaseT> getBaseInputType();
+
+    /**
+     * Returns the output type for the base translator.
+     *
+     * @return the output type for the base translator
+     */
+    public abstract Class<ObaseT> getBaseOutputType();
+
+    /**
+     * Returns the possible expansions of this factory.
+     *
+     * @return the possible expansions of this factory
+     */
+    protected Map<Pair<Type, Type>, TranslatorExpansion<IbaseT, ObaseT>> getExpansions() {
+        return Collections.emptyMap();
+    }
+
+    /**
+     * Returns the possible expansions of this factory.
+     *
+     * @return the possible expansions of this factory
+     */
+    protected Map<Type, Function<PreProcessor<IbaseT>, PreProcessor<?>>>
+            getPreprocessorExpansions() {
+        return Collections.singletonMap(getBaseInputType(), p -> p);
+    }
+
     /**
      * Returns the possible expansions of this factory.
      *
      * @return the possible expansions of this factory
      */
-    protected abstract Map<Pair<Type, Type>, Function<Translator<IbaseT, ObaseT>, Translator<?, ?>>>
-            getExpansions();
+    protected Map<Type, Function<PostProcessor<ObaseT>, PostProcessor<?>>>
+            getPostprocessorExpansions() {
+        return Collections.singletonMap(getBaseOutputType(), p -> p);
+    }
 
+    /** Represents {@link TranslatorOptions} by applying expansions to a base {@link Translator}. */
     final class ExpandedTranslatorOptions implements TranslatorOptions {
 
         private Translator<IbaseT, ObaseT> translator;
@@ -102,4 +199,15 @@ public abstract class ExpansionTranslatorFactory<IbaseT, ObaseT> implements Tran
             return newInstance(input, output, translator);
         }
     }
+
+    /**
+     * A function from a base translator to an expanded translator.
+     *
+     * @param <IbaseT> the base translator input type
+     * @param <ObaseT> the base translator output type
+     */
+    @FunctionalInterface
+    @SuppressWarnings({"PMD.GenericsNaming", "InterfaceTypeParameterName"})
+    public interface TranslatorExpansion<IbaseT, ObaseT>
+            extends Function<Translator<IbaseT, ObaseT>, Translator<?, ?>> {}
 }
diff --git a/bom/build.gradle b/bom/build.gradle
index 173aa612d..470aff1dd 100644
--- a/bom/build.gradle
+++ b/bom/build.gradle
@@ -19,6 +19,9 @@ dependencies {
         api "ai.djl:basicdataset:${version}"
         api "ai.djl:model-zoo:${version}"
         api "ai.djl:serving:${version}"
+        api "ai.djl.android:core:${version}"
+        api "ai.djl.android:onnxruntime:${version}"
+        api "ai.djl.android:pytorch-native:${version}"
         api "ai.djl.aws:aws-ai:${version}"
         api "ai.djl.audio:audio:${version}"
         api "ai.djl.dlr:dlr-engine:${version}"
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index fef0a915e..8f9dcbc0b 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -322,6 +322,13 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gatherNd(NDArray index) {
+        MxOpParams params = new MxOpParams();
+        return manager.invoke("gather_nd", new NDList(this, index), params).singletonOrThrow();
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray take(NDManager manager, NDArray index) {
@@ -1091,6 +1098,12 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         return reshape(new Shape(Math.toIntExact(size())));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray flatten(int startDim, int endDim) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
index ca4b83a3b..4a36afdce 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/imageclassification/PpWordRotateTranslatorFactory.java
@@ -14,20 +14,11 @@ package ai.djl.paddlepaddle.zoo.cv.imageclassification;
 
 import ai.djl.Model;
 import ai.djl.modality.Classifications;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.translator.ImageClassificationTranslatorFactory;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link PpWordRotateTranslatorFactory} instance. */
@@ -35,20 +26,8 @@ public class PpWordRotateTranslatorFactory extends ImageClassificationTranslator
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        if (input == Image.class && output == Classifications.class) {
-            return (Translator<I, O>) new PpWordRotateTranslator();
-        } else if (input == Path.class && output == Classifications.class) {
-            return (Translator<I, O>) new FileTranslator<>(new PpWordRotateTranslator());
-        } else if (input == URL.class && output == Classifications.class) {
-            return (Translator<I, O>) new UrlTranslator<>(new PpWordRotateTranslator());
-        } else if (input == InputStream.class && output == Classifications.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(new PpWordRotateTranslator());
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(new PpWordRotateTranslator());
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, Classifications> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return new PpWordRotateTranslator();
     }
 }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
index c55987644..99203e9c1 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpFaceDetectionTranslatorFactory.java
@@ -13,21 +13,12 @@
 package ai.djl.paddlepaddle.zoo.cv.objectdetection;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
 import ai.djl.modality.cv.translator.ObjectDetectionTranslatorFactory;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link PpFaceDetectionTranslator} instance. */
@@ -35,21 +26,8 @@ public class PpFaceDetectionTranslatorFactory extends ObjectDetectionTranslatorF
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        PpFaceDetectionTranslator translator = new PpFaceDetectionTranslator(arguments);
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return new PpFaceDetectionTranslator(arguments);
     }
 }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
index 2abbbca0e..b16b9f36b 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/objectdetection/PpWordDetectionTranslatorFactory.java
@@ -13,21 +13,12 @@
 package ai.djl.paddlepaddle.zoo.cv.objectdetection;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
 import ai.djl.modality.cv.translator.ObjectDetectionTranslatorFactory;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link PpWordDetectionTranslator} instance. */
@@ -35,21 +26,8 @@ public class PpWordDetectionTranslatorFactory extends ObjectDetectionTranslatorF
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        PpWordDetectionTranslator translator = new PpWordDetectionTranslator(arguments);
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return new PpWordDetectionTranslator(arguments);
     }
 }
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
index 5223e1385..f8f341549 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/src/main/java/ai/djl/paddlepaddle/zoo/cv/wordrecognition/PpWordRecognitionTranslatorFactory.java
@@ -13,60 +13,25 @@
 package ai.djl.paddlepaddle.zoo.cv.wordrecognition;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.modality.cv.translator.BaseImageTranslatorFactory;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
-import ai.djl.util.Pair;
 
-import java.io.InputStream;
-import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 
 /** An {@link TranslatorFactory} that creates a {@link PpWordRecognitionTranslator} instance. */
-public class PpWordRecognitionTranslatorFactory implements TranslatorFactory {
-
-    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
-
-    static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, String.class));
-        SUPPORTED_TYPES.add(new Pair<>(Path.class, String.class));
-        SUPPORTED_TYPES.add(new Pair<>(URL.class, String.class));
-        SUPPORTED_TYPES.add(new Pair<>(InputStream.class, String.class));
-    }
+public class PpWordRecognitionTranslatorFactory extends BaseImageTranslatorFactory<String> {
 
     /** {@inheritDoc} */
     @Override
-    public Set<Pair<Type, Type>> getSupportedTypes() {
-        return SUPPORTED_TYPES;
+    protected Translator<Image, String> buildBaseTranslator(Model model, Map<String, ?> arguments) {
+        return new PpWordRecognitionTranslator();
     }
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        PpWordRecognitionTranslator translator = new PpWordRecognitionTranslator();
-        if (input == Image.class && output == String.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == String.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == String.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == String.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    public Class<String> getBaseOutputType() {
+        return String.class;
     }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index cf2f35f96..d51cc0acd 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -261,11 +261,45 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     @Override
     public NDArray gather(NDArray index, int axis) {
         if (!(index instanceof PtNDArray)) {
-            throw new IllegalArgumentException("Only PtNDArray is supported.");
+            throw new IllegalArgumentException("Only PtNDArray index is supported.");
         }
         return JniUtils.gather(this, (PtNDArray) index, axis);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gatherNd(NDArray index) {
+        if (!(index instanceof PtNDArray)) {
+            throw new IllegalArgumentException("Only PtNDArray index is supported.");
+        }
+        Shape indexShape = index.getShape();
+        Shape dataShape = getShape();
+        int indexingDepth = (int) indexShape.get(0);
+        if (indexingDepth > dataShape.dimension()) {
+            throw new IllegalArgumentException(
+                    "Indexing rank "
+                            + indexShape.get(0)
+                            + " exceeds the data rank "
+                            + dataShape.dimension());
+        }
+        // Row-first order, the linear index is accumulated from z->y->x.
+        // For example, dataShape = (3, 2, 3), indexShape = (2, 3, 3)
+        // The method is: indexLinear = index[1] + index[0] * dataShape[1], row-first order
+        // indexLinear has shape (3, 3), is from combining the index along 0 axis.
+        // Each number in indexLinear is an indexing to an element in data (3, 2, ...).
+        // data is flattened to be (3*2, ...) which can be indexed by indexLinear.
+        // Finally, reshape the output to (3, 3, ...). Thus
+        // totalShape = indexShape.slice(1).addAll(dataShape.slice(indexingDepth));
+        NDArray indexLinear = index.get("{}, ...", indexingDepth - 1);
+        long dim = 1;
+        for (int i = indexingDepth - 2; i > -1; i--) {
+            dim = dim * dataShape.get(i + 1);
+            indexLinear = indexLinear.addi(index.get("{}, ...", i).muli(dim));
+        }
+        NDArray dataFlatten = this.flatten(0, indexingDepth - 1);
+        return dataFlatten.get(indexLinear);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray take(NDManager manager, NDArray index) {
@@ -1016,6 +1050,12 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.flatten(this, 0, -1);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray flatten(int startDim, int endDim) {
+        return JniUtils.flatten(this, startDim, endDim);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArray reshape(Shape shape) {
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
index c0ac429fa..b552ea995 100644
--- a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/cv/objectdetection/PtSsdTranslatorFactory.java
@@ -13,21 +13,12 @@
 package ai.djl.pytorch.zoo.cv.objectdetection;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
 import ai.djl.modality.cv.translator.ObjectDetectionTranslatorFactory;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link PtSsdTranslator} instance. */
@@ -35,21 +26,8 @@ public class PtSsdTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        PtSsdTranslator translator = PtSsdTranslator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return PtSsdTranslator.builder(arguments).build();
     }
 }
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index bd785b3fe..0cf130c1e 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -209,6 +209,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gatherNd(NDArray index) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray take(NDManager manager, NDArray index) {
@@ -1128,6 +1134,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         return reshape(new Shape(-1));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray flatten(int startDim, int endDim) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
index 2b0b7e52e..a97a0b25d 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/main/java/ai/djl/tensorflow/zoo/cv/objectdetction/TfSsdTranslatorFactory.java
@@ -13,21 +13,12 @@
 package ai.djl.tensorflow.zoo.cv.objectdetction;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.ImageServingTranslator;
 import ai.djl.modality.cv.translator.ObjectDetectionTranslatorFactory;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 
-import java.io.InputStream;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.Map;
 
 /** An {@link TranslatorFactory} that creates a {@link TfSsdTranslator} instance. */
@@ -35,21 +26,8 @@ public class TfSsdTranslatorFactory extends ObjectDetectionTranslatorFactory {
 
     /** {@inheritDoc} */
     @Override
-    @SuppressWarnings("unchecked")
-    public <I, O> Translator<I, O> newInstance(
-            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        TfSsdTranslator translator = TfSsdTranslator.builder(arguments).build();
-        if (input == Image.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == DetectedObjects.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
-        }
-        throw new IllegalArgumentException("Unsupported input/output types.");
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return TfSsdTranslator.builder(arguments).build();
     }
 }


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java
index 941e2c3dc..5427cebe7 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java
@@ -17,9 +17,7 @@ import ai.djl.modality.Classifications;
 import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.BasicTranslator;
 import ai.djl.translate.Translator;
 
 import org.testng.Assert;
@@ -56,15 +54,15 @@ public class ImageClassificationTranslatorFactoryTest {
 
             Translator<Path, Classifications> translator2 =
                     factory.newInstance(Path.class, Classifications.class, model, arguments);
-            Assert.assertTrue(translator2 instanceof FileTranslator);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
 
             Translator<URL, Classifications> translator3 =
                     factory.newInstance(URL.class, Classifications.class, model, arguments);
-            Assert.assertTrue(translator3 instanceof UrlTranslator);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
 
             Translator<InputStream, Classifications> translator4 =
                     factory.newInstance(InputStream.class, Classifications.class, model, arguments);
-            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
 
             Translator<Input, Output> translator5 =
                     factory.newInstance(Input.class, Output.class, model, arguments);
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java
index 79df6bcc9..21eaedadc 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java
@@ -17,9 +17,7 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.BasicTranslator;
 import ai.djl.translate.Translator;
 
 import org.testng.Assert;
@@ -56,15 +54,15 @@ public class InstanceSegmentationTranslatorFactoryTest {
 
             Translator<Path, DetectedObjects> translator2 =
                     factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator2 instanceof FileTranslator);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
 
             Translator<URL, DetectedObjects> translator3 =
                     factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator3 instanceof UrlTranslator);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
 
             Translator<InputStream, DetectedObjects> translator4 =
                     factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
 
             Translator<Input, Output> translator5 =
                     factory.newInstance(Input.class, Output.class, model, arguments);
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java
index 938a5cd0f..f26b7f2ab 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java
@@ -17,9 +17,7 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.BasicTranslator;
 import ai.djl.translate.Translator;
 
 import org.testng.Assert;
@@ -51,15 +49,15 @@ public class SingleShotDetectionTranslatorFactoryTest {
 
             Translator<Path, DetectedObjects> translator2 =
                     factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator2 instanceof FileTranslator);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
 
             Translator<URL, DetectedObjects> translator3 =
                     factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator3 instanceof UrlTranslator);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
 
             Translator<InputStream, DetectedObjects> translator4 =
                     factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
 
             Translator<Input, Output> translator5 =
                     factory.newInstance(Input.class, Output.class, model, arguments);
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java
index fae1a595f..cfbcfb829 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java
@@ -17,9 +17,7 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.BasicTranslator;
 import ai.djl.translate.Translator;
 
 import org.testng.Assert;
@@ -56,15 +54,15 @@ public class YoloTranslatorFactoryTest {
 
             Translator<Path, DetectedObjects> translator2 =
                     factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator2 instanceof FileTranslator);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
 
             Translator<URL, DetectedObjects> translator3 =
                     factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator3 instanceof UrlTranslator);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
 
             Translator<InputStream, DetectedObjects> translator4 =
                     factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
 
             Translator<Input, Output> translator5 =
                     factory.newInstance(Input.class, Output.class, model, arguments);
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java
index a86c03fac..1ac8bf7b8 100644
--- a/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java
@@ -17,9 +17,7 @@ import ai.djl.modality.Input;
 import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.DetectedObjects;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.BasicTranslator;
 import ai.djl.translate.Translator;
 
 import org.testng.Assert;
@@ -56,15 +54,15 @@ public class YoloV5TranslatorFactoryTest {
 
             Translator<Path, DetectedObjects> translator2 =
                     factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator2 instanceof FileTranslator);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
 
             Translator<URL, DetectedObjects> translator3 =
                     factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator3 instanceof UrlTranslator);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
 
             Translator<InputStream, DetectedObjects> translator4 =
                     factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
-            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
 
             Translator<Input, Output> translator5 =
                     factory.newInstance(Input.class, Output.class, model, arguments);
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
index 081afd999..50709beed 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayShapesManipulationOpTest.java
@@ -70,6 +70,14 @@ public class NDArrayShapesManipulationOpTest {
             array = manager.create(new Shape(2, 0));
             expected = manager.create(new Shape(0));
             Assert.assertEquals(array.flatten(), expected);
+
+            // partial flatten
+            array =
+                    manager.create(
+                            new float[] {1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f}, new Shape(2, 2, 2));
+            expected =
+                    manager.create(new float[] {1, 2, 3, 4, 5, 6, 7, 8}).reshape(new Shape(4, 2));
+            Assert.assertEquals(array.flatten(0, 1), expected);
         }
     }
 
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 6a6a3fe1f..9e9b7753a 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
@@ -73,6 +73,26 @@ public class NDIndexTest {
         }
     }
 
+    @Test
+    public void testGatherNd() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray original = manager.arange(1f, 9f).reshape(2, 2, 2);
+            NDArray index = manager.create(new float[] {0, 1, 1, 0}, new Shape(2, 2));
+            NDArray actual = original.gatherNd(index);
+            NDArray expected = manager.create(new float[] {3, 4, 5, 6}, new Shape(2, 2));
+            Assert.assertEquals(actual, expected);
+
+            original = manager.arange(1f, 19f).reshape(3, 2, 3);
+            index = manager.create(new float[] {2, 1, 1, 2, 0, 1, 1, 0}, new Shape(2, 4, 1));
+            actual = original.gatherNd(index);
+            expected =
+                    manager.create(
+                            new float[] {13, 14, 15, 10, 11, 12, 10, 11, 12, 13, 14, 15},
+                            new Shape(4, 1, 3));
+            Assert.assertEquals(actual, expected);
+        }
+    }
+
     @Test
     public void testTake() {
         try (NDManager manager = NDManager.newBaseManager()) {