Repository: deepjavalibrary/djl
Pull request #2036 (2022-09-26 14:50:01)

Adds TabNet model for tabular dataset in modelzoo
TabNet model almost complete, except for its test and some minor bugs
What I've done in this pr:
1. add a SparseMax function in nn.core(already been merged)
2. add an activation function: tabNetGLU in nn.Activation and its corresponding test function testTabNetGLU in ActivationTest
3. add feature transformer, attention transformer and Decision Step in tabNet
4. tabNet is a Block which extends AbstractBlock and it's functions like forwardInternal, getOutputShapes and so on
5. tabNet has a builder class for users to use

TODO: test for tabNet


Patch:
diff --git a/android/README.md b/android/README.md
index 0786d5ce7..fd99b283b 100644
--- a/android/README.md
+++ b/android/README.md
@@ -18,8 +18,8 @@ In gradle, you can include the snapshot repository and add the 4 modules in your
 dependencies {
     implementation "ai.djl:api:0.19.0-SNAPSHOT"
     implementation "ai.djl.android:core:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.pytorch:pytorch-engine:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.android:pytorch-native:0.19.0-SNAPSHOT"
-    androidRuntimeOnly "ai.djl.android:onnxruntime:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.pytorch:pytorch-engine:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.android:pytorch-native:0.19.0-SNAPSHOT"
+    runtimeOnly "ai.djl.android:onnxruntime:0.19.0-SNAPSHOT"
 }
 ```
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/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/docs/development/dependency_management.md b/docs/development/dependency_management.md
index 04646a175..267995268 100644
--- a/docs/development/dependency_management.md
+++ b/docs/development/dependency_management.md
@@ -118,6 +118,8 @@ See [How to use DJL's BOM](../../bom/README.md#how-to-use-djls-bom) for detail.
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.tensorrt](https://search.maven.org/search?q=g:ai.djl.tensorrt)           | [tensorrt](../../engines/tensorrt/README.md#installation)                                                                  | TensorRT engine adapter                                                                          |
 |                                                                                  |                                                                                                                            |                                                                                                  |
+| [ai.djl.java](https://search.maven.org/search?q=g:ai.djl.java)                   | [java](https://github.com/deepjavalibrary/djl-serving/tree/master/engines/java/README.md#installation)                     | Java engine adapter                                                                              |
+|                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.python](https://search.maven.org/search?q=g:ai.djl.python)               | [python](https://github.com/deepjavalibrary/djl-serving/tree/master/engines/python/README.md#installation)                 | Python engine adapter                                                                            |
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.ml.lightgbm](https://search.maven.org/search?q=g:ai.djl.ml.lightgbm)     | [lightgbm](../../engines/ml/lightgbm/README.md#installation)                                                               | LightGBM engine adapter                                                                          |
@@ -129,7 +131,7 @@ See [How to use DJL's BOM](../../bom/README.md#how-to-use-djls-bom) for detail.
 |                                                                                  | [pytorch-native](../../android/pytorch-native/README.md)                                                                   | Contains DJL PyTorch Android native package                                                      |
 |                                                                                  | [onnxruntime](../../engines/onnxruntime/onnxruntime-android/README.md#installation)                                        | Contains DJL Onnxruntime Android package                                                         |
 |                                                                                  |                                                                                                                            |                                                                                                  |
-| [ai.djl.audio](https://search.maven.org/search?q=g:ai.djl.audio)                 | [audio](../../extensions/audio/README.md#installation)                                                                     | Contains classes that make it easy for DJL to access AWS services                                |
+| [ai.djl.audio](https://search.maven.org/search?q=g:ai.djl.audio)                 | [audio](../../extensions/audio/README.md#installation)                                                                     | DJL audio support extension                                                                      |
 |                                                                                  |                                                                                                                            |                                                                                                  |
 | [ai.djl.aws](https://search.maven.org/search?q=g:ai.djl.aws)                     | [aws-ai](../../extensions/aws-ai/README.md#installation)                                                                   | Contains classes that make it easy for DJL to access AWS services                                |
 |                                                                                  |                                                                                                                            |                                                                                                  |
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
index 78b620345..a7396dd8a 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtDeviceType.java
@@ -32,8 +32,10 @@ public final class PtDeviceType {
             return 0;
         } else if (Device.Type.GPU.equals(deviceType)) {
             return 1;
+        } else if ("mps".equals(deviceType)) {
+            return 13;
         } else {
-            throw new IllegalArgumentException("Unsupported device: " + device.toString());
+            throw new IllegalArgumentException("Unsupported device: " + device);
         }
     }
 
@@ -49,6 +51,8 @@ public final class PtDeviceType {
                 return Device.Type.CPU;
             case 1:
                 return Device.Type.GPU;
+            case 13:
+                return "mps";
             default:
                 throw new IllegalArgumentException("Unsupported deviceType: " + deviceType);
         }
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/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
new file mode 100644
index 000000000..8b4e2326f
--- /dev/null
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MpsTest.java
@@ -0,0 +1,39 @@
+/*
+ * 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.pytorch.integration;
+
+import ai.djl.Device;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+
+import org.testng.Assert;
+import org.testng.SkipException;
+import org.testng.annotations.Test;
+
+public class MpsTest {
+
+    @Test
+    public void testMps() {
+        if (!"aarch64".equals(System.getProperty("os.arch"))
+                || !System.getProperty("os.name").startsWith("Mac")) {
+            throw new SkipException("MPS test requires M1 macOS.");
+        }
+
+        Device device = Device.of("mps", -1);
+        try (NDManager manager = NDManager.newBaseManager(device)) {
+            NDArray array = manager.zeros(new Shape(1, 2));
+            Assert.assertEquals(array.getDevice().getDeviceType(), "mps");
+        }
+    }
+}
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;