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

public INDArray scoreArray(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype
    /* y_hat is -1 or 1
    hinge loss is max(0,1-y_hat*y)
     */
    INDArray output = activationFn.getActivation(preOutput.dup(), true);

    INDArray scoreArr = output.muli(labels); //y*yhat
    scoreArr.rsubi(1.0); //1 - y*yhat

    if (mask != null) {
        LossUtil.applyMask(scoreArr, mask);
    }
    return scoreArr; // 1 - y*yhat
}
 

@Test
public void testANNSearchReflexive() {
    rpLSH.makeIndex(inputs);
    int idx = (new Random(12345)).nextInt(100);
    INDArray row = inputs.getRow(idx).reshape(1, intDimensions);

    INDArray searchResults = rpLSH.search(row, 100);


    INDArray res = Nd4j.zeros(DataType.BOOL, searchResults.shape());
    Nd4j.getExecutioner().exec(new BroadcastEqualTo(searchResults, row, res, -1));
    res = res.castTo(DataType.FLOAT);

    assertEquals(
            String.format("Expected one search result to be the query %s, but found %s", row, searchResults),
            1.0f, res.min(-1).maxNumber().floatValue(), 1e-3f);
}
 

@Override
public INDArray getVarCache(DataType dataType) {
    INDArray ret;
    if(dataType == DataType.HALF){
        INDArray vc = varCache.castTo(DataType.HALF);
        ret = vc.mul(vc).rdivi(1.0).subi(eps);
    } else {
        ret = varCache.mul(varCache).rdivi(1.0).subi(eps);
    }
    if(dataType == DataType.HALF){
        //Buffer is FP32
        return ret.castTo(DataType.HALF);
    }
    return ret;
}
 

protected INDArray scoreArray(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype

    INDArray output = activationFn.getActivation(preOutput.dup(), true);
    if(activationFn instanceof ActivationSoftmax && softmaxClipEps > 0.0){
        BooleanIndexing.replaceWhere(output, softmaxClipEps, Conditions.lessThan(softmaxClipEps));
        BooleanIndexing.replaceWhere(output, 1.0-softmaxClipEps, Conditions.greaterThan(1.0-softmaxClipEps));
    }
    INDArray scoreArr = Transforms.log(output, false).muli(labels);

    //Weighted loss function
    if (weights != null) {
        if (weights.length() != scoreArr.size(1)) {
            throw new IllegalStateException("Weights vector (length " + weights.length()
                            + ") does not match output.size(1)=" + preOutput.size(1));
        }
        scoreArr.muliRowVector(weights.castTo(scoreArr.dataType()));
    }

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

@Test
public void testMultipleNumpyJobsParallel(){
    PythonContextManager.deleteNonMainContexts();
    String code1 =(dataType == DataType.BOOL)?"z = x":"z = x + y";
    PythonJob job1 = new PythonJob("job1", code1, false);

    String code2 =(dataType == DataType.BOOL)?"z = y":"z = x - y";
    PythonJob job2 = new PythonJob("job2", code2, false);

    List<PythonVariable> inputs = new ArrayList<>();
    INDArray x = Nd4j.ones(dataType, 2, 3);
    INDArray y = Nd4j.zeros(dataType, 2, 3);
    INDArray z1 = (dataType == DataType.BOOL)?x:x.add(y);
    z1 = (dataType == DataType.BFLOAT16)? z1.castTo(DataType.FLOAT): z1;
    INDArray z2 = (dataType == DataType.BOOL)?y:x.sub(y);
    z2 = (dataType == DataType.BFLOAT16)? z2.castTo(DataType.FLOAT): z2;
    PythonType<INDArray> arrType = PythonTypes.get("numpy.ndarray");
    inputs.add(new PythonVariable<>("x", arrType, x));
    inputs.add(new PythonVariable<>("y", arrType, y));


    List<PythonVariable> outputs = new ArrayList<>();

    outputs.add(new PythonVariable<>("z", arrType));

    job1.exec(inputs, outputs);

    assertEquals(z1, outputs.get(0).getValue());


    job2.exec(inputs, outputs);

    assertEquals(z2, outputs.get(0).getValue());

}
 

@Test
public void testConversion(){
    INDArray arr = Nd4j.zeros(dataType, shape);
    PythonObject npArr = PythonTypes.convert(arr);
    INDArray arr2 = PythonTypes.<INDArray>getPythonTypeForPythonObject(npArr).toJava(npArr);
    if (dataType == DataType.BFLOAT16){
        arr = arr.castTo(DataType.FLOAT);
    }
    Assert.assertEquals(arr,arr2);
}
 

@Test
    public void testNDArrayTextSerializer() throws Exception {
        for(char order : new char[]{'c', 'f'}) {
            Nd4j.factory().setOrder(order);
            for (DataType globalDT : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
                Nd4j.setDefaultDataTypes(globalDT, globalDT);

                Nd4j.getRandom().setSeed(12345);
                INDArray in = Nd4j.rand(DataType.DOUBLE, 3, 4).muli(20).subi(10);

                val om = new ObjectMapper();

                for (DataType dt : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF, DataType.LONG, DataType.INT, DataType.SHORT,
                        DataType.BYTE, DataType.UBYTE, DataType.BOOL, DataType.UTF8}) {

                    INDArray arr;
                    if(dt == DataType.UTF8){
                        arr = Nd4j.create("aaaaa", "bbbb", "ccc", "dd", "e", "f", "g", "h", "i", "j", "k", "l").reshape('c', 3, 4);
                    } else {
                        arr = in.castTo(dt);
                    }

                    TestClass tc = new TestClass(arr);

                    String s = om.writeValueAsString(tc);
//                    System.out.println(dt);
//                    System.out.println(s);
//                    System.out.println("\n\n\n");

                    TestClass deserialized = om.readValue(s, TestClass.class);
                    assertEquals(dt.toString(), tc, deserialized);
                }
            }
        }
    }
 

@Override
public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype

    /*
    gradient is 0 if yhaty is >= 1
    else gradient is gradient of the loss function = (1-yhaty) wrt preOutput = -y*derivative_of_yhat wrt preout
    */

    INDArray bitMaskRowCol = scoreArray(labels, preOutput, activationFn, mask);
    /*
        bit mask is 0 if 1-sigma(y*yhat) is neg
        bit mask is 1 if 1-sigma(y*yhat) is +ve
     */
    BooleanIndexing.replaceWhere(bitMaskRowCol, 0.0, Conditions.lessThan(0.0));
    BooleanIndexing.replaceWhere(bitMaskRowCol, 1.0, Conditions.greaterThan(0.0));

    INDArray dLda = labels.neg().muli(bitMaskRowCol);

    if (mask != null && LossUtil.isPerOutputMasking(dLda, mask)) {
        //For *most* activation functions: we don't actually need to mask dL/da in addition to masking dL/dz later
        //but: some, like softmax, require both (due to dL/dz_i being a function of dL/da_j, for i != j)
        //We could add a special case for softmax (activationFn instanceof ActivationSoftmax) but that would be
        // error prone - though buy us a tiny bit of performance
        LossUtil.applyMask(dLda, mask);
    }

    INDArray gradients = activationFn.backprop(preOutput, dLda).getFirst(); //TODO activation functions with parameters

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

    return gradients;
}
 

@Override
public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype
    INDArray output = activationFn.getActivation(preOutput.dup(), true);

    INDArray outSubLabels = output.sub(labels);
    INDArray dLda = Nd4j.getExecutioner().exec(new Sign(outSubLabels));

    if (weights != null) {
        dLda.muliRowVector(weights.castTo(dLda.dataType()));
    }

    if (mask != null && LossUtil.isPerOutputMasking(dLda, mask)) {
        //For *most* activation functions: we don't actually need to mask dL/da in addition to masking dL/dz later
        //but: some, like softmax, require both (due to dL/dz_i being a function of dL/da_j, for i != j)
        //We could add a special case for softmax (activationFn instanceof ActivationSoftmax) but that would be
        // error prone - but buy us a tiny bit of performance
        LossUtil.applyMask(dLda, mask);
    }

    //dL/dz
    INDArray gradients = activationFn.backprop(preOutput, dLda).getFirst(); //TODO activation function param gradients

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

    return gradients;
}
 

@Override
public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype
    INDArray output = activationFn.getActivation(preOutput.dup(), true);

    INDArray dLda = output.subi(labels).muli(2);

    if (weights != null) {
        dLda.muliRowVector(weights.castTo(dLda.dataType()));
    }

    if (mask != null && LossUtil.isPerOutputMasking(dLda, mask)) {
        //For *most* activation functions: we don't actually need to mask dL/da in addition to masking dL/dz later
        //but: some, like softmax, require both (due to dL/dz_i being a function of dL/da_j, for i != j)
        //We could add a special case for softmax (activationFn instanceof ActivationSoftmax) but that would be
        // error prone - but buy us a tiny bit of performance
        LossUtil.applyMask(dLda, mask);
    }

    INDArray gradients = activationFn.backprop(preOutput, dLda).getFirst(); //TODO handle activation function parameter gradients

    //Loss function with masking
    if (mask != null) {
        LossUtil.applyMask(gradients, mask);
    }

    return gradients;
}
 

private INDArray transform(INDArray raw) {
    long[] shape = raw.shape();

    // before accessing the raw pointer, we need to make sure that array is actual on the host side
    Nd4j.getAffinityManager().ensureLocation(raw, AffinityManager.Location.HOST);

    Mat ocvmat = new Mat((int)shape[0], (int)shape[1], CV_32FC(3), raw.data().pointer());
    Mat cvmat = new Mat(shape[0], shape[1], CV_8UC(3));
    ocvmat.convertTo(cvmat, CV_8UC(3), 255.0, 0.0);
    cvtColor(cvmat, cvmat, COLOR_RGB2GRAY);
    Mat resized = new Mat(conf.getRescaledHeight(), conf.getRescaledWidth(), CV_8UC(1));
    resize(cvmat, resized, new Size(conf.getRescaledWidth(), conf.getRescaledHeight()));
    //   show(resized);
    //   waitKP();
    //Crop by croppingHeight, croppingHeight
    Mat cropped = resized.apply(new Rect(conf.getOffsetX(), conf.getOffsetY(), conf.getCroppingWidth(),
                    conf.getCroppingHeight()));
    //System.out.println(conf.getCroppingWidth() + " " + cropped.data().asBuffer().array().length);

    INDArray out = null;
    try {
        out = new NativeImageLoader(conf.getCroppingHeight(), conf.getCroppingWidth()).asMatrix(cropped);
    } catch (IOException e) {
        e.printStackTrace();
    }
    //System.out.println(out.shapeInfoToString());
    out = out.reshape(1, conf.getCroppingHeight(), conf.getCroppingWidth());
    INDArray compressed = out.castTo(DataType.UBYTE);
    return compressed;
}
 

private void calculate(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask, INDArray scoreOutput, INDArray gradientOutput) {
    if (scoreOutput == null && gradientOutput == null) {
        throw new IllegalArgumentException("You have to provide at least one of scoreOutput or gradientOutput!");
    }
    if (labels.size(1) != preOutput.size(1)) {
        throw new IllegalArgumentException(
                "Labels array numColumns (size(1) = " + labels.size(1) + ") does not match output layer"
                        + " number of outputs (nOut = " + preOutput.size(1) + ") ");

    }
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype
    final INDArray postOutput = activationFn.getActivation(preOutput.dup(), true);

    final INDArray positive = labels;
    final INDArray negative = labels.eq(0.0).castTo(Nd4j.defaultFloatingPointType());
    final INDArray normFactor = negative.sum(true,1).castTo(Nd4j.defaultFloatingPointType()).muli(positive.sum(true,1));


    long examples = positive.size(0);
    for (int i = 0; i < examples; i++) {
        final INDArray locCfn = postOutput.getRow(i, true);
        final long[] shape = locCfn.shape();

        final INDArray locPositive = positive.getRow(i, true);
        final INDArray locNegative = negative.getRow(i, true);
        final Double locNormFactor = normFactor.getDouble(i);

        final int outSetSize = locNegative.sumNumber().intValue();
        if(outSetSize == 0 || outSetSize == locNegative.columns()){
            if (scoreOutput != null) {
                scoreOutput.getRow(i, true).assign(0);
            }

            if (gradientOutput != null) {
                gradientOutput.getRow(i, true).assign(0);
            }
        }else {
            final INDArray operandA = Nd4j.ones(shape[1], shape[0]).mmul(locCfn);
            final INDArray operandB = operandA.transpose();

            final INDArray pairwiseSub = Transforms.exp(operandA.sub(operandB));

            final INDArray selection = locPositive.transpose().mmul(locNegative);

            final INDArray classificationDifferences = pairwiseSub.muli(selection).divi(locNormFactor);

            if (scoreOutput != null) {
                if (mask != null) {
                    final INDArray perLabel = classificationDifferences.sum(0);
                    LossUtil.applyMask(perLabel, mask.getRow(i, true));
                    perLabel.sum(scoreOutput.getRow(i, true), 0);
                } else {
                    classificationDifferences.sum(scoreOutput.getRow(i, true), 0, 1);
                }
            }

            if (gradientOutput != null) {
                gradientOutput.getRow(i, true).assign(classificationDifferences.sum(true, 0).addi(classificationDifferences.sum(true,1).transposei().negi()));
            }
        }
    }

    if (gradientOutput != null) {
        gradientOutput.assign(activationFn.backprop(preOutput.dup(), gradientOutput).getFirst());
        //multiply with masks, always
        if (mask != null) {
            LossUtil.applyMask(gradientOutput, mask);
        }
    }
}
 

@Override
public INDArray computeGradient(INDArray labels, INDArray preOutput, IActivation activationFn, INDArray mask) {
    if(!labels.equalShapes(preOutput)){
        Preconditions.throwEx("Labels and preOutput must have equal shapes: got shapes %s vs %s", labels.shape(), preOutput.shape());
    }
    INDArray grad;
    INDArray output = activationFn.getActivation(preOutput.dup(), true);
    labels = labels.castTo(preOutput.dataType());   //No-op if already correct dtype

    if (activationFn instanceof ActivationSoftmax) {

        if (mask != null && LossUtil.isPerOutputMasking(output, mask)) {
            throw new UnsupportedOperationException("Per output masking for MCXENT + softmax: not supported");
        }

        //Weighted loss function
        if (weights != null) {
            if (weights.length() != output.size(1)) {
                throw new IllegalStateException("Weights vector (length " + weights.length()
                                + ") does not match output.size(1)=" + output.size(1));
            }
            INDArray temp = labels.mulRowVector(weights.castTo(labels.dataType()));
            INDArray col = temp.sum(true,1);
            grad = output.mulColumnVector(col).sub(temp);
        } else {
            grad = output.subi(labels);
        }
    } else {
        INDArray dLda = output.rdivi(labels).negi();

        grad = activationFn.backprop(preOutput, dLda).getFirst(); //TODO activation function with weights

        //Weighted loss function
        if (weights != null) {
            if (weights.length() != output.size(1)) {
                throw new IllegalStateException("Weights vector (length " + weights.length()
                                + ") does not match output.size(1)=" + output.size(1));
            }
            grad.muliRowVector(weights.castTo(grad.dataType()));
        }
    }

    //Loss function with masking
    if (mask != null) {
        LossUtil.applyMask(grad, mask);
    }

    return grad;
}
 

@Test
public void testAsBytes() {
    INDArray orig = Nd4j.linspace(DataType.INT, 0, 10, 1);

    for (DataType dt : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF, DataType.BFLOAT16,
            DataType.LONG, DataType.INT, DataType.SHORT, DataType.BYTE, DataType.BOOL,
            DataType.UINT64, DataType.UINT32, DataType.UINT16, DataType.UBYTE}) {
        INDArray arr = orig.castTo(dt);

        byte[] b = arr.data().asBytes();        //NOTE: BIG ENDIAN

        if(ByteOrder.nativeOrder().equals(ByteOrder.LITTLE_ENDIAN)) {
            //Switch from big endian (as defined by asBytes which uses big endian) to little endian
            int w = dt.width();
            if (w > 1) {
                int len = b.length / w;
                for (int i = 0; i < len; i++) {
                    for (int j = 0; j < w / 2; j++) {
                        byte temp = b[(i + 1) * w - j - 1];
                        b[(i + 1) * w - j - 1] = b[i * w + j];
                        b[i * w + j] = temp;
                    }
                }
            }
        }

        INDArray arr2 = Nd4j.create(dt, arr.shape());
        ByteBuffer bb = arr2.data().pointer().asByteBuffer();
        bb.position(0);
        bb.put(b);

        Nd4j.getAffinityManager().tagLocation(arr2, AffinityManager.Location.HOST);

        assertEquals(arr.toString(), arr2.toString());
        assertEquals(arr, arr2);

        //Sanity check on data buffer getters:
        DataBuffer db = arr.data();
        DataBuffer db2 = arr2.data();
        for( int i=0; i<10; i++ ){
            assertEquals(db.getDouble(i), db2.getDouble(i), 0);
            assertEquals(db.getFloat(i), db2.getFloat(i), 0);
            assertEquals(db.getInt(i), db2.getInt(i), 0);
            assertEquals(db.getLong(i), db2.getLong(i), 0);
            assertEquals(db.getNumber(i), db2.getNumber(i));
        }

        assertArrayEquals(db.getDoublesAt(0, 10), db2.getDoublesAt(0, 10), 0);
        assertArrayEquals(db.getFloatsAt(0, 10), db2.getFloatsAt(0, 10), 0);
        assertArrayEquals(db.getIntsAt(0, 10), db2.getIntsAt(0, 10));
        assertArrayEquals(db.getLongsAt(0, 10), db2.getLongsAt(0, 10));
    }
}
 

@Test
public void testLocallyConnected() {
    for (DataType globalDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
        Nd4j.setDefaultDataTypes(globalDtype, globalDtype);
        for (DataType networkDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
            assertEquals(globalDtype, Nd4j.dataType());
            assertEquals(globalDtype, Nd4j.defaultFloatingPointType());

            INDArray[] in = null;
            for (int test = 0; test < 2; test++) {
                String msg = "Global dtype: " + globalDtype + ", network dtype: " + networkDtype + ", test=" + test;

                ComputationGraphConfiguration.GraphBuilder b = new NeuralNetConfiguration.Builder()
                        .dataType(networkDtype)
                        .seed(123)
                        .updater(new NoOp())
                        .weightInit(WeightInit.XAVIER)
                        .convolutionMode(ConvolutionMode.Same)
                        .graphBuilder();

                INDArray label;
                switch (test) {
                    case 0:
                        b.addInputs("in")
                                .addLayer("1", new LSTM.Builder().nOut(5).build(), "in")
                                .addLayer("2", new LocallyConnected1D.Builder().kernelSize(2).nOut(4).build(), "1")
                                .addLayer("out", new RnnOutputLayer.Builder().nOut(10).build(), "2")
                                .setOutputs("out")
                                .setInputTypes(InputType.recurrent(5, 2));
                        in = new INDArray[]{Nd4j.rand(networkDtype, 2, 5, 2)};
                        label = TestUtils.randomOneHotTimeSeries(2, 10, 2);
                        break;
                    case 1:
                        b.addInputs("in")
                                .addLayer("1", new ConvolutionLayer.Builder().kernelSize(2, 2).nOut(5).convolutionMode(ConvolutionMode.Same).build(), "in")
                                .addLayer("2", new LocallyConnected2D.Builder().kernelSize(2, 2).nOut(5).build(), "1")
                                .addLayer("out", new OutputLayer.Builder().nOut(10).build(), "2")
                                .setOutputs("out")
                                .setInputTypes(InputType.convolutional(8, 8, 1));
                        in = new INDArray[]{Nd4j.rand(networkDtype, 2, 1, 8, 8)};
                        label = TestUtils.randomOneHot(2, 10).castTo(networkDtype);
                        break;
                    default:
                        throw new RuntimeException();
                }

                ComputationGraph net = new ComputationGraph(b.build());
                net.init();

                INDArray out = net.outputSingle(in);
                assertEquals(msg, networkDtype, out.dataType());
                Map<String, INDArray> ff = net.feedForward(in, false);
                for (Map.Entry<String, INDArray> e : ff.entrySet()) {
                    if (e.getKey().equals("in"))
                        continue;
                    String s = msg + " - layer: " + e.getKey();
                    assertEquals(s, networkDtype, e.getValue().dataType());
                }

                net.setInputs(in);
                net.setLabels(label);
                net.computeGradientAndScore();

                net.fit(new MultiDataSet(in, new INDArray[]{label}));

                logUsedClasses(net);

                //Now, test mismatched dtypes for input/labels:
                for (DataType inputLabelDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
                    INDArray[] in2 = new INDArray[in.length];
                    for (int i = 0; i < in.length; i++) {
                        in2[i] = in[i].castTo(inputLabelDtype);
                    }
                    INDArray label2 = label.castTo(inputLabelDtype);
                    net.output(in2);
                    net.setInputs(in2);
                    net.setLabels(label2);
                    net.computeGradientAndScore();

                    net.fit(new MultiDataSet(in2, new INDArray[]{label2}));
                }
            }
        }
    }
}
 

public static INDArray maskedPoolingEpsilonTimeSeries(PoolingType poolingType, INDArray input, INDArray mask,
                INDArray epsilon2d, int pnorm) {

    if (input.rank() != 3) {
        throw new IllegalArgumentException("Expect rank 3 input activation array: got " + input.rank());
    }
    if (mask.rank() != 2) {
        throw new IllegalArgumentException("Expect rank 2 array for mask: got " + mask.rank());
    }
    if (epsilon2d.rank() != 2) {
        throw new IllegalArgumentException("Expected rank 2 array for errors: got " + epsilon2d.rank());
    }

    //Mask: [minibatch, tsLength]
    //Epsilon: [minibatch, vectorSize]

    mask = mask.castTo(input.dataType());

    switch (poolingType) {
        case MAX:
            INDArray negInfMask = mask.rsub(1.0);
            BooleanIndexing.replaceWhere(negInfMask, Double.NEGATIVE_INFINITY, Conditions.equals(1.0));

            INDArray withInf = Nd4j.createUninitialized(input.dataType(), input.shape());
            Nd4j.getExecutioner().exec(new BroadcastAddOp(input, negInfMask, withInf, 0, 2));
            //At this point: all the masked out steps have value -inf, hence can't be the output of the MAX op

            INDArray isMax = Nd4j.exec(new IsMax(withInf, withInf.ulike(), 2))[0];

            return Nd4j.getExecutioner().exec(new BroadcastMulOp(isMax, epsilon2d, isMax, 0, 1));
        case AVG:
        case SUM:
            //if out = sum(in,dims) then dL/dIn = dL/dOut -> duplicate to each step and mask
            //if out = avg(in,dims) then dL/dIn = 1/N * dL/dOut
            //With masking: N differs for different time series

            INDArray out = Nd4j.createUninitialized(input.dataType(), input.shape(), 'f');

            //Broadcast copy op, then divide and mask to 0 as appropriate
            Nd4j.getExecutioner().exec(new BroadcastCopyOp(out, epsilon2d, out, 0, 1));
            Nd4j.getExecutioner().exec(new BroadcastMulOp(out, mask, out, 0, 2));

            if (poolingType == PoolingType.SUM) {
                return out;
            }

            INDArray nEachTimeSeries = mask.sum(1); //[minibatchSize,tsLength] -> [minibatchSize,1]
            Nd4j.getExecutioner().exec(new BroadcastDivOp(out, nEachTimeSeries, out, 0));

            return out;

        case PNORM:
            //Similar to average and sum pooling: there's no N term here, so we can just set the masked values to 0
            INDArray masked2 = Nd4j.createUninitialized(input.dataType(), input.shape());
            Nd4j.getExecutioner().exec(new BroadcastMulOp(input, mask, masked2, 0, 2));

            INDArray abs = Transforms.abs(masked2, true);
            Transforms.pow(abs, pnorm, false);
            INDArray pNorm = Transforms.pow(abs.sum(2), 1.0 / pnorm);

            INDArray numerator;
            if (pnorm == 2) {
                numerator = input.dup();
            } else {
                INDArray absp2 = Transforms.pow(Transforms.abs(input, true), pnorm - 2, false);
                numerator = input.mul(absp2);
            }

            INDArray denom = Transforms.pow(pNorm, pnorm - 1, false);
            denom.rdivi(epsilon2d);
            Nd4j.getExecutioner().execAndReturn(new BroadcastMulOp(numerator, denom, numerator, 0, 1));
            Nd4j.getExecutioner().exec(new BroadcastMulOp(numerator, mask, numerator, 0, 2)); //Apply mask

            return numerator;
        default:
            throw new UnsupportedOperationException("Unknown or not supported pooling type: " + poolingType);
    }
}
 

@Override
public void eval(INDArray labelsArr, INDArray predictionsArr, INDArray maskArr) {
    Triple<INDArray,INDArray, INDArray> p = BaseEvaluation.reshapeAndExtractNotMasked(labelsArr, predictionsArr, maskArr, axis);
    INDArray labels = p.getFirst();
    INDArray predictions = p.getSecond();
    INDArray maskArray = p.getThird();

    if(labels.dataType() != predictions.dataType())
        labels = labels.castTo(predictions.dataType());

    if (!initialized) {
        initialize((int) labels.size(1));
    }
    //References for the calculations is this section:
    //https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm
    //https://en.wikipedia.org/wiki/Pearson_product-moment_correlation_coefficient#For_a_sample
    //Doing online calculation of means, sum of squares, etc.

    if (columnNames.size() != labels.size(1) || columnNames.size() != predictions.size(1)) {
        throw new IllegalArgumentException(
                        "Number of the columns of labels and predictions must match specification ("
                                        + columnNames.size() + "). Got " + labels.size(1) + " and "
                                        + predictions.size(1));
    }

    if (maskArray != null) {
        //Handle per-output masking. We are assuming *binary* masks here
        labels = labels.mul(maskArray);
        predictions = predictions.mul(maskArray);
    }

    labelsSumPerColumn.addi(labels.sum(0).castTo(labelsSumPerColumn.dataType()));

    INDArray error = predictions.sub(labels);
    INDArray absErrorSum = Nd4j.getExecutioner().exec(new ASum(error, 0));
    INDArray squaredErrorSum = error.mul(error).sum(0);

    sumAbsErrorsPerColumn.addi(absErrorSum.castTo(labelsSumPerColumn.dataType()));
    sumSquaredErrorsPerColumn.addi(squaredErrorSum.castTo(labelsSumPerColumn.dataType()));

    sumOfProducts.addi(labels.mul(predictions).sum(0).castTo(labelsSumPerColumn.dataType()));

    sumSquaredLabels.addi(labels.mul(labels).sum(0).castTo(labelsSumPerColumn.dataType()));
    sumSquaredPredicted.addi(predictions.mul(predictions).sum(0).castTo(labelsSumPerColumn.dataType()));


    val nRows = labels.size(0);

    INDArray newExampleCountPerColumn;
    if (maskArray == null) {
        newExampleCountPerColumn = exampleCountPerColumn.add(nRows);
    } else {
        newExampleCountPerColumn = exampleCountPerColumn.add(maskArray.sum(0).castTo(labelsSumPerColumn.dataType()));
    }
    currentMean.muliRowVector(exampleCountPerColumn).addi(labels.sum(0).castTo(labelsSumPerColumn.dataType())).diviRowVector(newExampleCountPerColumn);
    currentPredictionMean.muliRowVector(exampleCountPerColumn).addi(predictions.sum(0).castTo(labelsSumPerColumn.dataType()))
                    .divi(newExampleCountPerColumn);
    exampleCountPerColumn = newExampleCountPerColumn;

    sumLabels.addi(labels.sum(0).castTo(labelsSumPerColumn.dataType()));
}
 

public static INDArray maskedPoolingTimeSeries(PoolingType poolingType, INDArray toReduce, INDArray mask,
                int pnorm, DataType dataType) {
    if (toReduce.rank() != 3) {
        throw new IllegalArgumentException("Expect rank 3 array: got " + toReduce.rank());
    }
    if (mask.rank() != 2) {
        throw new IllegalArgumentException("Expect rank 2 array for mask: got " + mask.rank());
    }

    toReduce = toReduce.castTo(dataType);
    mask = mask.castTo(dataType);

    //Sum pooling: easy. Multiply by mask, then sum as normal
    //Average pooling: as above, but do a broadcast element-wise divi by mask.sum(1)
    //Max pooling: set to -inf if mask is 0, then do max as normal

    switch (poolingType) {
        case MAX:
            INDArray negInfMask = mask.castTo(dataType).rsub(1.0);
            BooleanIndexing.replaceWhere(negInfMask, Double.NEGATIVE_INFINITY, Conditions.equals(1.0));

            INDArray withInf = Nd4j.createUninitialized(dataType, toReduce.shape());
            Nd4j.getExecutioner().exec(new BroadcastAddOp(toReduce, negInfMask, withInf, 0, 2));
            //At this point: all the masked out steps have value -inf, hence can't be the output of the MAX op

            return withInf.max(2);
        case AVG:
        case SUM:
            INDArray masked = Nd4j.createUninitialized(dataType, toReduce.shape());
            Nd4j.getExecutioner().exec(new BroadcastMulOp(toReduce, mask, masked, 0, 2));
            INDArray summed = masked.sum(2);
            if (poolingType == PoolingType.SUM) {
                return summed;
            }

            INDArray maskCounts = mask.sum(1);
            summed.diviColumnVector(maskCounts);
            return summed;
        case PNORM:
            //Similar to average and sum pooling: there's no N term here, so we can just set the masked values to 0
            INDArray masked2 = Nd4j.createUninitialized(dataType, toReduce.shape());
            Nd4j.getExecutioner().exec(new BroadcastMulOp(toReduce, mask, masked2, 0, 2));

            INDArray abs = Transforms.abs(masked2, true);
            Transforms.pow(abs, pnorm, false);
            INDArray pNorm = abs.sum(2);

            return Transforms.pow(pNorm, 1.0 / pnorm);
        default:
            throw new UnsupportedOperationException("Unknown or not supported pooling type: " + poolingType);
    }
}
 

@Test
public void testEmbeddingDtypes() {
    for (DataType globalDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
        Nd4j.setDefaultDataTypes(globalDtype, globalDtype);
        for (DataType networkDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
            for (boolean frozen : new boolean[]{false, true}) {
                for (int test = 0; test < 3; test++) {
                    assertEquals(globalDtype, Nd4j.dataType());
                    assertEquals(globalDtype, Nd4j.defaultFloatingPointType());

                    String msg = "Global dtype: " + globalDtype + ", network dtype: " + networkDtype + ", test=" + test;

                    ComputationGraphConfiguration.GraphBuilder conf = new NeuralNetConfiguration.Builder()
                            .dataType(networkDtype)
                            .seed(123)
                            .updater(new NoOp())
                            .weightInit(new WeightInitDistribution(new UniformDistribution(-6, 6)))
                            .graphBuilder()
                            .addInputs("in")
                            .setOutputs("out");

                    INDArray input;
                    if (test == 0) {
                        if (frozen) {
                            conf.layer("0", new FrozenLayer(new EmbeddingLayer.Builder().nIn(5).nOut(5).build()), "in");
                        } else {
                            conf.layer("0", new EmbeddingLayer.Builder().nIn(5).nOut(5).build(), "in");
                        }
                        input = Nd4j.rand(networkDtype, 10, 1).muli(5).castTo(DataType.INT);
                        conf.setInputTypes(InputType.feedForward(1));
                    } else if (test == 1) {
                        if (frozen) {
                            conf.layer("0", new FrozenLayer(new EmbeddingSequenceLayer.Builder().nIn(5).nOut(5).build()), "in");
                        } else {
                            conf.layer("0", new EmbeddingSequenceLayer.Builder().nIn(5).nOut(5).build(), "in");
                        }
                        conf.layer("gp", new GlobalPoolingLayer.Builder(PoolingType.PNORM).pnorm(2).poolingDimensions(2).build(), "0");
                        input = Nd4j.rand(networkDtype, 10, 1, 5).muli(5).castTo(DataType.INT);
                        conf.setInputTypes(InputType.recurrent(1));
                    } else {
                        conf.layer("0", new RepeatVector.Builder().repetitionFactor(5).nOut(5).build(), "in");
                        conf.layer("gp", new GlobalPoolingLayer.Builder(PoolingType.SUM).build(), "0");
                        input = Nd4j.rand(networkDtype, 10, 5);
                        conf.setInputTypes(InputType.feedForward(5));
                    }

                    conf.appendLayer("el", new ElementWiseMultiplicationLayer.Builder().nOut(5).build())
                            .appendLayer("ae", new AutoEncoder.Builder().nOut(5).build())
                            .appendLayer("prelu", new PReLULayer.Builder().nOut(5).inputShape(5).build())
                            .appendLayer("out", new OutputLayer.Builder().nOut(10).build());

                    ComputationGraph net = new ComputationGraph(conf.build());
                    net.init();

                    INDArray label = Nd4j.zeros(networkDtype, 10, 10);

                    INDArray out = net.outputSingle(input);
                    assertEquals(msg, networkDtype, out.dataType());
                    Map<String, INDArray> ff = net.feedForward(input, false);
                    for (Map.Entry<String, INDArray> e : ff.entrySet()) {
                        if (e.getKey().equals("in"))
                            continue;
                        String s = msg + " - layer: " + e.getKey();
                        assertEquals(s, networkDtype, e.getValue().dataType());
                    }

                    net.setInput(0, input);
                    net.setLabels(label);
                    net.computeGradientAndScore();

                    net.fit(new DataSet(input, label));

                    logUsedClasses(net);

                    //Now, test mismatched dtypes for input/labels:
                    for (DataType inputLabelDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
                        INDArray in2 = input.castTo(inputLabelDtype);
                        INDArray label2 = label.castTo(inputLabelDtype);
                        net.output(in2);
                        net.setInput(0, in2);
                        net.setLabels(label2);
                        net.computeGradientAndScore();

                        net.fit(new DataSet(in2, label2));
                    }
                }
            }
        }
    }
}
 

@Test
public void testAttentionDTypes() {
    for (DataType globalDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
        Nd4j.setDefaultDataTypes(globalDtype, globalDtype);
        for (DataType networkDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
            assertEquals(globalDtype, Nd4j.dataType());
            assertEquals(globalDtype, Nd4j.defaultFloatingPointType());

            String msg = "Global dtype: " + globalDtype + ", network dtype: " + networkDtype;

            int mb = 3;
            int nIn = 3;
            int nOut = 5;
            int tsLength = 4;
            int layerSize = 8;
            int numQueries = 6;

            INDArray in = Nd4j.rand(networkDtype, new long[]{mb, nIn, tsLength});
            INDArray labels = TestUtils.randomOneHot(mb, nOut);

            MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                    .dataType(networkDtype)
                    .activation(Activation.TANH)
                    .updater(new NoOp())
                    .weightInit(WeightInit.XAVIER)
                    .list()
                    .layer(new LSTM.Builder().nOut(layerSize).build())
                    .layer(new SelfAttentionLayer.Builder().nOut(8).nHeads(2).projectInput(true).build())
                    .layer(new LearnedSelfAttentionLayer.Builder().nOut(8).nHeads(2).nQueries(numQueries).projectInput(true).build())
                    .layer(new RecurrentAttentionLayer.Builder().nIn(layerSize).nOut(layerSize).nHeads(1).projectInput(false).hasBias(false).build())
                    .layer(new GlobalPoolingLayer.Builder().poolingType(PoolingType.MAX).build())
                    .layer(new OutputLayer.Builder().nOut(nOut).activation(Activation.SOFTMAX)
                            .lossFunction(LossFunctions.LossFunction.MCXENT).build())
                    .setInputType(InputType.recurrent(nIn))
                    .build();

            MultiLayerNetwork net = new MultiLayerNetwork(conf);
            net.init();

            INDArray out = net.output(in);
            assertEquals(msg, networkDtype, out.dataType());
            List<INDArray> ff = net.feedForward(in);
            for (int i = 0; i < ff.size(); i++) {
                String s = msg + " - layer " + (i - 1) + " - " + (i == 0 ? "input" : net.getLayer(i - 1).conf().getLayer().getClass().getSimpleName());
                assertEquals(s, networkDtype, ff.get(i).dataType());
            }

            net.setInput(in);
            net.setLabels(labels);
            net.computeGradientAndScore();

            net.fit(new DataSet(in, labels));

            logUsedClasses(net);

            //Now, test mismatched dtypes for input/labels:
            for (DataType inputLabelDtype : new DataType[]{DataType.DOUBLE, DataType.FLOAT, DataType.HALF}) {
                INDArray in2 = in.castTo(inputLabelDtype);
                INDArray label2 = labels.castTo(inputLabelDtype);
                net.output(in2);
                net.setInput(in2);
                net.setLabels(label2);
                net.computeGradientAndScore();

                net.fit(new DataSet(in2, label2));
            }
        }
    }
}