/*
 *  ******************************************************************************
 *  *
 *  *
 *  * This program and the accompanying materials are made available under the
 *  * terms of the Apache License, Version 2.0 which is available at
 *  * https://www.apache.org/licenses/LICENSE-2.0.
 *  *
 *  *  See the NOTICE file distributed with this work for additional
 *  *  information regarding copyright ownership.
 *  * Unless required by applicable law or agreed to in writing, software
 *  * distributed under the License 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.
 *  *
 *  * SPDX-License-Identifier: Apache-2.0
 *  *****************************************************************************
 */
package org.eclipse.deeplearning4j.dl4jcore.nn.layers.feedforward.embedding;

import lombok.EqualsAndHashCode;
import org.deeplearning4j.BaseDL4JTest;
import org.eclipse.deeplearning4j.dl4jcore.TestUtils;
import org.deeplearning4j.nn.api.Layer;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.RNNFormat;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.conf.layers.EmbeddingLayer;
import org.deeplearning4j.nn.conf.layers.EmbeddingSequenceLayer;
import org.deeplearning4j.nn.conf.preprocessor.FeedForwardToRnnPreProcessor;
import org.deeplearning4j.nn.conf.preprocessor.RnnToFeedForwardPreProcessor;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.params.DefaultParamInitializer;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.nn.weights.embeddings.EmbeddingInitializer;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.nd4j.common.tests.tags.NativeTag;
import org.nd4j.common.tests.tags.TagNames;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.activations.impl.ActivationIdentity;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Sgd;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.List;
import java.util.Map;
import java.util.Random;
import static org.junit.jupiter.api.Assertions.*;
import org.junit.jupiter.api.DisplayName;

@DisplayName("Embedding Layer Test")
@NativeTag
@Tag(TagNames.DL4J_OLD_API)
class EmbeddingLayerTest extends BaseDL4JTest {

    @Test
    @DisplayName("Test Embedding Layer Config")
    void testEmbeddingLayerConfig() {
        for (boolean hasBias : new boolean[] { true, false }) {
            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(0, new EmbeddingLayer.Builder().hasBias(hasBias).nIn(10).nOut(5).build()).layer(1, new OutputLayer.Builder().nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).build();
            MultiLayerNetwork net = new MultiLayerNetwork(conf);
            net.init();
            Layer l0 = net.getLayer(0);
            assertEquals(org.deeplearning4j.nn.layers.feedforward.embedding.EmbeddingLayer.class, l0.getClass());
            assertEquals(10, ((FeedForwardLayer) l0.conf().getLayer()).getNIn());
            assertEquals(5, ((FeedForwardLayer) l0.conf().getLayer()).getNOut());
            INDArray weights = l0.getParam(DefaultParamInitializer.WEIGHT_KEY);
            INDArray bias = l0.getParam(DefaultParamInitializer.BIAS_KEY);
            assertArrayEquals(new long[] { 10, 5 }, weights.shape());
            if (hasBias) {
                assertArrayEquals(new long[] {  5 }, bias.shape());
            }
        }
    }

    @Test
    @DisplayName("Test Embedding Sequence Layer Config")
    void testEmbeddingSequenceLayerConfig() {
        int inputLength = 6;
        int nIn = 10;
        int embeddingDim = 5;
        int nout = 4;
        for (boolean hasBias : new boolean[] { true, false }) {
            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(new EmbeddingSequenceLayer.Builder().hasBias(hasBias).inputLength(inputLength).nIn(nIn).nOut(embeddingDim).build()).layer(new RnnOutputLayer.Builder().nIn(embeddingDim).nOut(nout).activation(Activation.SOFTMAX).build()).build();
            MultiLayerNetwork net = new MultiLayerNetwork(conf);
            net.init();
            Layer l0 = net.getLayer(0);
            assertEquals(org.deeplearning4j.nn.layers.feedforward.embedding.EmbeddingSequenceLayer.class, l0.getClass());
            assertEquals(10, ((FeedForwardLayer) l0.conf().getLayer()).getNIn());
            assertEquals(5, ((FeedForwardLayer) l0.conf().getLayer()).getNOut());
            INDArray weights = l0.getParam(DefaultParamInitializer.WEIGHT_KEY);
            INDArray bias = l0.getParam(DefaultParamInitializer.BIAS_KEY);
            assertArrayEquals(new long[] { 10, 5 }, weights.shape());
            if (hasBias) {
                assertArrayEquals(new long[] {  5 }, bias.shape());
            }
        }
    }

    @Test
    @DisplayName("Test Embedding Longer Sequences Forward Pass")
    void testEmbeddingLongerSequencesForwardPass() {
        int nClassesIn = 10;
        int inputLength = 6;
        int embeddingDim = 5;
        int nOut = 4;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(new EmbeddingSequenceLayer.Builder().inputLength(inputLength).hasBias(true).nIn(nClassesIn).nOut(embeddingDim).build()).layer(new RnnOutputLayer.Builder().nIn(embeddingDim).nOut(nOut).activation(Activation.SOFTMAX).build()).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        net.init();
        int batchSize = 3;
        INDArray inEmbedding = Nd4j.create(batchSize, inputLength);
        Random r = new Random(12345);
        for (int i = 0; i < batchSize; i++) {
            int classIdx = r.nextInt(nClassesIn);
            inEmbedding.putScalar(i, classIdx);
        }
        INDArray output = net.output(inEmbedding);
        assertArrayEquals(new long[] { batchSize, nOut, inputLength }, output.shape());
    }

    @Test
    @DisplayName("Test Embedding Single Sequence Forward Pass")
    void testEmbeddingSingleSequenceForwardPass() {
        int nClassesIn = 10;
        int embeddingDim = 5;
        int nOut = 4;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(new EmbeddingSequenceLayer.Builder().inputLength(1).hasBias(true).nIn(nClassesIn).nOut(embeddingDim).build()).layer(new RnnOutputLayer.Builder().nIn(embeddingDim).nOut(nOut).activation(Activation.SOFTMAX).build()).build();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(0, new DenseLayer.Builder().nIn(nClassesIn).nOut(5).activation(Activation.IDENTITY).build()).layer(1, new OutputLayer.Builder().nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).inputPreProcessor(0, new RnnToFeedForwardPreProcessor()).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
        net.init();
        net2.init();
        net2.setParams(net.params().dup());
        int batchSize = 3;
        INDArray inEmbedding = Nd4j.create(batchSize, 1);
        INDArray inOneHot = Nd4j.create(batchSize, nClassesIn, 1);
        Random r = new Random(12345);
        for (int i = 0; i < batchSize; i++) {
            int classIdx = r.nextInt(nClassesIn);
            inEmbedding.putScalar(i, classIdx);
            inOneHot.putScalar(new int[] { i, classIdx, 0 }, 1.0);
        }
        List<INDArray> activationsDense = net2.feedForward(inOneHot, false);
        List<INDArray> activationEmbedding = net.feedForward(inEmbedding, false);
        INDArray actD1 = activationsDense.get(1);
        INDArray actE1 = activationEmbedding.get(1).reshape(batchSize, embeddingDim);
        assertEquals(actD1, actE1);
        INDArray actD2 = activationsDense.get(2);
        INDArray actE2 = activationEmbedding.get(2).reshape(batchSize, nOut);
        assertEquals(actD2, actE2);
    }

    @Test
    @DisplayName("Test Embedding Forward Pass")
    void testEmbeddingForwardPass() {
        // With the same parameters, embedding layer should have same activations as the equivalent one-hot representation
        // input with a DenseLayer
        int nClassesIn = 10;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(0, new EmbeddingLayer.Builder().hasBias(true).nIn(nClassesIn).nOut(5).build()).layer(1, new OutputLayer.Builder().nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).build();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(0, new DenseLayer.Builder().nIn(nClassesIn).nOut(5).activation(Activation.IDENTITY).build()).layer(1, new OutputLayer.Builder().nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
        net.init();
        net2.init();
        net2.setParams(net.params().dup());
        int batchSize = 3;
        INDArray inEmbedding = Nd4j.create(batchSize, 1);
        INDArray inOneHot = Nd4j.create(batchSize, nClassesIn);
        Random r = new Random(12345);
        for (int i = 0; i < batchSize; i++) {
            int classIdx = r.nextInt(nClassesIn);
            inEmbedding.putScalar(i, classIdx);
            inOneHot.putScalar(new int[] { i, classIdx }, 1.0);
        }
        List<INDArray> activationsEmbedding = net.feedForward(inEmbedding, false);
        List<INDArray> activationsDense = net2.feedForward(inOneHot, false);
        for (int i = 1; i < 3; i++) {
            INDArray actE = activationsEmbedding.get(i);
            INDArray actD = activationsDense.get(i);
            assertEquals(actE, actD);
        }
    }

    @Test
    @DisplayName("Test Embedding Backward Pass")
    void testEmbeddingBackwardPass() {
        // With the same parameters, embedding layer should have same activations as the equivalent one-hot representation
        // input with a DenseLayer
        int nClassesIn = 10;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(0, new EmbeddingLayer.Builder().hasBias(true).nIn(nClassesIn).nOut(5).build()).layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.MCXENT).nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).build();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().activation(Activation.TANH).weightInit(WeightInit.XAVIER).list().layer(new DenseLayer.Builder().nIn(nClassesIn).nOut(5).activation(Activation.IDENTITY).build()).layer(new OutputLayer.Builder(LossFunctions.LossFunction.MCXENT).nIn(5).nOut(4).activation(Activation.SOFTMAX).build()).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
        net.init();
        net2.init();
        net2.setParams(net.params().dup());
        int batchSize = 3;
        INDArray inEmbedding = Nd4j.create(batchSize, 1);
        INDArray inOneHot = Nd4j.create(batchSize, nClassesIn);
        INDArray outLabels = Nd4j.create(batchSize, 4);
        Random r = new Random(12345);
        for (int i = 0; i < batchSize; i++) {
            int classIdx = r.nextInt(nClassesIn);
            inEmbedding.putScalar(i, classIdx);
            inOneHot.putScalar(new int[] { i, classIdx }, 1.0);
            int labelIdx = r.nextInt(4);
            outLabels.putScalar(new int[] { i, labelIdx }, 1.0);
        }
        net.setInput(inEmbedding);
        net2.setInput(inOneHot);
        net.setLabels(outLabels);
        net2.setLabels(outLabels);
        net.computeGradientAndScore();
        net2.computeGradientAndScore();
        assertEquals(net2.score(), net.score(), 1e-6);
        Map<String, INDArray> gradient = net.gradient().gradientForVariable();
        Map<String, INDArray> gradient2 = net2.gradient().gradientForVariable();
        assertEquals(gradient.size(), gradient2.size());
        for (String s : gradient.keySet()) {
            assertEquals(gradient2.get(s), gradient.get(s));
        }
    }

    @Test
    @DisplayName("Test Embedding Sequence Backward Pass")
    void testEmbeddingSequenceBackwardPass() {
        int nClassesIn = 10;
        int embeddingDim = 5;
        int nOut = 4;
        int inputLength = 1;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(new EmbeddingSequenceLayer.Builder().inputLength(inputLength).hasBias(true).nIn(nClassesIn).nOut(embeddingDim).build()).layer(new RnnOutputLayer.Builder().nIn(embeddingDim).nOut(nOut).activation(Activation.SOFTMAX).build()).setInputType(InputType.recurrent(nClassesIn, inputLength, RNNFormat.NCW)).build();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().activation(Activation.TANH).list().layer(new DenseLayer.Builder().nIn(nClassesIn).nOut(embeddingDim).activation(Activation.IDENTITY).build()).layer(new RnnOutputLayer.Builder().nIn(embeddingDim).nOut(nOut).activation(Activation.SOFTMAX).build()).setInputType(InputType.recurrent(nClassesIn, inputLength, RNNFormat.NCW)).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
        net.init();
        net2.init();
        net2.setParams(net.params().dup());
        int batchSize = 3;
        INDArray inEmbedding = Nd4j.create(batchSize, 1);
        INDArray inOneHot = Nd4j.create(batchSize, nClassesIn, 1);
        INDArray outLabels = Nd4j.create(batchSize, 4, 1);
        Random r = new Random(1337);
        for (int i = 0; i < batchSize; i++) {
            int classIdx = r.nextInt(nClassesIn);
            inEmbedding.putScalar(i, classIdx);
            inOneHot.putScalar(new int[] { i, classIdx, 0 }, 1.0);
            int labelIdx = r.nextInt(4);
            outLabels.putScalar(new int[] { i, labelIdx, 0 }, 1.0);
        }
        net.setInput(inEmbedding);
        net2.setInput(inOneHot);
        net.setLabels(outLabels);
        net2.setLabels(outLabels);
        net.computeGradientAndScore();
        net2.computeGradientAndScore();
        // System.out.println(net.score() + "\t" + net2.score());
        assertEquals(net2.score(), net.score(), 1e-6);
        Map<String, INDArray> gradient = net.gradient().gradientForVariable();
        Map<String, INDArray> gradient2 = net2.gradient().gradientForVariable();
        assertEquals(gradient.size(), gradient2.size());
        for (String s : gradient.keySet()) {
            assertEquals(gradient2.get(s), gradient.get(s));
        }
    }

    @Test
    @DisplayName("Test Embedding Layer RNN")
    void testEmbeddingLayerRNN() {
        int nClassesIn = 10;
        int batchSize = 3;
        int timeSeriesLength = 8;
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().activation(Activation.TANH).dataType(DataType.DOUBLE).list().layer(0, new EmbeddingLayer.Builder().hasBias(true).nIn(nClassesIn).nOut(5).build()).layer(1, new LSTM.Builder().nIn(5).nOut(7).activation(Activation.SOFTSIGN).build()).layer(2, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MCXENT).nIn(7).nOut(4).activation(Activation.SOFTMAX).build()).inputPreProcessor(0, new RnnToFeedForwardPreProcessor()).inputPreProcessor(1, new FeedForwardToRnnPreProcessor()).setInputType(InputType.recurrent(nClassesIn, timeSeriesLength, RNNFormat.NCW)).build();
        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().activation(Activation.TANH).weightInit(WeightInit.XAVIER).dataType(DataType.DOUBLE).list().layer(0, new DenseLayer.Builder().nIn(nClassesIn).nOut(5).activation(Activation.IDENTITY).build()).layer(1, new LSTM.Builder().nIn(5).nOut(7).activation(Activation.SOFTSIGN).build()).layer(2, new RnnOutputLayer.Builder(LossFunctions.LossFunction.MCXENT).nIn(7).nOut(4).activation(Activation.SOFTMAX).build()).inputPreProcessor(0, new RnnToFeedForwardPreProcessor()).inputPreProcessor(1, new FeedForwardToRnnPreProcessor()).setInputType(InputType.recurrent(nClassesIn, timeSeriesLength, RNNFormat.NCW)).build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
        net.init();
        net2.init();
        net2.setParams(net.params().dup());
        ;
        INDArray inEmbedding = Nd4j.create(batchSize, 1, timeSeriesLength);
        INDArray inOneHot = Nd4j.create(batchSize, nClassesIn, timeSeriesLength);
        INDArray outLabels = Nd4j.create(batchSize, 4, timeSeriesLength);
        Random r = new Random(12345);
        for (int i = 0; i < batchSize; i++) {
            for (int j = 0; j < timeSeriesLength; j++) {
                int classIdx = r.nextInt(nClassesIn);
                inEmbedding.putScalar(new int[] { i, 0, j }, classIdx);
                inOneHot.putScalar(new int[] { i, classIdx, j }, 1.0);
                int labelIdx = r.nextInt(4);
                outLabels.putScalar(new int[] { i, labelIdx, j }, 1.0);
            }
        }
        net.setInput(inEmbedding);
        net2.setInput(inOneHot);
        net.setLabels(outLabels);
        net2.setLabels(outLabels);
        net.computeGradientAndScore();
        net2.computeGradientAndScore();
        // System.out.println(net.score() + "\t" + net2.score());
        assertEquals(net2.score(), net.score(), 1e-5);
        Map<String, INDArray> gradient = net.gradient().gradientForVariable();
        Map<String, INDArray> gradient2 = net2.gradient().gradientForVariable();
        assertEquals(gradient.size(), gradient2.size());
        for (String s : gradient.keySet()) {
            assertEquals(gradient2.get(s), gradient.get(s));
        }
    }

    @Test
    @DisplayName("Test Embedding Layer With Masking")
    void testEmbeddingLayerWithMasking() {
        // Idea: have masking on the input with an embedding and dense layers on input
        // Ensure that the parameter gradients for the inputs don't depend on the inputs when inputs are masked
        int[] miniBatchSizes = { 1, 2, 5 };
        int nIn = 2;
        Random r = new Random(12345);
        int numInputClasses = 10;
        int timeSeriesLength = 5;
        for (DataType maskDtype : new DataType[] { DataType.FLOAT, DataType.DOUBLE, DataType.INT }) {
            for (int nExamples : miniBatchSizes) {
                Nd4j.getRandom().setSeed(12345);
                MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).updater(new Sgd(0.1)).seed(12345).list().layer(0, new EmbeddingLayer.Builder().hasBias(true).activation(Activation.TANH).nIn(numInputClasses).nOut(5).build()).layer(1, new DenseLayer.Builder().activation(Activation.TANH).nIn(5).nOut(4).build()).layer(2, new LSTM.Builder().activation(Activation.TANH).nIn(4).nOut(3).build()).layer(3, new RnnOutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).inputPreProcessor(0, new RnnToFeedForwardPreProcessor()).inputPreProcessor(2, new FeedForwardToRnnPreProcessor()).setInputType(InputType.recurrent(numInputClasses, timeSeriesLength, RNNFormat.NCW)).build();
                MultiLayerNetwork net = new MultiLayerNetwork(conf);
                net.init();
                MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).updater(new Sgd(0.1)).seed(12345).list().layer(0, new DenseLayer.Builder().activation(Activation.TANH).nIn(numInputClasses).nOut(5).build()).layer(1, new DenseLayer.Builder().activation(Activation.TANH).nIn(5).nOut(4).build()).layer(2, new LSTM.Builder().activation(Activation.TANH).nIn(4).nOut(3).build()).layer(3, new RnnOutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).inputPreProcessor(0, new RnnToFeedForwardPreProcessor()).inputPreProcessor(2, new FeedForwardToRnnPreProcessor()).setInputType(InputType.recurrent(numInputClasses, timeSeriesLength, RNNFormat.NCW)).build();
                MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
                net2.init();
                net2.setParams(net.params().dup());
                INDArray inEmbedding = Nd4j.zeros(nExamples, 1, timeSeriesLength);
                INDArray inDense = Nd4j.zeros(nExamples, numInputClasses, timeSeriesLength);
                INDArray labels = Nd4j.zeros(nExamples, 4, timeSeriesLength);
                for (int i = 0; i < nExamples; i++) {
                    for (int j = 0; j < timeSeriesLength; j++) {
                        int inIdx = r.nextInt(numInputClasses);
                        inEmbedding.putScalar(new int[] { i, 0, j }, inIdx);
                        inDense.putScalar(new int[] { i, inIdx, j }, 1.0);
                        int outIdx = r.nextInt(4);
                        labels.putScalar(new int[] { i, outIdx, j }, 1.0);
                    }
                }
                INDArray inputMask = Nd4j.zeros(maskDtype, nExamples, timeSeriesLength);
                for (int i = 0; i < nExamples; i++) {
                    for (int j = 0; j < timeSeriesLength; j++) {
                        inputMask.putScalar(new int[] { i, j }, (r.nextBoolean() ? 1.0 : 0.0));
                    }
                }
                net.setLayerMaskArrays(inputMask, null);
                net2.setLayerMaskArrays(inputMask, null);
                List<INDArray> actEmbedding = net.feedForward(inEmbedding, false);
                List<INDArray> actDense = net2.feedForward(inDense, false);
                for (int i = 1; i < actEmbedding.size(); i++) {
                    assertEquals(actDense.get(i), actEmbedding.get(i));
                }
                net.setLabels(labels);
                net2.setLabels(labels);
                net.computeGradientAndScore();
                net2.computeGradientAndScore();
                // System.out.println(net.score() + "\t" + net2.score());
                assertEquals(net2.score(), net.score(), 1e-5);
                Map<String, INDArray> gradients = net.gradient().gradientForVariable();
                Map<String, INDArray> gradients2 = net2.gradient().gradientForVariable();
                assertEquals(gradients.keySet(), gradients2.keySet());
                for (String s : gradients.keySet()) {
                    assertEquals(gradients2.get(s), gradients.get(s));
                }
            }
        }
    }

    @Test
    @DisplayName("Test W 2 V Inits")
    void testW2VInits() {
        Nd4j.setDefaultDataTypes(DataType.FLOAT, DataType.FLOAT);
        for (int i = 0; i < 2; i++) {
            INDArray vectors = Nd4j.linspace(1, 15, 15, DataType.FLOAT).reshape(5, 3);
            EmbeddingLayer el;
            if (i == 0) {
                el = new EmbeddingLayer.Builder().weightInit(vectors).build();
            } else {
                el = new EmbeddingLayer.Builder().weightInit(new WordVectorsMockup()).build();
            }
            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(12345).list().layer(el).layer(new DenseLayer.Builder().activation(Activation.TANH).nIn(3).nOut(3).build()).layer(new OutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).build();
            MultiLayerNetwork net = new MultiLayerNetwork(conf);
            net.init();
            INDArray w = net.getParam("0_W");
            assertEquals(vectors, w);
            TestUtils.testModelSerialization(net);
            // Test same thing for embedding sequence layer:
            EmbeddingSequenceLayer esl;
            if (i == 0) {
                esl = new EmbeddingSequenceLayer.Builder().weightInit(vectors).build();
            } else {
                esl = new EmbeddingSequenceLayer.Builder().weightInit(new WordVectorsMockup()).build();
            }
            conf = new NeuralNetConfiguration.Builder().seed(12345).list().layer(esl).layer(new GlobalPoolingLayer()).layer(new DenseLayer.Builder().activation(Activation.TANH).nIn(3).nOut(3).build()).layer(new OutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).build();
            net = new MultiLayerNetwork(conf);
            net.init();
            w = net.getParam("0_W");
            assertEquals(vectors, w);
            TestUtils.testModelSerialization(net);
        }
    }

    @Test
    @DisplayName("Test Embedding Sequence Layer With Masking")
    void testEmbeddingSequenceLayerWithMasking() {
        // Idea: have masking on the input with an embedding and dense layers on input
        // Ensure that the parameter gradients for the inputs don't depend on the inputs when inputs are masked
        int[] miniBatchSizes = { 1, 3 };
        int nIn = 2;
        Random r = new Random(12345);
        int numInputClasses = 10;
        int timeSeriesLength = 5;
        for (DataType maskDtype : new DataType[] { DataType.FLOAT, DataType.DOUBLE, DataType.INT }) {
            for (DataType inLabelDtype : new DataType[] { DataType.FLOAT, DataType.DOUBLE, DataType.INT }) {
                for (int inputRank : new int[] { 2, 3 }) {
                    for (int nExamples : miniBatchSizes) {
                        Nd4j.getRandom().setSeed(12345);
                        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).updater(new Sgd(0.1)).seed(12345).list().layer(0, new EmbeddingSequenceLayer.Builder().hasBias(true).activation(Activation.TANH).nIn(numInputClasses).nOut(5).build()).layer(1, new DenseLayer.Builder().activation(Activation.TANH).nIn(5).nOut(4).build()).layer(2, new LSTM.Builder().activation(Activation.TANH).nIn(4).nOut(3).build()).layer(3, new RnnOutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).setInputType(InputType.recurrent(numInputClasses, timeSeriesLength, RNNFormat.NCW)).build();
                        MultiLayerNetwork net = new MultiLayerNetwork(conf);
                        net.init();
                        MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).updater(new Sgd(0.1)).seed(12345).list().layer(0, new DenseLayer.Builder().activation(Activation.TANH).nIn(numInputClasses).nOut(5).build()).layer(1, new DenseLayer.Builder().activation(Activation.TANH).nIn(5).nOut(4).build()).layer(2, new LSTM.Builder().activation(Activation.TANH).nIn(4).nOut(3).dataFormat(RNNFormat.NCW).build()).layer(3, new RnnOutputLayer.Builder().lossFunction(LossFunctions.LossFunction.MSE).nIn(3).nOut(4).build()).setInputType(InputType.recurrent(numInputClasses, 1, RNNFormat.NCW)).build();
                        MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
                        net2.init();
                        net2.setParams(net.params().dup());
                        INDArray inEmbedding = Nd4j.zeros(inLabelDtype, inputRank == 2 ? new long[] { nExamples, timeSeriesLength } : new long[] { nExamples, 1, timeSeriesLength });
                        INDArray inDense = Nd4j.zeros(inLabelDtype, nExamples, numInputClasses, timeSeriesLength);
                        INDArray labels = Nd4j.zeros(inLabelDtype, nExamples, 4, timeSeriesLength);
                        for (int i = 0; i < nExamples; i++) {
                            for (int j = 0; j < timeSeriesLength; j++) {
                                int inIdx = r.nextInt(numInputClasses);
                                inEmbedding.putScalar(inputRank == 2 ? new int[] { i, j } : new int[] { i, 0, j }, inIdx);
                                inDense.putScalar(new int[] { i, inIdx, j }, 1.0);
                                int outIdx = r.nextInt(4);
                                labels.putScalar(new int[] { i, outIdx, j }, 1.0);
                            }
                        }
                        INDArray inputMask = Nd4j.zeros(maskDtype, nExamples, timeSeriesLength);
                        for (int i = 0; i < nExamples; i++) {
                            for (int j = 0; j < timeSeriesLength; j++) {
                                inputMask.putScalar(new int[] { i, j }, (r.nextBoolean() ? 1.0 : 0.0));
                            }
                        }
                        net.setLayerMaskArrays(inputMask, null);
                        net2.setLayerMaskArrays(inputMask, null);
                        List<INDArray> actEmbedding = net.feedForward(inEmbedding, false);
                        List<INDArray> actDense = net2.feedForward(inDense, false);
                        for (int i = 2; i < actEmbedding.size(); i++) {
                            // Start from layer 2: EmbeddingSequence is 3d, first dense is 2d (before reshape)
                            assertEquals(actDense.get(i), actEmbedding.get(i));
                        }
                        net.setLabels(labels);
                        net2.setLabels(labels);
                        net.computeGradientAndScore();
                        net2.computeGradientAndScore();
                        assertEquals(net2.score(), net.score(), 1e-5);
                        Map<String, INDArray> gradients = net.gradient().gradientForVariable();
                        Map<String, INDArray> gradients2 = net2.gradient().gradientForVariable();
                        assertEquals(gradients.keySet(), gradients2.keySet());
                        for (String s : gradients.keySet()) {
                            assertEquals(gradients2.get(s), gradients.get(s));
                        }
                    }
                }
            }
        }
    }

    @EqualsAndHashCode
    @DisplayName("Word Vectors Mockup")
    private static class WordVectorsMockup implements EmbeddingInitializer {

        @Override
        public void loadWeightsInto(INDArray array) {
            INDArray vectors = Nd4j.linspace(1, 15, 15, DataType.FLOAT).reshape(5, 3);
            array.assign(vectors);
        }

        @Override
        public long vocabSize() {
            return 5;
        }

        @Override
        public int vectorSize() {
            return 3;
        }

        @Override
        public boolean jsonSerializable() {
            return true;
        }
    }

    @Test
    @DisplayName("Test Embedding Default Activation")
    void testEmbeddingDefaultActivation() {
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().list().layer(new EmbeddingLayer.Builder().nIn(10).nOut(10).build()).layer(new EmbeddingSequenceLayer.Builder().nIn(10).nOut(10).build()).build();
        EmbeddingLayer l = (EmbeddingLayer) conf.getConf(0).getLayer();
        assertEquals(new ActivationIdentity(), l.getActivationFn());
        EmbeddingSequenceLayer l2 = (EmbeddingSequenceLayer) conf.getConf(1).getLayer();
        assertEquals(new ActivationIdentity(), l2.getActivationFn());
    }

    @Test
    @DisplayName("Test Embedding Weight Init")
    void testEmbeddingWeightInit() {
        // https://github.com/eclipse/deeplearning4j/issues/8663
        // The embedding layer weight initialization should be independent of the vocabulary size (nIn setting)
        for (WeightInit wi : new WeightInit[] { WeightInit.XAVIER, WeightInit.RELU, WeightInit.XAVIER_UNIFORM, WeightInit.LECUN_NORMAL }) {
            for (boolean seq : new boolean[] { false, true }) {
                Nd4j.getRandom().setSeed(12345);
                MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(12345).list().layer(seq ? new EmbeddingSequenceLayer.Builder().weightInit(wi).nIn(100).nOut(100).build() : new EmbeddingLayer.Builder().weightInit(wi).nIn(100).nOut(100).build()).build();
                MultiLayerNetwork net = new MultiLayerNetwork(conf);
                net.init();
                Nd4j.getRandom().setSeed(12345);
                MultiLayerConfiguration conf2 = new NeuralNetConfiguration.Builder().seed(12345).list().layer(seq ? new EmbeddingSequenceLayer.Builder().weightInit(wi).nIn(100).nOut(100).build() : new EmbeddingLayer.Builder().weightInit(wi).nIn(100).nOut(100).build()).build();
                MultiLayerNetwork net2 = new MultiLayerNetwork(conf2);
                net2.init();
                Nd4j.getRandom().setSeed(12345);
                MultiLayerConfiguration conf3 = new NeuralNetConfiguration.Builder().seed(12345).list().layer(seq ? new EmbeddingSequenceLayer.Builder().weightInit(wi).nIn(100000).nOut(100).build() : new EmbeddingLayer.Builder().weightInit(wi).nIn(100000).nOut(100).build()).build();
                MultiLayerNetwork net3 = new MultiLayerNetwork(conf3);
                net3.init();
                INDArray p1 = net.params();
                INDArray p2 = net2.params();
                INDArray p3 = net3.params();
                boolean eq = p1.equalsWithEps(p2, 1e-4);
                String str = (seq ? "EmbeddingSequenceLayer" : "EmbeddingLayer") + " - " + wi;
                assertTrue(eq,str + " p1/p2 params not equal");
                double m1 = p1.meanNumber().doubleValue();
                double s1 = p1.stdNumber().doubleValue();
                double m3 = p3.meanNumber().doubleValue();
                double s3 = p3.stdNumber().doubleValue();
                assertEquals( m1, m3, 0.1,str);
                assertEquals(s1, s3, 0.1,str);
                double re = relErr(s1, s3);
                assertTrue( re < 0.05,str + " - " + re);
            }
        }
    }

    public static double relErr(double d1, double d2) {
        if (d1 == 0.0 && d2 == 0.0)
            return 0.0;
        return Math.abs(d1 - d2) / (Math.abs(d1) + Math.abs(d2));
    }
}
