Java Code Examples for org.nd4j.linalg.api.ndarray.INDArray#put()

@Override
public INDArray doForward(boolean training, LayerWorkspaceMgr workspaceMgr) {
    if (!canDoForward()) {
        throw new IllegalStateException("Cannot do forward pass: inputs not set");
    }
    // inputs[index] => {batchSize, numberOfEmbeds}
    INDArray left = inputs[0];
    INDArray right = inputs[1];
    long size = inputs[0].shape()[0];
    INDArray value = workspaceMgr.createUninitialized(ArrayType.ACTIVATIONS, size);
    // 求两个行向量的点积
    for (int index = 0; index < size; index++) {
        INDArray product = left.getRow(index).mmul(right.getRow(index).transpose());
        value.put(index, product);
    }
    // outputs[index] => {batchSize, 1}
    return Shape.newShapeNoCopy(value, new long[] { value.length(), 1L }, value.ordering() == 'f');
}
 

@Override
public void doForward() {
    GlobalMatrix inputKeyMatrix = GlobalMatrix.class.cast(inputKeyValues[0].getKey());
    GlobalMatrix inputValueMatrix = GlobalMatrix.class.cast(inputKeyValues[0].getValue());
    Nd4jMatrix outputKeyMatrix = Nd4jMatrix.class.cast(outputKeyValue.getKey());
    Nd4jMatrix outputValueMatrix = Nd4jMatrix.class.cast(outputKeyValue.getValue());

    {
        INDArray outputData = outputKeyMatrix.getArray();
        int cursor = 0;
        for (MathMatrix component : inputKeyMatrix.getComponentMatrixes()) {
            Nd4jMatrix nd4j = Nd4jMatrix.class.cast(component);
            INDArray array = nd4j.getArray();
            if (orientation) {
                outputData.put(new INDArrayIndex[] { NDArrayIndex.all(), NDArrayIndex.interval(cursor, cursor + array.columns()) }, array);
                cursor += array.columns();
            } else {
                outputData.put(new INDArrayIndex[] { NDArrayIndex.interval(cursor, cursor + array.rows()), NDArrayIndex.all() }, array);
                cursor += array.rows();
            }
        }
    }

    outputValueMatrix.setValues(0F);
}
 

public static void main(String[] args) {
    INDArray qTable = buildQTable();
    for (int i = 0; i < MAX_EPISODES; i++) {
        int stepCounter = 0;
        int S = 0;
        boolean terminated = false;
        updateEnv(S, i, stepCounter);
        while (!terminated) {
            int A = chooseAction(S, qTable);
            int[] result = getEnvBack(S, ACTIONS[A]);
            int S_ = result[0], R = result[1];
            double qPredict = qTable.getDouble(S, A), qTarget;
            if (S_ != Integer.MAX_VALUE) {
                qTarget = R + GAMMA * qTable.getRow(S_).maxNumber().doubleValue();
            } else {
                qTarget = R;
                terminated = true;
            }
            qTable.put(S, A, qPredict + ALPHA * (qTarget - qPredict));
            S = S_;
            stepCounter++;
            updateEnv(S, i, stepCounter);
        }
    }
    System.out.println(qTable);
}
 

private void growCapacity(int idx) {
    if(container == null) {
        container = Nd4j.create(10L);
    }
    else if(idx >= container.length()) {
        long max = Math.max(container.length() * 2L,idx);
        INDArray newContainer = Nd4j.create(max);
        newContainer.put(new INDArrayIndex[]{NDArrayIndex.interval(0,container.length())},container);
        container = newContainer;
    }
}
 

@Test
public void testFlatIndexPointInterval() {
    INDArray zeros = Nd4j.zeros(1, 4);
    INDArrayIndex x = NDArrayIndex.point(0);
    INDArrayIndex y = NDArrayIndex.interval(1, 2, true);
    INDArray value = Nd4j.ones(1, 2);
    zeros.put(new INDArrayIndex[] {x, y}, value);

    INDArray assertion = Nd4j.create(new double[] {0.0, 1.0, 1.0, 0.0});
    assertEquals(assertion, zeros);
}
 

/**
 * Implement column formatted images
 * @param img the image to process
 * @param kh the kernel height
 * @param kw the kernel width
 * @param sy the stride along y
 * @param sx the stride along x
 * @param ph the padding width
 * @param pw the padding height
 * @param pval the padding value
 * @param coverAll whether to cover the whole image or not
 * @return the column formatted image
 *
 */
public static INDArray im2col(INDArray img, int kh, int kw, int sy, int sx, int ph, int pw, int pval,
                boolean coverAll) {
    //number of images
    long n = img.size(0);
    //number of channels (depth)
    long c = img.size(1);
    //image height
    long h = img.size(2);
    //image width
    long w = img.size(3);
    long outHeight = outSize(h, kh, sy, ph, coverAll);
    long outWidth = outSize(w, kw, sx, pw, coverAll);
    INDArray padded = Nd4j.pad(img, new int[][] {{0, 0}, {0, 0}, {ph, ph + sy - 1}, {pw, pw + sx - 1}},
                    Nd4j.PadMode.CONSTANT);
    INDArray ret = Nd4j.create(n, c, kh, kw, outHeight, outWidth);
    for (int i = 0; i < kh; i++) {
        //offset for the row based on the stride and output height
        long iLim = i + sy * outHeight;
        for (int j = 0; j < kw; j++) {
            //offset for the column based on stride and output width
            long jLim = j + sx * outWidth;
            INDArray get = padded.get(NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.interval(i, sy, iLim),
                            NDArrayIndex.interval(j, sx, jLim));
            ret.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(i),
                            NDArrayIndex.point(j), NDArrayIndex.all(), NDArrayIndex.all()}, get);
        }
    }
    return ret;
}
 

@Test
public void testPutRowIndexing() {
    INDArray arr = Nd4j.ones(1, 10);
    INDArray row = Nd4j.create(1, 10);

    arr.putRow(0, row); //OK
    arr.put(new INDArrayIndex[] {point(0), all()}, row); //Exception
    assertEquals(arr, row);
}
 

@Test
public void testScalarPut() {
    List<List<Integer>> indices = new ArrayList<>();
    indices.add(Arrays.asList(0));
    indices.add(Arrays.asList(1));
    indices.add(Arrays.asList(0));
    indices.add(Arrays.asList(0));

    INDArray linspace = Nd4j.linspace(1,16,16).reshape('c',2,2,2,2);
    linspace.put(indices,Nd4j.scalar(99.0));
    assertEquals(99.0,linspace.getDouble(0,1,0,0),1e-1);
}
 

private Pair<INDArray, INDArray> paddingFeatures(List<List<Writable>> list, int minExamples,
		int longestVector, SubsetDetails details, int[] vectorLength) {
	if (list.isEmpty()) {
		throw new ZeroLengthSequenceException("Zero length sequence encountered");
	}

	// We do not take into account the label in this vector, just the features
	INDArray arr = Nd4j.create(new int[] { list.size(), longestVector }, 'f');

	INDArray maskArray = Nd4j.zeros(list.size(), longestVector);

	for (int i = 0; i < list.size(); i++)
	{
		List<Writable> record = list.get(i);
		Writable w = null;
		if (details.subsetStart < 0)
		{
			w = record.get(0);
		}
		else
		{
			w = record.get(details.subsetStart);
		}

		if (w instanceof NDArrayWritable)
		{
			INDArray value = ((NDArrayWritable) w).get();
			arr.put(new INDArrayIndex[] { NDArrayIndex.point(i), NDArrayIndex.interval(0, value.length()) }, value);
		} 
		else
		{
			arr.putScalar((long) 0, (long) 0, w.toDouble());
		}

		// Masking array entries at start (for align end)
		for (int t2 = 0; t2 < vectorLength[i]; t2++) {
			maskArray.putScalar(i, t2, 1.0);
		}
	}

	return new Pair<>(arr, maskArray);
}
 

@Override
public INDArray doForward(boolean training, LayerWorkspaceMgr workspaceMgr) {
    // stacking along dimension 0
    // inputs[] is an array of INDArray (e.g.: shape of 3 x [nExamples, nSize])
    // what we want to do is make a stacked output (e.g.: [3 x nExamples, nSize])
    lastInputShapes = null;
    int nStack = inputs.length;
    val inShape = inputs[0].shape();
    val outShape = new long[inShape.length];

    // create the new shape
    outShape[0] = nStack * inShape[0];
    for (int i = 1; i < inShape.length; i++) {
        outShape[i] = inShape[i];
    }

    boolean variableLengthTS = false;
    if (inShape.length == 3) {
        //RNN data - check for variable length time series
        long minLength = inputs[0].size(2);
        long maxLength = minLength;
        for (int i = 1; i < inputs.length; i++) {
            long thisLength = inputs[i].size(2);
            minLength = Math.min(minLength, thisLength);
            maxLength = Math.max(maxLength, thisLength);
        }
        variableLengthTS = (minLength != maxLength);

        if (!variableLengthTS) {
            try(MemoryWorkspace ws = workspaceMgr.notifyScopeBorrowed(ArrayType.ACTIVATIONS)) {
                return Nd4j.concat(0, inputs);
            }
        }

        outShape[2] = maxLength;
        INDArray out = workspaceMgr.create(ArrayType.ACTIVATIONS, inputs[0].dataType(), outShape);
        long numExamples = inputs[0].size(0);
        lastInputShapes = new long[inputs.length][0];
        for (int i = 0; i < inputs.length; i++) {
            out.put(new INDArrayIndex[] {NDArrayIndex.interval(i * numExamples, (i + 1) * numExamples),
                            NDArrayIndex.all(), NDArrayIndex.interval(0, inputs[i].size(2))}, inputs[i]);
            lastInputShapes[i] = inputs[i].shape();
        }

        return out;
    } else {
        try(MemoryWorkspace ws = workspaceMgr.notifyScopeBorrowed(ArrayType.ACTIVATIONS)) {
            return Nd4j.concat(0, inputs);
        }
    }
}
 

/**
 * Merge this ROC instance with another.
 * This ROC instance is modified, by adding the stats from the other instance.
 *
 * @param other ROC instance to combine with this one
 */
@Override
public void merge(ROC other) {
    if (this.thresholdSteps != other.thresholdSteps) {
        throw new UnsupportedOperationException(
                "Cannot merge ROC instances with different numbers of threshold steps ("
                        + this.thresholdSteps + " vs. " + other.thresholdSteps + ")");
    }
    this.countActualPositive += other.countActualPositive;
    this.countActualNegative += other.countActualNegative;
    this.auc = null;
    this.auprc = null;
    this.rocCurve = null;
    this.prCurve = null;


    if (isExact) {
        if (other.exampleCount == 0) {
            return;
        }

        if (this.exampleCount == 0) {
            this.exampleCount = other.exampleCount;
            this.probAndLabel = other.probAndLabel;
            return;
        }

        if (this.exampleCount + other.exampleCount > this.probAndLabel.size(0)) {
            //Allocate new array
            val newSize = this.probAndLabel.size(0) + Math.max(other.probAndLabel.size(0), exactAllocBlockSize);
            INDArray newProbAndLabel = Nd4j.create(DataType.DOUBLE, newSize, 2);
            newProbAndLabel.put(new INDArrayIndex[]{interval(0, exampleCount), all()}, probAndLabel.get(interval(0, exampleCount), all()));
            probAndLabel = newProbAndLabel;
        }

        INDArray toPut = other.probAndLabel.get(interval(0, other.exampleCount), all());
        probAndLabel.put(new INDArrayIndex[]{
                        interval(exampleCount, exampleCount + other.exampleCount), all()},
                toPut);
    } else {
        for (Double d : this.counts.keySet()) {
            CountsForThreshold cft = this.counts.get(d);
            CountsForThreshold otherCft = other.counts.get(d);
            cft.countTruePositive += otherCft.countTruePositive;
            cft.countFalsePositive += otherCft.countFalsePositive;
        }
    }

    this.exampleCount += other.exampleCount;
}
 

@Override
public DataSet next(int num) {
  // Check if next() call is valid - throws appropriate exceptions
  checkIfNextIsValid();

  // Collect
  List<String> sentences = new ArrayList<>(num);
  List<Double> labelsRaw = new ArrayList<>(num);
  collectData(num, sentences, labelsRaw);
  final int numDocuments = sentences.size();

  // Tokenize sentences
  List<List<String>> tokenizedSentences = tokenizeSentences(sentences);

  // Get longest sentence length
  int maxSentenceLength = tokenizedSentences.stream().mapToInt(List::size).max().getAsInt();

  // Truncate maximum sentence length
  if (maxSentenceLength > truncateLength || maxSentenceLength == 0) {
    maxSentenceLength = truncateLength;
  }

  // Init feature/label arrays
  int[] featureShape = {numDocuments, wordVectorSize, maxSentenceLength};
  int[] labelShape = {numDocuments, data.numClasses(), maxSentenceLength};
  INDArray features = Nd4j.create(featureShape, 'f');
  INDArray labels = Nd4j.create(labelShape, 'f');
  INDArray featuresMask = Nd4j.zeros(numDocuments, maxSentenceLength);
  INDArray labelsMask = Nd4j.zeros(numDocuments, maxSentenceLength);

  for (int i = 0; i < numDocuments; i++) {
    List<String> tokens = tokenizedSentences.get(i);

    // Check for empty document
    if (tokens.isEmpty()) {
      continue;
    }

    // Get the last index of the current document (truncated)
    int lastIdx = Math.min(tokens.size(), maxSentenceLength);

    // Get all wordvectors in batch
    List<String> truncatedTokenList = tokens.subList(0, lastIdx);
    final INDArray vectors = wordVectors.getWordVectors(truncatedTokenList).transpose();

    /*
     * Put wordvectors into features array at the following indices:
     * 1) Document (i)
     * 2) All vector elements which is equal to NDArrayIndex.interval(0, vectorSize)
     * 3) All elements between 0 and the length of the current sequence
     */
    INDArrayIndex[] indices = {point(i), all(), interval(0, lastIdx)};
    features.put(indices, vectors);

    // Assign "1" to each position where a feature is present, that is, in the interval of
    // [0, lastIdx)
    featuresMask.get(point(i), interval(0, lastIdx)).assign(1);

    // Put the labels in the labels and labelsMask arrays
    // Differ between classification and regression task
    if (data.numClasses() == 1) { // Regression
      double val = labelsRaw.get(i);
      labels.putScalar(new int[]{i, 0, lastIdx - 1}, val);
    } else if (data.numClasses() > 1) { // Classification
      // One-Hot-Encoded class
      int idx = labelsRaw.get(i).intValue();
      // Set label
      labels.putScalar(new int[]{i, idx, lastIdx - 1}, 1.0);
    } else {
      throw new RuntimeException("Could not detect classification or regression task.");
    }

    // Set final timestep for this example to 1.0 to show that an output exists here
    int[] lastTimestepIndex = {i, lastIdx - 1};
    labelsMask.putScalar(lastTimestepIndex, 1.0);
  }

  // Cache the dataset
  final DataSet ds = new DataSet(features, labels, featuresMask, labelsMask);

  // Move cursor
  cursor += ds.numExamples();
  return ds;
}
 

/**
 * This test fails, because it highlights current mechanics on SpecifiedIndex stuff.
 * Internally there's
 *
 * @throws Exception
 */
@Test
public void testSliceAssign1() {
    INDArray array = Nd4j.zeros(4, 4);

    INDArray patch = Nd4j.create(new float[] {1e-5f, 1e-5f, 1e-5f});

    INDArray slice = array.slice(1);
    int[] idx = new int[] {0, 1, 3};
    INDArrayIndex[] range = new INDArrayIndex[] {new SpecifiedIndex(idx)};

    INDArray subarray = slice.get(range);

    //System.out.println("Subarray: " + Arrays.toString(subarray.data().asFloat()) + " isView: " + subarray.isView());

    slice.put(range, patch);

    //System.out.println("Array after being patched: " + Arrays.toString(array.data().asFloat()));

    assertFalse(BooleanIndexing.and(array, Conditions.equals(0f)));
}
 

@Test
    public void testConcat3d() {
        INDArray first = Nd4j.linspace(1, 24, 24, Nd4j.dataType()).reshape('c', 2, 3, 4);
        INDArray second = Nd4j.linspace(24, 36, 12, Nd4j.dataType()).reshape('c', 1, 3, 4);
        INDArray third = Nd4j.linspace(36, 48, 12, Nd4j.dataType()).reshape('c', 1, 3, 4);

        //ConcatV2, dim 0
        INDArray exp = Nd4j.create(2 + 1 + 1, 3, 4);
        exp.put(new INDArrayIndex[] {NDArrayIndex.interval(0, 2), NDArrayIndex.all(), NDArrayIndex.all()}, first);
        exp.put(new INDArrayIndex[] {NDArrayIndex.point(2), NDArrayIndex.all(), NDArrayIndex.all()}, second);
        exp.put(new INDArrayIndex[] {NDArrayIndex.point(3), NDArrayIndex.all(), NDArrayIndex.all()}, third);

        INDArray concat0 = Nd4j.concat(0, first, second, third);

        assertEquals(exp, concat0);

        //ConcatV2, dim 1
        second = Nd4j.linspace(24, 32, 8, Nd4j.dataType()).reshape('c', 2, 1, 4);
        for (int i = 0; i < second.tensorsAlongDimension(1); i++) {
            INDArray secondTad = second.tensorAlongDimension(i, 1);
//            System.out.println(second.tensorAlongDimension(i, 1));
        }

        third = Nd4j.linspace(32, 48, 16).reshape('c', 2, 2, 4);
        exp = Nd4j.create(2, 3 + 1 + 2, 4);
        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(0, 3), NDArrayIndex.all()}, first);
        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.point(3), NDArrayIndex.all()}, second);
        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(4, 6), NDArrayIndex.all()}, third);

        INDArray concat1 = Nd4j.concat(1, first, second, third);

        assertEquals(exp, concat1);

        //ConcatV2, dim 2
        second = Nd4j.linspace(24, 36, 12).reshape('c', 2, 3, 2);
        third = Nd4j.linspace(36, 42, 6).reshape('c', 2, 3, 1);
        exp = Nd4j.create(2, 3, 4 + 2 + 1);

        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.interval(0, 4)}, first);
        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.interval(4, 6)}, second);
        exp.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(6)}, third);

        INDArray concat2 = Nd4j.concat(2, first, second, third);

        assertEquals(exp, concat2);
    }
 

@Test
    public void testGet() {
//        System.out.println("Testing sub-array put and get with a 3D array ...");

        INDArray arr = Nd4j.linspace(0, 124, 125).reshape(5, 5, 5);

        /*
         * Extract elements with the following indices:
         *
         * (2,1,1) (2,1,2) (2,1,3)
         * (2,2,1) (2,2,2) (2,2,3)
         * (2,3,1) (2,3,2) (2,3,3)
         */

        int slice = 2;

        int iStart = 1;
        int jStart = 1;

        int iEnd = 4;
        int jEnd = 4;

        // Method A: Element-wise.

        INDArray subArr_A = Nd4j.create(new int[] {3, 3});

        for (int i = iStart; i < iEnd; i++) {
            for (int j = jStart; j < jEnd; j++) {

                double val = arr.getDouble(slice, i, j);
                int[] sub = new int[] {i - iStart, j - jStart};

                subArr_A.putScalar(sub, val);
            }
        }

        // Method B: Using NDArray get and put with index classes.

        INDArray subArr_B = Nd4j.create(new int[] {3, 3});

        INDArrayIndex ndi_Slice = NDArrayIndex.point(slice);
        INDArrayIndex ndi_J = NDArrayIndex.interval(jStart, jEnd);
        INDArrayIndex ndi_I = NDArrayIndex.interval(iStart, iEnd);

        INDArrayIndex[] whereToGet = new INDArrayIndex[] {ndi_Slice, ndi_I, ndi_J};

        INDArray whatToPut = arr.get(whereToGet);
//        System.out.println(whatToPut);
        INDArrayIndex[] whereToPut = new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all()};

        subArr_B.put(whereToPut, whatToPut);

        assertEquals(subArr_A, subArr_B);

//        System.out.println("... done");
    }
 

@Test
    public void testEvalMasking() {
        int miniBatch = 5;
        int nOut = 3;
        int tsLength = 6;

        INDArray labels = Nd4j.zeros(miniBatch, nOut, tsLength);
        INDArray predicted = Nd4j.zeros(miniBatch, nOut, tsLength);

        Nd4j.getRandom().setSeed(12345);
        Random r = new Random(12345);
        for (int i = 0; i < miniBatch; i++) {
            for (int j = 0; j < tsLength; j++) {
                INDArray rand = Nd4j.rand(1, nOut);
                rand.divi(rand.sumNumber());
                predicted.put(new INDArrayIndex[] {NDArrayIndex.point(i), all(), NDArrayIndex.point(j)},
                                rand);
                int idx = r.nextInt(nOut);
                labels.putScalar(new int[] {i, idx, j}, 1.0);
            }
        }

        //Create a longer labels/predicted with mask for first and last time step
        //Expect masked evaluation to be identical to original evaluation
        INDArray labels2 = Nd4j.zeros(miniBatch, nOut, tsLength + 2);
        labels2.put(new INDArrayIndex[] {all(), all(),
                        interval(1, tsLength + 1)}, labels);
        INDArray predicted2 = Nd4j.zeros(miniBatch, nOut, tsLength + 2);
        predicted2.put(new INDArrayIndex[] {all(), all(),
                        interval(1, tsLength + 1)}, predicted);

        INDArray labelsMask = Nd4j.ones(miniBatch, tsLength + 2);
        for (int i = 0; i < miniBatch; i++) {
            labelsMask.putScalar(new int[] {i, 0}, 0.0);
            labelsMask.putScalar(new int[] {i, tsLength + 1}, 0.0);
        }

        Evaluation evaluation = new Evaluation();
        evaluation.evalTimeSeries(labels, predicted);

        Evaluation evaluation2 = new Evaluation();
        evaluation2.evalTimeSeries(labels2, predicted2, labelsMask);

//        System.out.println(evaluation.stats());
//        System.out.println(evaluation2.stats());
        evaluation.stats();
        evaluation2.stats();

        assertEquals(evaluation.accuracy(), evaluation2.accuracy(), 1e-12);
        assertEquals(evaluation.f1(), evaluation2.f1(), 1e-12);

        assertMapEquals(evaluation.falsePositives(), evaluation2.falsePositives());
        assertMapEquals(evaluation.falseNegatives(), evaluation2.falseNegatives());
        assertMapEquals(evaluation.truePositives(), evaluation2.truePositives());
        assertMapEquals(evaluation.trueNegatives(), evaluation2.trueNegatives());

        for (int i = 0; i < nOut; i++)
            assertEquals(evaluation.classCount(i), evaluation2.classCount(i));
    }
 

@Test
public void testMergingWithPerOutputMasking() {

    //Test 2d mask merging, 2d data
    //features
    INDArray f2d1 = Nd4j.create(new double[] {1, 2, 3});
    INDArray f2d2 = Nd4j.create(new double[][] {{4, 5, 6}, {7, 8, 9}});
    //labels
    INDArray l2d1 = Nd4j.create(new double[] {1.5, 2.5, 3.5});
    INDArray l2d2 = Nd4j.create(new double[][] {{4.5, 5.5, 6.5}, {7.5, 8.5, 9.5}});
    //feature masks
    INDArray fm2d1 = Nd4j.create(new double[] {0, 1, 1});
    INDArray fm2d2 = Nd4j.create(new double[][] {{1, 0, 1}, {0, 1, 0}});
    //label masks
    INDArray lm2d1 = Nd4j.create(new double[] {1, 1, 0});
    INDArray lm2d2 = Nd4j.create(new double[][] {{1, 0, 0}, {0, 1, 1}});

    DataSet mds2d1 = new DataSet(f2d1, l2d1, fm2d1, lm2d1);
    DataSet mds2d2 = new DataSet(f2d2, l2d2, fm2d2, lm2d2);
    DataSet merged = DataSet.merge(Arrays.asList(mds2d1, mds2d2));

    INDArray expFeatures2d = Nd4j.create(new double[][] {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}});
    INDArray expLabels2d = Nd4j.create(new double[][] {{1.5, 2.5, 3.5}, {4.5, 5.5, 6.5}, {7.5, 8.5, 9.5}});
    INDArray expFM2d = Nd4j.create(new double[][] {{0, 1, 1}, {1, 0, 1}, {0, 1, 0}});
    INDArray expLM2d = Nd4j.create(new double[][] {{1, 1, 0}, {1, 0, 0}, {0, 1, 1}});

    DataSet dsExp2d = new DataSet(expFeatures2d, expLabels2d, expFM2d, expLM2d);
    assertEquals(dsExp2d, merged);

    //Test 4d features, 2d labels, 2d masks
    INDArray f4d1 = Nd4j.create(1, 3, 5, 5);
    INDArray f4d2 = Nd4j.create(2, 3, 5, 5);
    DataSet ds4d1 = new DataSet(f4d1, l2d1, null, lm2d1);
    DataSet ds4d2 = new DataSet(f4d2, l2d2, null, lm2d2);
    DataSet merged4d = DataSet.merge(Arrays.asList(ds4d1, ds4d2));
    assertEquals(expLabels2d, merged4d.getLabels());
    assertEquals(expLM2d, merged4d.getLabelsMaskArray());

    //Test 3d mask merging, 3d data
    INDArray f3d1 = Nd4j.create(1, 3, 4);
    INDArray f3d2 = Nd4j.create(1, 3, 3);
    INDArray l3d1 = Nd4j.getExecutioner().exec(new BernoulliDistribution(Nd4j.create(1, 3, 4), 0.5));
    INDArray l3d2 = Nd4j.getExecutioner().exec(new BernoulliDistribution(Nd4j.create(2, 3, 3), 0.5));
    INDArray lm3d1 = Nd4j.getExecutioner().exec(new BernoulliDistribution(Nd4j.create(1, 3, 4), 0.5));
    INDArray lm3d2 = Nd4j.getExecutioner().exec(new BernoulliDistribution(Nd4j.create(2, 3, 3), 0.5));
    DataSet ds3d1 = new DataSet(f3d1, l3d1, null, lm3d1);
    DataSet ds3d2 = new DataSet(f3d2, l3d2, null, lm3d2);

    INDArray expLabels3d = Nd4j.create(3, 3, 4);
    expLabels3d.put(new INDArrayIndex[] {NDArrayIndex.point(0), NDArrayIndex.all(), NDArrayIndex.interval(0, 4)},
                    l3d1);
    expLabels3d.put(new INDArrayIndex[] {NDArrayIndex.interval(1, 2, true), NDArrayIndex.all(),
                    NDArrayIndex.interval(0, 3)}, l3d2);
    INDArray expLM3d = Nd4j.create(3, 3, 4);
    expLM3d.put(new INDArrayIndex[] {NDArrayIndex.point(0), NDArrayIndex.all(), NDArrayIndex.interval(0, 4)},
                    lm3d1);
    expLM3d.put(new INDArrayIndex[] {NDArrayIndex.interval(1, 2, true), NDArrayIndex.all(),
                    NDArrayIndex.interval(0, 3)}, lm3d2);


    DataSet merged3d = DataSet.merge(Arrays.asList(ds3d1, ds3d2));
    assertEquals(expLabels3d, merged3d.getLabels());
    assertEquals(expLM3d, merged3d.getLabelsMaskArray());

    //Test 3d features, 2d masks, 2d output (for example: RNN -> global pooling w/ per-output masking)
    DataSet ds3d2d1 = new DataSet(f3d1, l2d1, null, lm2d1);
    DataSet ds3d2d2 = new DataSet(f3d2, l2d2, null, lm2d2);
    DataSet merged3d2d = DataSet.merge(Arrays.asList(ds3d2d1, ds3d2d2));

    assertEquals(expLabels2d, merged3d2d.getLabels());
    assertEquals(expLM2d, merged3d2d.getLabelsMaskArray());
}
 

private INDArray computeTempImage(INDArray image, int numBoxes, MtcnnUtil.PadResult padResult, int size) throws IOException {

		//  tempimg = np.zeros(shape=(size, size, 3, num_boxes))
		INDArray tempImg = Nd4j.zeros(new int[] { size, size, CHANNEL_COUNT, numBoxes }, C_ORDERING);

		opencv_core.Size newSize = new opencv_core.Size(size, size);

		for (int k = 0; k < numBoxes; k++) {
			//tmp = np.zeros((int(stage_status.tmph[k]), int(stage_status.tmpw[k]), 3))
			INDArray tmp = Nd4j.zeros(new int[] { padResult.getTmph().getInt(k), padResult.getTmpw().getInt(k), CHANNEL_COUNT }, C_ORDERING);

			// tmp[stage_status.dy[k] - 1:stage_status.edy[k], stage_status.dx[k] - 1:stage_status.edx[k], :] = \
			//   img[stage_status.y[k] - 1:stage_status.ey[k], stage_status.x[k] - 1:stage_status.ex[k], :]
			tmp.put(new INDArrayIndex[] {
							interval(padResult.getDy().getInt(k) - 1, padResult.getEdy().getInt(k)),
							interval(padResult.getDx().getInt(k) - 1, padResult.getEdx().getInt(k)),
							all() },
					image.get(
							interval(padResult.getY().getInt(k) - 1, padResult.getEy().getInt(k)),
							interval(padResult.getX().getInt(k) - 1, padResult.getEx().getInt(k)),
							all()));

			// if tmp.shape[0] > 0 and tmp.shape[1] > 0 or tmp.shape[0] == 0 and tmp.shape[1] == 0:
			//    tempimg[:, :, :, k] = cv2.resize(tmp, (size, size), interpolation=cv2.INTER_AREA)
			if ((tmp.shape()[0] > 0 && tmp.shape()[1] > 0) || (tmp.shape()[0] == 0 && tmp.shape()[1] == 0)) {

				INDArray resizedImage = resize(tmp.permutei(2, 0, 1).dup(), newSize)
						.get(point(0), all(), all(), all()).permutei(1, 2, 0).dup();

				tempImg.put(new INDArrayIndex[] { all(), all(), all(), point(k) }, resizedImage);
			}
			else {
				return Nd4j.empty();
			}
		}

		// tempimg = (tempimg - 127.5) * 0.0078125
		tempImg = tempImg.subi(127.5).muli(0.0078125);

		// tempimg1 = np.transpose(tempimg, (3, 1, 0, 2))
		INDArray tempImg1 = tempImg.permutei(3, 1, 0, 2).dup();

		return tempImg1;
	}
 

private DataSet vectorize(List<? extends Sample> aData, Object2IntMap<String> aTagset,
        boolean aIncludeLabels)
    throws IOException
{
    // vectorize is pretty fast taking around 1-2ms
    
    // long start = System.currentTimeMillis();
    int maxSentenceLength = traits.getMaxSentenceLength();
    
    // Create data for training
    int embeddingSize = wordVectors.dimensions(); 
    INDArray featureVec = Nd4j.create(aData.size(), embeddingSize, maxSentenceLength);

    // Tags are using a 1-hot encoding
    INDArray labelVec = Nd4j.create(aData.size(), traits.getMaxTagsetSize(), maxSentenceLength);
    
    // Sentences have variable length, so we we need to mask positions not used in short
    // sentences.
    INDArray featureMask = Nd4j.zeros(aData.size(), maxSentenceLength);
    INDArray labelMask = Nd4j.zeros(aData.size(), maxSentenceLength);
    
    // Get word vectors for each word in review, and put them in the training data
    int sampleIdx = 0;
    for (Sample sample : aData) {
        List<String> tokens = sample.getSentence();
        List<String> labels = sample.getTags();
        for (int t = 0; t < Math.min(tokens.size(), maxSentenceLength); t++) {
            String word = tokens.get(t);
            INDArray vector = Nd4j.create(wordVectors.vectorize(word));

            if (vector == null) {
                vector = randUnk;
            }

            featureVec.put(new INDArrayIndex[] { point(sampleIdx), all(), point(t) }, vector);
            featureMask.putScalar(new int[] { sampleIdx, t }, 1.0);

            // exclude padding labels from training
            // compare instances to avoid collision with possible no_label user label
            if (labels != null && labels.get(t) != NO_LABEL) {
                labelMask.putScalar(new int[] { sampleIdx, t }, 1.0);
            }

            if (aIncludeLabels && labels != null) {
                String label = labels.get(t);
                // do not add padding label no_label as predictable label
                if (label != NO_LABEL) {
                    aTagset.computeIfAbsent(label, key -> aTagset.size());
                    labelVec.putScalar(sampleIdx, aTagset.get(label), t, 1.0);
                }
            }
        }
        
        sampleIdx++;
    }

    // log.trace("Vectorizing took {}ms", System.currentTimeMillis() - start);
    
    return new DataSet(featureVec, labelVec, featureMask, labelMask);
}
 

/**
 * This method returns the gradient of the cost function with respect to the
 * output from the previous layer.  For this cost function, the gradient
 * is derived from Bishop's paper "Mixture Density Networks" (1994) which
 * gives an elegant closed-form expression for the derivatives with respect
 * to each of the output components.
 * @param labels Labels to train on.
 * @param preOutput Output of neural network before applying the final activation function.
 * @param activationFn Activation function of output layer.
 * @param mask Mask to apply to gradients.
 * @return Gradient of cost function with respect to preOutput parameters.
 */
@Override
public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    long nSamples = labels.size(0);

    INDArray output = activationFn.getActivation(preOutput.dup(), false);

    MixtureDensityComponents mdc = extractComponents(output);

    INDArray gradient = Nd4j.zeros(nSamples, preOutput.columns());

    INDArray labelsMinusMu = labelsMinusMu(labels, mdc.mu);
    INDArray labelsMinusMuSquared = labelsMinusMu.mul(labelsMinusMu).sum(2);

    // This computes pi_i, see Bishop equation (30).
    // See http://www.plsyard.com/dealing-overflow-and-underflow-in-softmax-function/
    // this post for why we calculate the pi_i in this way.
    // With the exponential function here, we have to be very careful
    // about overflow/underflow considerations even with
    // fairly intermediate values.  Subtracting the max
    // here helps to ensure over/underflow does not happen here.
    // This isn't exactly a softmax because there's an 'alpha' coefficient
    // here, but the technique works, nonetheless.
    INDArray variance = mdc.sigma.mul(mdc.sigma);
    INDArray minustwovariance = variance.mul(2).negi();
    INDArray normalPart = mdc.alpha.div(Transforms.pow(mdc.sigma.mul(SQRT_TWO_PI), mLabelWidth));
    INDArray exponent = labelsMinusMuSquared.div(minustwovariance);
    INDArray exponentMax = exponent.max(1);
    exponent.subiColumnVector(exponentMax);
    INDArray pi = Transforms.exp(exponent).muli(normalPart);
    INDArray piDivisor = pi.sum(1);
    pi.diviColumnVector(piDivisor);

    // See Bishop equation (35)
    //INDArray dLdZAlpha = Nd4j.zeros(nSamples, nLabelsPerSample, mMixturesPerLabel); //mdc.alpha.sub(pi);
    INDArray dLdZAlpha = mdc.alpha.sub(pi);
    // See Bishop equation (38)
    INDArray dLdZSigma = (labelsMinusMuSquared.div(variance).subi(mLabelWidth)).muli(-1).muli(pi);
    // See Bishop equation (39)

    // This turned out to be way less efficient than
    // the simple 'for' loop here.
    //INDArray dLdZMu = pi
    //        .div(variance)
    //        .reshape(nSamples, mMixtures, 1)
    //        .repeat(2, mLabelWidth)
    //        .muli(labelsMinusMu)
    //        .negi()
    //        .reshape(nSamples, mMixtures * mLabelWidth);

    INDArray dLdZMu = Nd4j.create(nSamples, mMixtures, mLabelWidth);
    for (int k = 0; k < mLabelWidth; k++) {
        dLdZMu.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(), NDArrayIndex.point(k)},
                        labelsMinusMu.get(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.all(),
                                        NDArrayIndex.point(k)}).muli(pi).divi(variance).negi());
    }
    dLdZMu = dLdZMu.reshape(nSamples, mMixtures * mLabelWidth);

    // Place components of gradient into gradient holder.
    gradient.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(0, mMixtures)}, dLdZAlpha);
    gradient.put(new INDArrayIndex[] {NDArrayIndex.all(), NDArrayIndex.interval(mMixtures, mMixtures * 2)},
                    dLdZSigma);
    gradient.put(new INDArrayIndex[] {NDArrayIndex.all(),
                    NDArrayIndex.interval(mMixtures * 2, (mLabelWidth + 2) * mMixtures)}, dLdZMu);

    INDArray gradients = activationFn.backprop(preOutput, gradient).getFirst();

    if (mask != null) {
        LossUtil.applyMask(gradients, mask);
    }

    return gradients;
}