Repository: deepjavalibrary/djl
Pull request #1714 (2022-06-14 15:12:23)

add android test code
## Description ##

Add android model loading test code


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1714?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1714](https://codecov.io/gh/deepjavalibrary/djl/pull/1714?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (7e31369) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `1.64%`.
> The diff coverage is `69.69%`.

```diff
@@             Coverage Diff              @@
##             master    #1714      +/-   ##
============================================
- Coverage     72.08%   70.44%   -1.65%     
- Complexity     5126     5515     +389     
============================================
  Files           473      523      +50     
  Lines         21970    24446    +2476     
  Branches       2351     2646     +295     
============================================
+ Hits          15838    17222    +1384     
- Misses         4925     5915     +990     
- Partials       1207     1309     +102     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1714?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/1714/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/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/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/1714/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/1714/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/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <ø> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `8.33% <0.00%> (-0.50%)` | :arrow_down: |
| [.../modality/cv/translator/YoloTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/YoloV5Translator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvci5qYXZh) | `5.38% <0.00%> (-0.31%)` | :arrow_down: |
| ... and [287 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1714/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

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


Patch:
diff --git a/api/src/main/java/ai/djl/nn/core/LinearCollection.java b/api/src/main/java/ai/djl/nn/core/LinearCollection.java
new file mode 100644
index 000000000..5d69fe1c4
--- /dev/null
+++ b/api/src/main/java/ai/djl/nn/core/LinearCollection.java
@@ -0,0 +1,243 @@
+/*
+ * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.nn.core;
+
+import ai.djl.Device;
+import ai.djl.MalformedModelException;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.AbstractBlock;
+import ai.djl.nn.Block;
+import ai.djl.nn.Parameter;
+import ai.djl.training.ParameterStore;
+import ai.djl.util.PairList;
+import ai.djl.util.Preconditions;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.util.Collections;
+
+/**
+ * A LinearCollection block applies \(m\) linear transformations \(Y_i = X_i W_i + b_i\) for each
+ * \(i \in [1, \ldots, m]\) and \(m = \prod_{j=1}^t s_j\). \(t\) is typically 1, so compared to a
+ * {@link Linear} block, the involved shapes have typically one split dimension added which
+ * separates the different linear transformations from each other. Another difference to a {@link
+ * Linear} block is that the weight is not transposed (to align with the internally used algebraic
+ * operation {@link NDArray#matMul(NDArray)} ). There is currently only a single batch dimension
+ * \(x_1\) supported.
+ *
+ * <p>It has the following shapes:
+ *
+ * <ul>
+ *   <li>input X: [x_1, s_1, s_2, …, s_t, input_dim]
+ *   <li>weight W: [s_1, s_2, …, s_t, input_dim, units]
+ *   <li>Bias b: [s_1, s_2, …, s_t, units]
+ *   <li>output Y: [x_1, s_1, s_2, …, s_t, units]
+ * </ul>
+ *
+ * <p>The LinearCollection block should be constructed using {@link LinearCollection.Builder}.
+ */
+public class LinearCollection extends AbstractBlock {
+
+    private static final byte VERSION = 1;
+
+    private long units;
+    private long inputFeatures;
+
+    private Shape inputShape;
+
+    private Parameter weight;
+    private Parameter bias;
+
+    private int[] shiftBatchAxis;
+
+    private int[] reverseShiftBatchAxis;
+
+    LinearCollection(Builder builder) {
+        super(VERSION);
+        units = builder.units;
+        weight =
+                addParameter(
+                        Parameter.builder()
+                                .setName("weight")
+                                .setType(Parameter.Type.WEIGHT)
+                                .build());
+        if (builder.bias) {
+            bias =
+                    addParameter(
+                            Parameter.builder()
+                                    .setName("bias")
+                                    .setType(Parameter.Type.BIAS)
+                                    .build());
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected NDList forwardInternal(
+            ParameterStore parameterStore,
+            NDList inputs,
+            boolean training,
+            PairList<String, Object> params) {
+        NDArray input = inputs.singletonOrThrow();
+        Device device = input.getDevice();
+        NDArray weightArr = parameterStore.getValue(weight, device, training);
+        NDArray biasArr = parameterStore.getValue(bias, device, training);
+        return linear(input, weightArr, biasArr);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Shape[] getOutputShapes(Shape[] inputs) {
+        return new Shape[] {inputs[0].slice(0, inputs[0].dimension() - 1).add(units)};
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public PairList<String, Shape> describeInput() {
+        return new PairList<>(
+                Collections.singletonList("linearInput"), Collections.singletonList(inputShape));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected void beforeInitialize(Shape... inputShapes) {
+        super.beforeInitialize(inputShapes);
+        Preconditions.checkArgument(inputShapes.length == 1, "Linear block only support 1 input");
+        Shape input = inputShapes[0];
+        inputFeatures = input.slice(1, input.dimension()).size();
+        inputShape = input.slice(0, 1);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(Shape[] inputShapes) {
+        Shape input = inputShapes[0];
+        weight.setShape(input.slice(1, input.dimension()).add(units));
+        if (bias != null) {
+            bias.setShape(input.slice(1, input.dimension() - 1).add(units));
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected void saveMetadata(DataOutputStream os) throws IOException {
+        os.writeLong(units);
+        os.writeLong(inputFeatures);
+        os.write(inputShape.getEncoded());
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void loadMetadata(byte loadVersion, DataInputStream is)
+            throws IOException, MalformedModelException {
+        switch (loadVersion) {
+            case VERSION:
+                units = is.readLong();
+                inputFeatures = is.readLong();
+                break;
+            default:
+                throw new MalformedModelException("Unsupported encoding version: " + loadVersion);
+        }
+        inputShape = Shape.decode(is);
+    }
+
+    /**
+     * Applies linear transformations to the incoming data.
+     *
+     * @param input input X: [x_1, s_1, s_2, …, s_t, input_dim]
+     * @param weight weight W: [s_1, s_2, …, s_t, input_dim, units]
+     * @param bias bias b: [s_1, s_2, …, s_t, units]
+     * @return output Y: [x_1, s_1, s_2, …, s_t, units]
+     */
+    public NDList linear(NDArray input, NDArray weight, NDArray bias) {
+        if (shiftBatchAxis == null) {
+            // as the batch axis is the first axis in the shape of the input resp. output
+            // arrays, it needs to be shifted in order to bring the split axes in front resp. back
+            // again to be suitable for matMul;
+            // in case there is only one split axis, the transpose array (1,0,2) could be used for
+            // both shifts, but for the general case we calculate the transpose arrays here
+            int dim = input.getShape().dimension();
+            shiftBatchAxis = new int[dim];
+            reverseShiftBatchAxis = new int[dim];
+            for (int d = 0; d < dim - 2; d++) {
+                shiftBatchAxis[d] = d + 1;
+                reverseShiftBatchAxis[d + 1] = d;
+            }
+            shiftBatchAxis[dim - 1] = dim - 1;
+            reverseShiftBatchAxis[dim - 1] = dim - 1;
+            shiftBatchAxis[dim - 2] = 0;
+            reverseShiftBatchAxis[0] = dim - 2;
+        }
+        NDArray resultArr =
+                input.transpose(shiftBatchAxis).matMul(weight).transpose(reverseShiftBatchAxis);
+        if (bias != null) {
+            resultArr.addi(bias);
+        }
+        return new NDList(resultArr);
+    }
+
+    /**
+     * Creates a builder to build a {@code Linear}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct a {@link LinearCollection} type of {@link Block}. */
+    public static final class Builder {
+
+        private long units;
+        private boolean bias = true;
+
+        Builder() {}
+
+        /**
+         * Sets the number of output channels.
+         *
+         * @param units the number of desired output channels
+         * @return this Builder
+         */
+        public Builder setUnits(long units) {
+            this.units = units;
+            return this;
+        }
+
+        /**
+         * Sets the optional parameter that indicates whether to include a bias vector with default
+         * value of true.
+         *
+         * @param bias whether to use a bias vector parameter
+         * @return this Builder
+         */
+        public Builder optBias(boolean bias) {
+            this.bias = bias;
+            return this;
+        }
+
+        /**
+         * Returns the constructed {@code Linear}.
+         *
+         * @return the constructed {@code Linear}
+         * @throws IllegalArgumentException if all required parameters (outChannels) have not been
+         *     set
+         */
+        public LinearCollection build() {
+            Preconditions.checkArgument(units > 0, "You must specify unit");
+            return new LinearCollection(this);
+        }
+    }
+}


Test patch:
diff --git a/android/core/src/androidTest/java/ai/djl/android/core/ModelLoadingTest.java b/android/core/src/androidTest/java/ai/djl/android/core/ModelLoadingTest.java
new file mode 100644
index 000000000..1ae1b73a9
--- /dev/null
+++ b/android/core/src/androidTest/java/ai/djl/android/core/ModelLoadingTest.java
@@ -0,0 +1,55 @@
+/*
+ * 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.android.core;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.Classifications;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.modality.cv.translator.ImageClassificationTranslator;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+
+/* Integration test to check if model and inference runs fine */
+public class ModelLoadingTest {
+
+    @Test
+    public void testModelLoading() throws IOException, ModelException, TranslateException {
+        ImageClassificationTranslator translator = ImageClassificationTranslator.builder()
+                .addTransform(new ToTensor())
+                .optFlag(Image.Flag.GRAYSCALE)
+                .optApplySoftmax(true)
+                .build();
+        Criteria<Image, Classifications> criteria =
+                Criteria.builder()
+                        .setTypes(Image.class, Classifications.class)
+                        .optModelUrls("https://djl-ai.s3.amazonaws.com/resources/demo/pytorch/doodle_mobilenet.zip")
+                        .optTranslator(translator)
+                        .build();
+        Image image = ImageFactory.getInstance().fromUrl("https://resources.djl.ai/images/0.png");
+        try (ZooModel<Image, Classifications> model = ModelZoo.loadModel(criteria);
+             Predictor<Image, Classifications> predictor = model.newPredictor()) {
+            Classifications result = predictor.predict(image);
+            Assert.assertEquals("dog", result.best().getClassName());
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
index a30f8f382..83afde91f 100644
--- a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
@@ -21,6 +21,7 @@ import ai.djl.modality.nlp.embedding.TrainableWordEmbedding;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.LayoutType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
@@ -34,6 +35,7 @@ import ai.djl.nn.convolutional.Conv2d;
 import ai.djl.nn.convolutional.Conv2dTranspose;
 import ai.djl.nn.convolutional.Conv3d;
 import ai.djl.nn.core.Linear;
+import ai.djl.nn.core.LinearCollection;
 import ai.djl.nn.norm.BatchNorm;
 import ai.djl.nn.norm.Dropout;
 import ai.djl.nn.norm.GhostBatchNorm;
@@ -57,6 +59,7 @@ import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.util.Arrays;
+import java.util.stream.IntStream;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -177,6 +180,165 @@ public class BlockCoreTest {
         }
     }
 
+    @Test
+    public void testLinearCollection() throws IOException, MalformedModelException {
+
+        // used to initialize both weight and bias
+        Initializer initializer =
+                (m, s, t) -> m.arange(0, s.size(), 1, DataType.FLOAT32).reshape(s);
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.l2Loss())
+                        .optInitializer(initializer, Parameter.Type.WEIGHT)
+                        .optInitializer(initializer, Parameter.Type.BIAS);
+        int featureSize = 5;
+        int outSize = 7;
+        int batchSize = 11;
+        // number of linear transformations modeled by LinearCollection
+        int splitSize = 6;
+
+        // following ndarrays will be scoped to modelLinearCollection's ndmanager so they can be
+        // used later in the for-loop
+        NDArray expectedWeightForLinearCollection;
+        NDArray expectedBiasForLinearCollection;
+        NDArray expectedLabelsForLinearCollection;
+        NDArray dataForLinearCollection;
+
+        try (NDManager sharedManager = NDManager.newBaseManager()) {
+
+            // test expected output when using LinearCollection
+            try (Model model = Model.newInstance("model")) {
+                Block block = LinearCollection.builder().setUnits(outSize).build();
+                model.setBlock(block);
+
+                try (Trainer trainer = model.newTrainer(config)) {
+                    Shape inputShape = new Shape(batchSize, splitSize, featureSize);
+                    trainer.initialize(inputShape);
+
+                    NDManager manager = trainer.getManager();
+                    dataForLinearCollection =
+                            sharedManager.arange(inputShape.size()).reshape(inputShape);
+                    NDArray actualLabels =
+                            trainer.forward(new NDList(dataForLinearCollection)).singletonOrThrow();
+                    NDArray actualWeight = block.getParameters().get("weight").getArray();
+                    NDArray actualBias = block.getParameters().get("bias").getArray();
+                    expectedWeightForLinearCollection =
+                            initializer.initialize(
+                                    sharedManager,
+                                    new Shape(splitSize, featureSize, outSize),
+                                    DataType.FLOAT32);
+                    expectedBiasForLinearCollection =
+                            initializer.initialize(
+                                    sharedManager, new Shape(splitSize, outSize), DataType.FLOAT32);
+                    Assert.assertEquals(expectedWeightForLinearCollection, actualWeight);
+                    Assert.assertEquals(expectedBiasForLinearCollection, actualBias);
+                    expectedLabelsForLinearCollection =
+                            dataForLinearCollection
+                                    .transpose(1, 0, 2)
+                                    .matMul(expectedWeightForLinearCollection)
+                                    .transpose(1, 0, 2)
+                                    .add(expectedBiasForLinearCollection);
+                    Assert.assertEquals(expectedLabelsForLinearCollection, actualLabels);
+
+                    testEncode(manager, block);
+                }
+            }
+
+            // test expected output can be reproduced by using splitSize-many Linear blocks
+            for (int splitIndex : IntStream.range(0, splitSize).toArray()) {
+                try (Model model = Model.newInstance("model")) {
+                    Block block = Linear.builder().setUnits(outSize).build();
+                    model.setBlock(block);
+
+                    try (Trainer trainer = model.newTrainer(config)) {
+
+                        // copy initial weight/bias values from LinearCollection
+                        Shape inputShape = new Shape(batchSize, featureSize);
+                        block.setInitializer(
+                                (m, s, t) -> {
+                                    NDArray w =
+                                            expectedWeightForLinearCollection
+                                                    .get(splitIndex)
+                                                    .transpose();
+                                    w.attach(m);
+                                    return w;
+                                },
+                                Parameter.Type.WEIGHT);
+                        block.setInitializer(
+                                (m, s, t) -> {
+                                    NDArray w = expectedBiasForLinearCollection.get(splitIndex);
+                                    w.attach(m);
+                                    return w;
+                                },
+                                Parameter.Type.BIAS);
+                        trainer.initialize(inputShape);
+
+                        NDManager manager = trainer.getManager();
+                        NDArray data = dataForLinearCollection.transpose(1, 0, 2).get(splitIndex);
+                        NDArray actualLabels = trainer.forward(new NDList(data)).singletonOrThrow();
+                        NDArray actualWeight = block.getParameters().get("weight").getArray();
+                        NDArray actualBias = block.getParameters().get("bias").getArray();
+                        NDArray expectedWeight =
+                                expectedWeightForLinearCollection.get(splitIndex).transpose();
+                        NDArray expectedBias = expectedBiasForLinearCollection.get(splitIndex);
+                        Assert.assertEquals(expectedWeight, actualWeight);
+                        Assert.assertEquals(expectedBias, actualBias);
+                        NDArray expectedLabels =
+                                data.dot(expectedWeight.transpose()).add(expectedBias);
+                        Assert.assertEquals(expectedLabels, actualLabels);
+                        Assert.assertEquals(
+                                expectedLabelsForLinearCollection
+                                        .transpose(1, 0, 2)
+                                        .get(splitIndex),
+                                actualLabels);
+
+                        testEncode(manager, block);
+                    }
+                }
+            }
+
+            // test expected output when using multiple split axis
+            try (Model model = Model.newInstance("model")) {
+                Block block = LinearCollection.builder().setUnits(outSize).build();
+                model.setBlock(block);
+
+                try (Trainer trainer = model.newTrainer(config)) {
+                    long[] splitAxis = {3, 1, 2};
+                    Assert.assertEquals(splitSize, new Shape(splitAxis).size());
+                    Shape inputShape = new Shape(batchSize).add(splitAxis).add(featureSize);
+                    trainer.initialize(inputShape);
+
+                    NDManager manager = trainer.getManager();
+                    NDArray data = dataForLinearCollection.reshape(inputShape);
+                    NDArray actualLabels = trainer.forward(new NDList(data)).singletonOrThrow();
+                    NDArray actualWeight = block.getParameters().get("weight").getArray();
+                    NDArray actualBias = block.getParameters().get("bias").getArray();
+                    NDArray expectedWeight =
+                            initializer.initialize(
+                                    manager,
+                                    new Shape(splitAxis).add(featureSize, outSize),
+                                    DataType.FLOAT32);
+                    NDArray expectedBias =
+                            initializer.initialize(
+                                    manager, new Shape(splitAxis).add(outSize), DataType.FLOAT32);
+                    Assert.assertEquals(expectedWeight, actualWeight);
+                    Assert.assertEquals(expectedBias, actualBias);
+                    NDArray expectedLabels =
+                            data.transpose(1, 2, 3, 0, 4)
+                                    .matMul(expectedWeight)
+                                    .transpose(3, 0, 1, 2, 4)
+                                    .add(expectedBias);
+                    Assert.assertEquals(expectedLabels, actualLabels);
+                    Assert.assertEquals(
+                            expectedLabelsForLinearCollection.reshape(
+                                    new Shape(batchSize).add(splitAxis).add(outSize)),
+                            actualLabels);
+
+                    testEncode(manager, block);
+                }
+            }
+        }
+    }
+
     @SuppressWarnings("try")
     @Test
     public void testBatchNorm() throws IOException, MalformedModelException {