Repository: deepjavalibrary/djl
Pull request #2045 (2022-09-28 23:02:13)

Handle error when forget to initialize a child block
This throws an error with a much clearer error message in case someone fails to initialize a child block. It should help greatly reduce the time it takes to debug the error in this situation.

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2045?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **69.63**% // Decreases project coverage by **`-2.45%`** :warning:
> Coverage data is based on head [(`efee4f4`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2045?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 68.39% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2045      +/-   ##
============================================
- Coverage     72.08%   69.63%   -2.46%     
- Complexity     5126     5965     +839     
============================================
  Files           473      598     +125     
  Lines         21970    26513    +4543     
  Branches       2351     2880     +529     
============================================
+ Hits          15838    18462    +2624     
- Misses         4925     6665    +1740     
- Partials       1207     1386     +179     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2045?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [.../cv/translator/InstanceSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (-86.59%)` | :arrow_down: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `7.14% <0.00%> (-11.04%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| ... and [512 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2045/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2045?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/modality/Classifications.java b/api/src/main/java/ai/djl/modality/Classifications.java
index ba87736e7..8d21bb5c6 100644
--- a/api/src/main/java/ai/djl/modality/Classifications.java
+++ b/api/src/main/java/ai/djl/modality/Classifications.java
@@ -14,6 +14,7 @@ package ai.djl.modality;
 
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.types.DataType;
+import ai.djl.translate.Ensembleable;
 import ai.djl.util.JsonSerializable;
 import ai.djl.util.JsonUtils;
 
@@ -29,6 +30,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.Iterator;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -36,7 +38,7 @@ import java.util.stream.Collectors;
  * {@code Classifications} is the container that stores the classification results for
  * classification on a single input.
  */
-public class Classifications implements JsonSerializable {
+public class Classifications implements JsonSerializable, Ensembleable<Classifications> {
 
     private static final long serialVersionUID = 1L;
 
@@ -93,6 +95,24 @@ public class Classifications implements JsonSerializable {
         this.topK = topK;
     }
 
+    /**
+     * Returns the classes that were classified into.
+     *
+     * @return the classes that were classified into
+     */
+    public List<String> getClassNames() {
+        return classNames;
+    }
+
+    /**
+     * Returns the list of probabilities for each class (matching the order of the class names).
+     *
+     * @return the list of probabilities for each class (matching the order of the class names)
+     */
+    public List<Double> getProbabilities() {
+        return probabilities;
+    }
+
     /**
      * Set the topK number of classes to be displayed.
      *
@@ -210,6 +230,30 @@ public class Classifications implements JsonSerializable {
         return sb.toString();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Classifications ensembleWith(Iterator<Classifications> it) {
+        int size = probabilities.size();
+        List<Double> newProbabilities = new ArrayList<>(size);
+        newProbabilities.addAll(probabilities);
+        int count = 1;
+        while (it.hasNext()) {
+            ++count;
+            Classifications c = it.next();
+            for (int i = 0; i < size; ++i) {
+                newProbabilities.set(i, newProbabilities.get(i) + c.probabilities.get(i));
+            }
+            if (!c.classNames.equals(classNames)) {
+                throw new IllegalArgumentException(
+                        "Found a classNames mismatch during ensembling. All input Classifications"
+                                + " should have the same classNames, but some were different");
+            }
+        }
+        final int total = count;
+        newProbabilities.replaceAll(p -> p / total);
+        return new Classifications(classNames, newProbabilities);
+    }
+
     /**
      * A {@code Classification} stores the classification result for a single class on a single
      * input.
diff --git a/api/src/main/java/ai/djl/nn/Parameter.java b/api/src/main/java/ai/djl/nn/Parameter.java
index 62b62137f..857e4c474 100644
--- a/api/src/main/java/ai/djl/nn/Parameter.java
+++ b/api/src/main/java/ai/djl/nn/Parameter.java
@@ -120,7 +120,8 @@ public class Parameter implements AutoCloseable {
      */
     public NDArray getArray() {
         if (!isInitialized()) {
-            throw new IllegalStateException("The array has not been initialized");
+            throw new UninitializedParameterException(
+                    "The array for parameter \"" + getName() + "\" has not been initialized");
         }
         return array;
     }
diff --git a/api/src/main/java/ai/djl/nn/UninitializedParameterException.java b/api/src/main/java/ai/djl/nn/UninitializedParameterException.java
new file mode 100644
index 000000000..20d3fe2ac
--- /dev/null
+++ b/api/src/main/java/ai/djl/nn/UninitializedParameterException.java
@@ -0,0 +1,58 @@
+/*
+ * 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.nn;
+
+/** Thrown to indicate that a {@link Parameter} was not initialized. */
+public class UninitializedParameterException extends RuntimeException {
+
+    private static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs a new exception with the specified detail message. The cause is not initialized,
+     * and may subsequently be initialized by a call to {@link #initCause}.
+     *
+     * @param message the detail message that is saved for later retrieval by the {@link
+     *     #getMessage()} method
+     */
+    public UninitializedParameterException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs a new exception with the specified detail message and cause.
+     *
+     * <p>Note that the detail message associated with {@code cause} is <i>not</i> automatically
+     * incorporated in this exception's detail message.
+     *
+     * @param message the detail message that is saved for later retrieval by the {@link
+     *     #getMessage()} method
+     * @param cause the cause that is saved for later retrieval by the {@link #getCause()} method. A
+     *     {@code null} value is permitted, and indicates that the cause is nonexistent or unknown
+     */
+    public UninitializedParameterException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    /**
+     * Constructs a new exception with the specified cause and a detail message of {@code
+     * (cause==null ? null : cause.toString())} which typically contains the class and detail
+     * message of {@code cause}. This constructor is useful for exceptions that are little more than
+     * wrappers for other throwables. For example, {@link java.security.PrivilegedActionException}.
+     *
+     * @param cause the cause that is saved for later retrieval by the {@link #getCause()} method. A
+     *     {@code null} value is permitted, and indicates that the cause is nonexistent or unknown
+     */
+    public UninitializedParameterException(Throwable cause) {
+        super(cause);
+    }
+}
diff --git a/api/src/main/java/ai/djl/nn/core/SparseMax.java b/api/src/main/java/ai/djl/nn/core/SparseMax.java
index f8e0e4f90..feab01650 100644
--- a/api/src/main/java/ai/djl/nn/core/SparseMax.java
+++ b/api/src/main/java/ai/djl/nn/core/SparseMax.java
@@ -10,7 +10,6 @@
  * 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.core;
 
 import ai.djl.ndarray.NDArray;
@@ -66,7 +65,7 @@ public class SparseMax extends AbstractBlock {
     @Override
     public Shape[] getOutputShapes(Shape[] inputShapes) {
         // the shape of input and output are the same
-        return new Shape[0];
+        return new Shape[] {inputShapes[0]};
     }
 
     /** {@inheritDoc} */
diff --git a/api/src/main/java/ai/djl/training/Trainer.java b/api/src/main/java/ai/djl/training/Trainer.java
index 393bca171..8358c0238 100644
--- a/api/src/main/java/ai/djl/training/Trainer.java
+++ b/api/src/main/java/ai/djl/training/Trainer.java
@@ -21,6 +21,7 @@ import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Parameter;
+import ai.djl.nn.UninitializedParameterException;
 import ai.djl.training.dataset.Batch;
 import ai.djl.training.dataset.Dataset;
 import ai.djl.training.evaluator.Evaluator;
@@ -121,7 +122,20 @@ public class Trainer implements AutoCloseable {
                 .forEach(
                         pair -> {
                             for (Device device : devices) {
-                                parameterStore.getValue(pair.getValue(), device, true);
+                                try {
+                                    parameterStore.getValue(pair.getValue(), device, true);
+                                } catch (UninitializedParameterException e) {
+                                    throw new IllegalStateException(
+                                            "Failed to initialize parameter: "
+                                                    + pair.getKey()
+                                                    + ".\n"
+                                                    + "If you are defining a Block extending"
+                                                    + " AbstractBlock, check that you are"
+                                                    + " initializing all child blocks as part of"
+                                                    + " the overload for"
+                                                    + " AbstractBlock.initializeChildBlocks().",
+                                            e);
+                                }
                             }
                         });
     }
diff --git a/api/src/main/java/ai/djl/training/loss/TabNetLoss.java b/api/src/main/java/ai/djl/training/loss/TabNetLoss.java
new file mode 100644
index 000000000..296f4b93d
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/loss/TabNetLoss.java
@@ -0,0 +1,44 @@
+/*
+ * 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.training.loss;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * Calculates the loss for tabNet. Actually, loss has been calculated through the forward function
+ * of tabNet. What's done here is just getting the loss function from prediction.
+ */
+public class TabNetLoss extends Loss {
+    /** Calculates the loss of a TabNet instance. */
+    public TabNetLoss() {
+        this("TabNetLoss");
+    }
+
+    /**
+     * Calculates the loss of a TabNet instance.
+     *
+     * @param name the name of the loss function
+     */
+    public TabNetLoss(String name) {
+        super(name);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        // loss is already calculated inside the forward of tabNet
+        // so here we just need to get it out from prediction
+        return predictions.get(1);
+    }
+}
diff --git a/api/src/main/java/ai/djl/translate/Ensembleable.java b/api/src/main/java/ai/djl/translate/Ensembleable.java
new file mode 100644
index 000000000..12e19f401
--- /dev/null
+++ b/api/src/main/java/ai/djl/translate/Ensembleable.java
@@ -0,0 +1,50 @@
+/*
+ * 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 java.util.Iterator;
+import java.util.List;
+
+/**
+ * Represents a class that can be ensembled (or averaged).
+ *
+ * <p>Typically, ensembling is used for the output of models/translators. By averaging multiple
+ * models, it is often possible to get greater accuracy then running each model individually.
+ */
+public interface Ensembleable<T> {
+
+    /**
+     * Creates an ensembled output with a list of outputs.
+     *
+     * @param iterator the outputs to ensemble with. It uses the caller class to determine how to
+     *     ensemble.
+     * @return the ensembled (averaged) output
+     */
+    T ensembleWith(Iterator<T> iterator);
+
+    /**
+     * Finds the ensemble of a list of outputs.
+     *
+     * @param outputs the outputs to ensemble.
+     * @param <T> the type of object to ensemble. Usually also the type returned
+     * @return the ensembled (averaged) output
+     */
+    static <T extends Ensembleable<T>> T ensemble(List<T> outputs) {
+        if (outputs.isEmpty()) {
+            return null;
+        }
+        Iterator<T> it = outputs.iterator();
+        T item = it.next();
+        return item.ensembleWith(it);
+    }
+}
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
index d4d27b9ed..9d8037cfa 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/zoo/OrtModelZoo.java
@@ -30,6 +30,7 @@ public class OrtModelZoo extends ModelZoo {
 
     OrtModelZoo() {
         addModel(REPOSITORY.model(CV.IMAGE_CLASSIFICATION, GROUP_ID, "resnet", "0.0.1"));
+        addModel(REPOSITORY.model(CV.OBJECT_DETECTION, GROUP_ID, "yolo5s", "0.0.1"));
         addModel(REPOSITORY.model(Tabular.SOFTMAX_REGRESSION, GROUP_ID, "iris_flowers", "0.0.1"));
     }
 
diff --git a/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java b/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java
index c3eb83148..ca2458bcc 100644
--- a/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java
+++ b/examples/src/main/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModel.java
@@ -88,7 +88,7 @@ public final class ObjectDetectionWithTensorflowSavedModel {
                         .setTypes(Image.class, DetectedObjects.class)
                         .optModelUrls(modelUrl)
                         // saved_model.pb file is in the subfolder of the model archive file
-                        .optModelName("ssd_mobilenet_v2_320x320_coco17_tpu-8/saved_model")
+                        .optModelName("saved_model")
                         .optTranslator(new MyTranslator())
                         .optEngine("TensorFlow")
                         .optProgress(new ProgressBar())
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
index 0c505e3ad..282d8d247 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TextEmbeddingTranslator.java
@@ -68,8 +68,7 @@ public class TextEmbeddingTranslator implements Translator<String, float[]> {
         NDManager manager = ctx.getNDManager();
         NDArray inputAttentionMask = manager.create(attentionMask).toType(DataType.FLOAT32, true);
         long[] shape = embeddings.getShape().getShape();
-        inputAttentionMask = inputAttentionMask.tile(shape[shape.length - 1]);
-        inputAttentionMask = inputAttentionMask.reshape(embeddings.getShape());
+        inputAttentionMask = inputAttentionMask.expandDims(-1).broadcast(shape);
         NDArray inputAttentionMaskSum = inputAttentionMask.sum(AXIS);
         NDArray clamp = inputAttentionMaskSum.clip(1e-9, 1e12);
         NDArray prod = embeddings.mul(inputAttentionMask);
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
new file mode 100644
index 000000000..dac0827ff
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
@@ -0,0 +1,608 @@
+/*
+ * 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.basicmodelzoo.tabular;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.AbstractBlock;
+import ai.djl.nn.Activation;
+import ai.djl.nn.Block;
+import ai.djl.nn.Blocks;
+import ai.djl.nn.LambdaBlock;
+import ai.djl.nn.ParallelBlock;
+import ai.djl.nn.SequentialBlock;
+import ai.djl.nn.core.Linear;
+import ai.djl.nn.core.SparseMax;
+import ai.djl.nn.norm.BatchNorm;
+import ai.djl.nn.norm.GhostBatchNorm;
+import ai.djl.training.ParameterStore;
+import ai.djl.util.Pair;
+import ai.djl.util.PairList;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * {@code TabNet} contains a generic implementation of TabNet adapted from
+ * https://towardsdatascience.com/implementing-tabnet-in-pytorch-fc977c383279 (Original author
+ * Samrat Thapa)
+ *
+ * <p>TabNet is a neural architecture for tabular dataset developed by the research team at Google
+ * Cloud AI. It was able to achieve state_of_the_art results on several datasets in both regression
+ * and classification problems. Another desirable feature of TabNet is interpretability. Contrary to
+ * most of deep learning, where the neural networks act like black boxes, we can interpret which
+ * features the models selects in case of TabNet.
+ *
+ * <p>see https://arxiv.org/pdf/1908.07442.pdf for more information about TabNet
+ */
+public final class TabNet extends AbstractBlock {
+    private static final byte VERSION = 1;
+    private Block firstStep;
+    private List<Block> steps;
+    private Block fullyConnected;
+    private Block batchNorm;
+    private int numD;
+    private int numA;
+
+    /**
+     * Creates a {@link TabNet} instance with given builder.
+     *
+     * @param builder the builder to create TabNet
+     */
+    private TabNet(Builder builder) {
+        super(VERSION);
+        batchNorm =
+                addChildBlock(
+                        "batchNorm",
+                        BatchNorm.builder()
+                                .optAxis(0)
+                                .optMomentum(builder.batchNormMomentum)
+                                .build());
+        List<Block> sharedBlocks = new ArrayList<>();
+        for (int i = 0; i < builder.numShared; i++) {
+            sharedBlocks.add(
+                    addChildBlock(
+                            "sharedfc" + i,
+                            Linear.builder().setUnits(2L * (builder.numA + builder.numD)).build()));
+        }
+
+        firstStep =
+                addChildBlock(
+                        "featureTransformer",
+                        featureTransformer(
+                                sharedBlocks,
+                                builder.numD + builder.numA,
+                                builder.numIndependent,
+                                builder.virtualBatchSize,
+                                builder.batchNormMomentum));
+
+        steps = new ArrayList<>();
+        for (int i = 0; i < builder.numSteps - 1; i++) {
+            steps.add(
+                    addChildBlock(
+                            "steps" + (i + 1),
+                            new DecisionStep(
+                                    builder.inputDim,
+                                    builder.numD,
+                                    builder.numA,
+                                    sharedBlocks,
+                                    builder.numIndependent,
+                                    builder.virtualBatchSize,
+                                    builder.batchNormMomentum)));
+        }
+        fullyConnected =
+                addChildBlock(
+                        "fullyConnected", Linear.builder().setUnits(builder.finalOutDim).build());
+        this.numD = builder.numD;
+        this.numA = builder.numA;
+    }
+
+    /**
+     * Applies tabNetGLU activation(which is mostly used in tabNet) on the input {@link NDArray}.
+     *
+     * @param array the input {@link NDArray}
+     * @param units the half number of the resultant features
+     * @return the {@link NDArray} after applying tabNetGLU function
+     */
+    public static NDArray tabNetGLU(NDArray array, int units) {
+        return array.get(":,:{}", units).mul(Activation.sigmoid(array.get(":, {}:", units)));
+    }
+
+    /**
+     * Applies tabNetGLU activation(which is mostly used in tabNet) on the input singleton {@link
+     * NDList}.
+     *
+     * @param arrays the input singleton {@link NDList}
+     * @param units the half number of the resultant features
+     * @return the singleton {@link NDList} after applying tabNetGLU function
+     */
+    public static NDList tabNetGLU(NDList arrays, int units) {
+        return new NDList(tabNetGLU(arrays.singletonOrThrow(), units));
+    }
+
+    /**
+     * Creates a {@link LambdaBlock} that applies the {@link #tabNetGLU(NDArray, int)} activation
+     * function in its forward function.
+     *
+     * @param units the half number of feature
+     * @return {@link LambdaBlock} that applies the {@link #tabNetGLU(NDArray, int)} activation
+     *     function
+     */
+    public static Block tabNetGLUBlock(int units) {
+        return new LambdaBlock(arrays -> tabNetGLU(arrays, units), "tabNetGLU");
+    }
+
+    /**
+     * Creates a FC-BN-GLU block used in tabNet. In order to do GLU, we double the dimension of the
+     * input features to the GLU using a fc layer.
+     *
+     * @param sharedBlock the shared fully connected layer
+     * @param outDim the output feature dimension
+     * @param virtualBatchSize the virtualBatchSize
+     * @param batchNormMomentum the momentum used for ghost batchNorm layer
+     * @return a FC-BN-GLU block
+     */
+    public static Block gluBlock(
+            Block sharedBlock, int outDim, int virtualBatchSize, float batchNormMomentum) {
+        SequentialBlock featureBlock = new SequentialBlock();
+        int units = 2 * outDim;
+        if (sharedBlock == null) {
+            featureBlock.add(Linear.builder().setUnits(units).build());
+        } else {
+            featureBlock.add(sharedBlock);
+        }
+        featureBlock
+                .add(
+                        GhostBatchNorm.builder()
+                                .optAxis(0)
+                                .optVirtualBatchSize(virtualBatchSize)
+                                .optMomentum(batchNormMomentum)
+                                .build())
+                .add(tabNetGLUBlock(outDim));
+
+        return featureBlock;
+    }
+
+    /**
+     * Creates a featureTransformer Block. The feature transformer is where all the selected
+     * features are processed to generate the final output.
+     *
+     * @param sharedBlocks the sharedBlocks of feature transformer
+     * @param outDim the output dimension of feature transformer
+     * @param numIndependent the number of independent blocks of feature transformer
+     * @param virtualBatchSize the virtual batch size for ghost batch norm
+     * @param batchNormMomentum the momentum for batch norm layer
+     * @return a feature transformer
+     */
+    public static Block featureTransformer(
+            List<Block> sharedBlocks,
+            int outDim,
+            int numIndependent,
+            int virtualBatchSize,
+            float batchNormMomentum) {
+        List<Block> allBlocks = new ArrayList<>();
+        if (!sharedBlocks.isEmpty()) {
+            for (Block sharedBlock : sharedBlocks) {
+                allBlocks.add(gluBlock(sharedBlock, outDim, virtualBatchSize, batchNormMomentum));
+            }
+        }
+        for (int i = 0; i < numIndependent; i++) {
+            allBlocks.add(gluBlock(null, outDim, virtualBatchSize, batchNormMomentum));
+        }
+
+        SequentialBlock featureBlocks = new SequentialBlock();
+        int startIndex = 0;
+        if (!sharedBlocks.isEmpty()) {
+            startIndex = 1;
+            featureBlocks.add(allBlocks.get(0));
+        }
+        for (int i = startIndex; i < allBlocks.size(); i++) {
+            featureBlocks.add(
+                    new ParallelBlock(
+                            ndLists -> {
+                                NDList unit = ndLists.get(0);
+                                NDList parallel = ndLists.get(1);
+                                return new NDList(
+                                        NDArrays.add(
+                                                        unit.singletonOrThrow(),
+                                                        parallel.singletonOrThrow())
+                                                .mul(Math.sqrt(0.5)));
+                            },
+                            Arrays.asList(allBlocks.get(i), Blocks.identityBlock())));
+        }
+        return featureBlocks;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected NDList forwardInternal(
+            ParameterStore parameterStore,
+            NDList inputs,
+            boolean training,
+            PairList<String, Object> params) {
+        NDManager manager = inputs.getManager();
+        NDArray input = inputs.singletonOrThrow();
+        NDArray x =
+                batchNorm.forward(parameterStore, new NDList(input), training).singletonOrThrow();
+        NDArray xa =
+                firstStep
+                        .forward(parameterStore, new NDList(x), training)
+                        .singletonOrThrow()
+                        .get(":," + this.numD + ":");
+        NDArray sparseLoss = null;
+        NDArray out = null;
+        NDArray priors = manager.ones(x.getShape());
+        for (Block step : steps) {
+            NDList tempRes = step.forward(parameterStore, new NDList(x, xa, priors), training);
+            NDArray xte = tempRes.get(0);
+            NDArray loss = tempRes.get(1);
+            if (out == null) {
+                out = Activation.relu(xte.get(":,:" + this.numD));
+            } else {
+                out = out.add(Activation.relu(xte.get(":,:" + this.numD)));
+            }
+            xa = xte.get(":," + this.numD + ":");
+            sparseLoss = sparseLoss == null ? loss : sparseLoss.add(loss);
+        }
+        NDArray finalOutput =
+                fullyConnected
+                        .forward(parameterStore, new NDList(out), training)
+                        .singletonOrThrow();
+        return new NDList(finalOutput, sparseLoss);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        Shape[] shapes = inputShapes;
+        Shape[] xShapes = batchNorm.getOutputShapes(shapes);
+
+        Shape[] xaShapes = firstStep.getOutputShapes(xShapes); // input shape for xa
+        xaShapes[0] = Shape.update(xaShapes[0], xaShapes[0].dimension() - 1, this.numA);
+        shapes =
+                new Shape[] {
+                    xShapes[0], xaShapes[0], xShapes[0]
+                }; // shape of priors should be the same as x
+        Shape outputShape = new Shape();
+        Shape lossShape = new Shape();
+        for (Block step : steps) {
+            Shape[] outputShapes = step.getOutputShapes(shapes);
+            outputShape = Shape.update(outputShapes[0], outputShapes[0].dimension() - 1, numD);
+            lossShape = outputShapes[1];
+        }
+        outputShape = fullyConnected.getOutputShapes(new Shape[] {outputShape})[0];
+        return new Shape[] {outputShape, lossShape};
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected void initializeChildBlocks(
+            NDManager manager, DataType dataType, Shape... inputShapes) {
+        Shape[] shapes = inputShapes;
+        batchNorm.initialize(manager, dataType, shapes);
+        Shape[] xShapes = batchNorm.getOutputShapes(shapes);
+        firstStep.initialize(manager, dataType, xShapes);
+        Shape[] xaShapes = firstStep.getOutputShapes(xShapes); // input shape for xa
+        xaShapes[0] = Shape.update(xaShapes[0], xaShapes[0].dimension() - 1, this.numD);
+        shapes =
+                new Shape[] {
+                    xShapes[0], xaShapes[0], xShapes[0]
+                }; // shape of priors should be the same as x
+        Shape outputShape = new Shape();
+        for (Block step : steps) {
+            step.initialize(manager, dataType, shapes);
+            Shape[] outputShapes = step.getOutputShapes(shapes);
+            outputShape = Shape.update(outputShapes[0], outputShapes[0].dimension() - 1, numD);
+        }
+        fullyConnected.initialize(manager, dataType, outputShape);
+    }
+
+    /**
+     * Creates a builder to build a {@link TabNet}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * AttentionTransformer is where the tabNet models learn the relationship between relevant
+     * features, and decides which features to pass on to the feature transformer of the current
+     * decision step.
+     */
+    public static final class AttentionTransformer extends AbstractBlock {
+        private static final Byte VERSION = 1;
+        private Block fullyConnected;
+        private Block batchNorm;
+        private Block sparseMax;
+
+        /**
+         * Creates an attentionTransformer Block with given parameters.
+         *
+         * @param inputDim the input Dimension of the TabNet
+         * @param virtualBatchSize the virtual batch size for ghost batchNorm
+         * @param batchNormMomentum the momentum for batchNorm layer
+         */
+        private AttentionTransformer(int inputDim, int virtualBatchSize, float batchNormMomentum) {
+            super(VERSION);
+            fullyConnected =
+                    addChildBlock("fullyConnected", Linear.builder().setUnits(inputDim).build());
+            batchNorm =
+                    addChildBlock(
+                            "ghostBatchNorm",
+                            GhostBatchNorm.builder()
+                                    .optAxis(0)
+                                    .optVirtualBatchSize(virtualBatchSize)
+                                    .optMomentum(batchNormMomentum)
+                                    .build());
+            sparseMax = addChildBlock("sparseMax", new SparseMax(-1, 8));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected NDList forwardInternal(
+                ParameterStore parameterStore,
+                NDList inputs,
+                boolean training,
+                PairList<String, Object> params) {
+            NDArray x = inputs.get(0);
+            NDArray priors = inputs.get(1);
+            NDList x1 = fullyConnected.forward(parameterStore, new NDList(x), training);
+            NDList x2 = batchNorm.forward(parameterStore, x1, training);
+            return sparseMax.forward(
+                    parameterStore, new NDList(x2.singletonOrThrow().mul(priors)), training);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Shape[] getOutputShapes(Shape[] inputShapes) {
+            Shape[] shapes = {inputShapes[0]};
+            for (Pair<String, Block> child : getChildren()) {
+                shapes = child.getValue().getOutputShapes(shapes);
+            }
+            return shapes;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected void initializeChildBlocks(
+                NDManager manager, DataType dataType, Shape... inputShapes) {
+            Shape[] shapes = {inputShapes[0]};
+            for (Block child : getChildren().values()) {
+                child.initialize(manager, dataType, shapes);
+                shapes = child.getOutputShapes(shapes);
+            }
+        }
+    }
+
+    /** DecisionStep is just combining featureTransformer and attentionTransformer together. */
+    public static final class DecisionStep extends AbstractBlock {
+        private static final Byte VERSION = 1;
+        private Block featureTransformer;
+        private Block attentionTransformer;
+
+        /**
+         * Creates a {@link DecisionStep} with given parameters.
+         *
+         * @param inputDim the number of input dimension for attentionTransformer
+         * @param numD the number of dimension except attentionTransformer
+         * @param numA the number of dimension for attentionTransformer
+         * @param shared the shared fullyConnected layers
+         * @param nInd the number of independent fullyConnected layers
+         * @param virtualBatchSize the virtual batch size
+         * @param batchNormMomentum the momentum for batchNorm layer
+         */
+        public DecisionStep(
+                int inputDim,
+                int numD,
+                int numA,
+                List<Block> shared,
+                int nInd,
+                int virtualBatchSize,
+                float batchNormMomentum) {
+            super(VERSION);
+            this.featureTransformer =
+                    addChildBlock(
+                            "featureTransformer",
+                            featureTransformer(
+                                    shared,
+                                    numD + numA,
+                                    nInd,
+                                    virtualBatchSize,
+                                    batchNormMomentum));
+            this.attentionTransformer =
+                    addChildBlock(
+                            "attentionTransformer",
+                            new AttentionTransformer(
+                                    inputDim, virtualBatchSize, batchNormMomentum));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected NDList forwardInternal(
+                ParameterStore parameterStore,
+                NDList inputs,
+                boolean training,
+                PairList<String, Object> params) {
+            NDArray x = inputs.get(0);
+            NDArray a = inputs.get(1);
+            NDArray priors = inputs.get(2);
+            NDList mask =
+                    attentionTransformer.forward(parameterStore, new NDList(a, priors), training);
+            NDArray sparseLoss =
+                    mask.singletonOrThrow()
+                            .mul(-1)
+                            .mul(NDArrays.add(mask.singletonOrThrow(), 1e-10).log())
+                            .mean();
+            NDList x1 = featureTransformer.forward(parameterStore, new NDList(x), training);
+            return new NDList(x1.singletonOrThrow(), sparseLoss);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Shape[] getOutputShapes(Shape[] inputShapes) {
+            Shape[] xShape = {inputShapes[0]};
+            Shape[] aShape = {inputShapes[1], inputShapes[2]};
+            Shape[] x1Shape = featureTransformer.getOutputShapes(xShape);
+            Shape[] lossShape = attentionTransformer.getOutputShapes(aShape);
+            return new Shape[] {x1Shape[0], lossShape[0]};
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected void initializeChildBlocks(
+                NDManager manager, DataType dataType, Shape... inputShapes) {
+            Shape[] xShape = {inputShapes[0]};
+            Shape[] aShape = {inputShapes[1], inputShapes[2]};
+            this.attentionTransformer.initialize(manager, dataType, aShape);
+            this.featureTransformer.initialize(manager, dataType, xShape);
+        }
+    }
+
+    /** The Builder to construct a {@link TabNet} object. */
+    public static class Builder {
+        int inputDim = 128;
+        int finalOutDim = 10;
+        int numD = 64;
+        int numA = 64;
+        int numShared = 2;
+        int numIndependent = 2;
+        int numSteps = 5;
+        int virtualBatchSize = 128;
+        float batchNormMomentum = 0.9f;
+
+        /**
+         * Sets the input dimension of TabNet.
+         *
+         * @param inputDim the input dimension
+         * @return this {@code Builder}
+         */
+        public Builder setInputDim(int inputDim) {
+            this.inputDim = inputDim;
+            return this;
+        }
+
+        /**
+         * Sets the output dimension for TabNet.
+         *
+         * @param outDim the output dimension
+         * @return this {@code Builder}
+         */
+        public Builder setOutDim(int outDim) {
+            this.finalOutDim = outDim;
+            return this;
+        }
+
+        /**
+         * Sets the number of dimension except attentionTransformer.
+         *
+         * @param numD the number of dimension except attentionTransformer
+         * @return this {@code Builder}
+         */
+        public Builder optNumD(int numD) {
+            this.numD = numD;
+            return this;
+        }
+
+        /**
+         * Sets the number of dimension for attentionTransformer.
+         *
+         * @param numA the number of dimension for attentionTransformer
+         * @return this {@code Builder}
+         */
+        public Builder optNumA(int numA) {
+            this.numA = numA;
+            return this;
+        }
+
+        /**
+         * Sets the number of shared fullyConnected layers.
+         *
+         * @param numShared the number of shared fullyConnected layers
+         * @return this {@code Builder}
+         */
+        public Builder optNumShared(int numShared) {
+            this.numShared = numShared;
+            return this;
+        }
+
+        /**
+         * Sets the number of independent fullyConnected layers.
+         *
+         * @param numIndependent the number of independent fullyConnected layers
+         * @return this {@code Builder}
+         */
+        public Builder optNumIndependent(int numIndependent) {
+            this.numIndependent = numIndependent;
+            return this;
+        }
+
+        /**
+         * Sets the number of decision steps for tabNet.
+         *
+         * @param numSteps the number of decision steps for tabNet
+         * @return this {@code Builder}
+         */
+        public Builder optNumSteps(int numSteps) {
+            this.numSteps = numSteps;
+            return this;
+        }
+
+        /**
+         * Sets the virtual batch size for ghost batch norm.
+         *
+         * @param virtualBatchSize the virtual batch size
+         * @return this {@code Builder}
+         */
+        public Builder optVirtualBatchSize(int virtualBatchSize) {
+            this.virtualBatchSize = virtualBatchSize;
+            return this;
+        }
+
+        /**
+         * Sets the momentum for batchNorm layer.
+         *
+         * @param batchNormMomentum the momentum for batchNormLayer
+         * @return this {@code Builder}
+         */
+        public Builder optBatchNormMomentum(float batchNormMomentum) {
+            this.batchNormMomentum = batchNormMomentum;
+            return this;
+        }
+
+        /**
+         * Builds an attentionTransformer with given parameter for test.
+         *
+         * @param units the number of test units
+         * @return an attentionTransformer Block
+         */
+        public Block buildAttentionTransformer(int units) {
+            return new AttentionTransformer(10, virtualBatchSize, batchNormMomentum);
+        }
+
+        /**
+         * Builds a TabNet with given {@code Builder}.
+         *
+         * @return a tabNetBlock
+         */
+        public Block build() {
+            return new TabNet(this);
+        }
+    }
+}
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
new file mode 100644
index 000000000..ff3201da0
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
@@ -0,0 +1,18 @@
+/*
+ * 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.
+ */
+
+/**
+ * Contains classes for the computer vision section ({@link ai.djl.Application.Tabular}) of the
+ * {@link ai.djl.basicmodelzoo.BasicModelZoo}.
+ */
+package ai.djl.basicmodelzoo.tabular;


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/EnsembleTest.java b/api/src/test/java/ai/djl/modality/EnsembleTest.java
new file mode 100644
index 000000000..2a49a314e
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/EnsembleTest.java
@@ -0,0 +1,53 @@
+/*
+ * 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;
+
+import ai.djl.translate.Ensembleable;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+public class EnsembleTest {
+
+    @Test
+    public void testEnsembleClassifications() {
+        Classifications actual = Ensembleable.ensemble(Collections.<Classifications>emptyList());
+        Assert.assertNull(actual);
+
+        List<String> classNames = Arrays.asList("a", "b", "c");
+        List<Classifications> list = new ArrayList<>();
+        for (int i = 0; i < 3; i++) {
+            double prob = i;
+            List<Double> probs = Arrays.asList(prob, prob, prob);
+            list.add(new Classifications(classNames, probs));
+        }
+        actual = Ensembleable.ensemble(list);
+        Assert.assertNotNull(actual);
+        Assert.assertEquals(actual.getClassNames(), classNames);
+        Assert.assertEquals(actual.getProbabilities(), Arrays.asList(1.0, 1.0, 1.0));
+
+        Assert.assertThrows(
+                () -> {
+                    List<String> badClassNames = Arrays.asList("b", "c", "d");
+                    List<Classifications> badList = new ArrayList<>();
+                    badList.addAll(list);
+                    badList.add(new Classifications(badClassNames, Arrays.asList(1.0, 1.0, 1.0)));
+                    Ensembleable.ensemble(badList);
+                });
+    }
+}
diff --git a/engines/onnxruntime/onnxruntime-engine/src/test/resources/mlrepo/model/cv/object_detection/ai/djl/onnxruntime/yolo5s/metadata.json b/engines/onnxruntime/onnxruntime-engine/src/test/resources/mlrepo/model/cv/object_detection/ai/djl/onnxruntime/yolo5s/metadata.json
new file mode 100644
index 000000000..cd9c0240e
--- /dev/null
+++ b/engines/onnxruntime/onnxruntime-engine/src/test/resources/mlrepo/model/cv/object_detection/ai/djl/onnxruntime/yolo5s/metadata.json
@@ -0,0 +1,38 @@
+{
+  "metadataVersion": "0.2",
+  "resourceType": "model",
+  "application": "cv/object_detection",
+  "groupId": "ai.djl.onnxruntime",
+  "artifactId": "yolo5s",
+  "name": "YoloV5s",
+  "description": "YoloV5 Model",
+  "website": "http://www.djl.ai/engines/onnxruntime/model-zoo",
+  "licenses": {
+    "license": {
+      "name": "The Apache License, Version 2.0",
+      "url": "https://www.apache.org/licenses/LICENSE-2.0"
+    }
+  },
+  "artifacts": [
+    {
+      "version": "0.0.1",
+      "snapshot": false,
+      "name": "yolo5s",
+      "arguments": {
+        "width": 640,
+        "height": 640,
+        "resize": true,
+        "normalize": true,
+        "translatorFactory": "ai.djl.modality.cv.translator.YoloV5TranslatorFactory"
+      },
+      "files": {
+        "model": {
+          "uri": "0.0.1/yolov5s.zip",
+          "name": "",
+          "sha1Hash": "c30c24c8b67e0105ad468753efd2e4921e732dae",
+          "size": 24704995
+        }
+      }
+    }
+  ]
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
new file mode 100644
index 000000000..5dd6d2344
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
@@ -0,0 +1,112 @@
+/*
+ * 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.integration.tests.model_zoo.tabular;
+
+import ai.djl.Model;
+import ai.djl.basicmodelzoo.tabular.TabNet;
+import ai.djl.engine.Engine;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.nn.Parameter;
+import ai.djl.testing.Assertions;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingConfig;
+import ai.djl.training.dataset.Batch;
+import ai.djl.training.initializer.Initializer;
+import ai.djl.training.loss.Loss;
+import ai.djl.training.loss.TabNetLoss;
+import ai.djl.translate.Batchifier;
+import ai.djl.util.PairList;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class TabNetTest {
+    @Test
+    public void testTabNetGLU() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.l2Loss())
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(TabNet.tabNetGLUBlock(1));
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(new Shape(4));
+                NDManager manager = trainer.getManager();
+                NDArray data = manager.create(new float[] {1, 2, 3, 4});
+                data = data.reshape(2, 2);
+                // expected calculated through pytorch
+                NDArray expected = manager.create(new float[] {0.8808f, 2.946f});
+                NDArray result = trainer.forward(new NDList(data)).singletonOrThrow().squeeze();
+                Assertions.assertAlmostEquals(result, expected);
+            }
+        }
+    }
+
+    @Test
+    public void testTrainingAndLogic() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(new TabNetLoss())
+                        .optDevices(Engine.getInstance().getDevices(2));
+
+        Block tabNet = TabNet.builder().setOutDim(10).build();
+        try (Model model = Model.newInstance("tabNet")) {
+            model.setBlock(tabNet);
+            try (Trainer trainer = model.newTrainer(config)) {
+                int batchSize = 1;
+                Shape inputShape = new Shape(batchSize, 128);
+                trainer.initialize(inputShape);
+                NDManager manager = trainer.getManager();
+                NDArray input = manager.randomUniform(0, 1, inputShape);
+                NDArray label = manager.ones(new Shape(batchSize, 10));
+                Batch batch =
+                        new Batch(
+                                manager.newSubManager(),
+                                new NDList(input),
+                                new NDList(label),
+                                batchSize,
+                                Batchifier.STACK,
+                                Batchifier.STACK,
+                                0,
+                                0);
+                PairList<String, Parameter> parameters = tabNet.getParameters();
+                EasyTrain.trainBatch(trainer, batch);
+                trainer.step();
+                // the gamma of batchNorm Layer
+                Assert.assertEquals(
+                        parameters.get(0).getValue().getArray().getShape(), new Shape(1));
+
+                // weight of shared fullyConnected Block0
+                Assert.assertEquals(
+                        parameters.get(4).getValue().getArray().getShape(), new Shape(256, 128));
+
+                // the parameter value of a shared fc Block should be the same
+                Assert.assertEquals(parameters.get(8).getValue(), parameters.get(4).getValue());
+                Assert.assertEquals(parameters.get(32).getValue(), parameters.get(4).getValue());
+
+                // fc's weight of attention Transformer of step01
+                Assert.assertEquals(
+                        parameters.get(56).getValue().getArray().getShape(), new Shape(128, 64));
+
+                // the final fc Block
+                Assert.assertEquals(
+                        parameters.get(152).getValue().getArray().getShape(), new Shape(10, 64));
+            }
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/package-info.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/package-info.java
new file mode 100644
index 000000000..5a3c8ac57
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains tests using the engine for {@link ai.djl.basicmodelzoo.tabular}. */
+package ai.djl.integration.tests.model_zoo.tabular;