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

@Override
public Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr) {
    assertInputSet(true);

    if(epsilon.dataType() != dataType){
        epsilon = epsilon.castTo(dataType);
    }

    INDArray outEpsilon;
    try(MemoryWorkspace ws = workspaceMgr.notifyScopeBorrowed(ArrayType.ACTIVATION_GRAD)){
        if (layerConf().getDataFormat() == RNNFormat.NCW) {
            outEpsilon = epsilon.sum(2);
        }else{
            outEpsilon = epsilon.sum(1);
        }
    }

    Gradient gradient = new DefaultGradient();
    return new Pair<>(gradient, outEpsilon);
}
 

/**
 * Validate that the operation is being applied on a floating point type INDArray
 *
 * @param opName    Operation name to print in the exception
 * @param inputName Name of the input to the op to validate
 * @param v         Variable to validate datatype for (input to operation)
 */
public static void validateFloatingPoint(String opName, String inputName, INDArray v) {
    if (v == null)
        return;
    if (!v.dataType().isFPType())
        throw new IllegalStateException("Input \"" + inputName + "\" for operation \"" + opName +
                "\" must be an floating point type; got array with non-floating point data type " + v.dataType());
}
 

/**
 * Validate that the operation is being applied on a boolean type INDArray
 *
 * @param opName Operation name to print in the exception
 * @param v      Variable to validate datatype for (input to operation)
 */
public static void validateBool(String opName, INDArray v) {
    if (v == null)
        return;
    if (v.dataType() != DataType.BOOL)
        throw new IllegalStateException("Cannot apply operation \"" + opName + "\" to array with non-boolean point data type " + v.dataType());
}
 

/**
 * Validate that the operation is being applied on an integer type INDArray
 *
 * @param opName    Operation name to print in the exception
 * @param inputName Name of the input to the op to validate
 * @param v         Variable to validate datatype for (input to operation)
 */
public static void validateInteger(String opName, String inputName, INDArray v) {
    if (v == null)
        return;
    if (!v.dataType().isIntType())
        throw new IllegalStateException("Input \"" + inputName + "\" for operation \"" + opName + "\" must be an integer" +
                " type; got array with non-integer data type " + v.dataType());
}
 

/**
 * Validate that the operation is being applied on an integer type INDArray
 *
 * @param opName Operation name to print in the exception
 * @param v      Variable to validate datatype for (input to operation)
 */
public static void validateInteger(String opName, INDArray v) {
    if (v == null)
        return;
    if (!v.dataType().isIntType())
        throw new IllegalStateException("Cannot apply operation \"" + opName + "\" to array with non-integer data type " + v.dataType());
}
 

public FusedBatchNorm(@NonNull INDArray x, @NonNull INDArray scale, @NonNull INDArray offset,
                      int dataFormat, int isTraining,
                      INDArray yOut, INDArray batchMeanOut, INDArray batchMeanVar) {
    addInputArgument(x, scale, offset);
    addIArgument(dataFormat, isTraining);
    if (yOut != null && batchMeanOut != null && batchMeanVar != null) {
        addOutputArgument(yOut, batchMeanOut, batchMeanVar);
    }
    this.outputDataType = x.dataType();
}
 

/**
 * Validate that the operation is being applied on a numerical INDArray (not boolean or utf8).
 * Some operations (such as sum, norm2, add(Number) etc) don't make sense when applied to boolean/utf8 arrays
 *
 * @param opName Operation name to print in the exception
 * @param v      Variable to validate datatype for (input to operation)
 */
public static void validateNumerical(String opName, String inputName, INDArray v) {
    if (v == null)
        return;
    if (v.dataType() == DataType.BOOL || v.dataType() == DataType.UTF8)
        throw new IllegalStateException("Input \"" + inputName + "\" for operation \"" + opName + "\" must be an numerical type type;" +
                " got array with non-integer data type " + v.dataType());
}
 

public PythonObject(NumpyArray npArray) {
    int numpyType;
    INDArray indArray = npArray.getNd4jArray();
    DataType dataType = indArray.dataType();

    switch (dataType) {
        case DOUBLE:
            numpyType = NPY_DOUBLE;
            break;
        case FLOAT:
        case BFLOAT16:
            numpyType = NPY_FLOAT;
            break;
        case SHORT:
            numpyType = NPY_SHORT;
            break;
        case INT:
            numpyType = NPY_INT;
            break;
        case LONG:
            numpyType = NPY_INT64;
            break;
        case UINT16:
            numpyType = NPY_USHORT;
            break;
        case UINT32:
            numpyType = NPY_UINT;
            break;
        case UINT64:
            numpyType = NPY_UINT64;
            break;
        case BOOL:
            numpyType = NPY_BOOL;
            break;
        case BYTE:
            numpyType = NPY_BYTE;
            break;
        case UBYTE:
            numpyType = NPY_UBYTE;
            break;
        case HALF:
            numpyType = NPY_HALF;
            break;
        default:
            throw new RuntimeException("Unsupported dtype: " + npArray.getDtype());
    }

    long[] shape = indArray.shape();
    INDArray inputArray = indArray;
    if(dataType == DataType.BFLOAT16) {
        log.warn("\n\nThe given nd4j array \n\n{}\n\n is of BFLOAT16 datatype. " +
                "Casting a copy of it to FLOAT and creating the respective numpy array from it.\n", indArray);
        inputArray = indArray.castTo(DataType.FLOAT);
    }

    //Sync to host memory in the case of CUDA, before passing the host memory pointer to Python
    if(inputArray.data() instanceof BaseDataBuffer){
        ((BaseDataBuffer)inputArray.data()).syncToPrimary();
    }

    nativePythonObject = PyArray_New(PyArray_Type(), shape.length, new SizeTPointer(shape),
            numpyType, null,
            inputArray.data().addressPointer(),
            0, NPY_ARRAY_CARRAY, null);

}
 

public static String shapeToStringShort(INDArray arr){
    long[] s = arr.shape();
    return arr.dataType() + "," + (s == null ? "[]" : Arrays.toString(s).replace(" ","")) + "," + arr.ordering();
}
 

public ZerosLike(INDArray in, INDArray out){
    this(in, out, in.dataType());
}
 

public static boolean isSameType(INDArray x, INDArray y) {
    return x.dataType() == y.dataType();
}
 

private static String createString(INDArray arr){
    StringBuilder sb = new StringBuilder();

    if(arr.isEmpty()){
        sb.append("Nd4j.empty(DataType.").append(arr.dataType()).append(");");
    } else {
        sb.append("Nd4j.createFromArray(");

        DataType dt = arr.dataType();
        switch (dt){
            case DOUBLE:
                double[] dArr = arr.dup().data().asDouble();
                sb.append(Arrays.toString(dArr).replaceAll("[\\[\\]]", ""));
                break;
            case FLOAT:
            case HALF:
            case BFLOAT16:
                float[] fArr = arr.dup().data().asFloat();
                sb.append(Arrays.toString(fArr)
                        .replaceAll(",", "f,")
                        .replaceAll("]", "f")
                        .replaceAll("[\\[\\]]", ""));
                break;
            case LONG:
            case UINT32:
            case UINT64:
                long[] lArr = arr.dup().data().asLong();
                sb.append(Arrays.toString(lArr)
                        .replaceAll(",", "L,")
                        .replaceAll("]", "L")
                        .replaceAll("[\\[\\]]", ""));
                break;
            case INT:
            case SHORT:
            case UBYTE:
            case BYTE:
            case UINT16:
            case BOOL:
                int[] iArr = arr.dup().data().asInt();
                sb.append(Arrays.toString(iArr).replaceAll("[\\[\\]]", ""));
                break;
            case UTF8:
                break;
            case COMPRESSED:
            case UNKNOWN:
                break;
        }

        sb.append(").reshape(").append(Arrays.toString(arr.shape()).replaceAll("[\\[\\]]", ""))
                .append(")");

        if(dt == DataType.HALF || dt == DataType.BFLOAT16 || dt == DataType.UINT32 || dt == DataType.UINT64 ||
                dt == DataType.SHORT || dt == DataType.UBYTE || dt == DataType.BYTE || dt == DataType.UINT16 || dt == DataType.BOOL){
            sb.append(".cast(DataType.").append(arr.dataType()).append(")");
        }
    }

    return sb.toString();
}
 

public static INDArray maskedPoolingConvolution(PoolingType poolingType, INDArray toReduce, INDArray mask, int pnorm, DataType dataType) {
    if(mask.rank() != 4){
        //TODO BETTER ERROR MESSAGE EXPLAINING FORMAT
        //TODO ALSO HANDLE LEGACY FORMAT WITH WARNING WHERE POSSIBLE
        throw new IllegalStateException("Expected rank 4 mask array: Got array with shape " + Arrays.toString(mask.shape()));
    }

    mask = mask.castTo(dataType);   //no-op if already correct dtype

    // [minibatch, channels, h, w] data with a mask array of shape [minibatch, 1, X, Y]
    // where X=(1 or inH) and Y=(1 or inW)

    //General case: must be equal or 1 on each dimension
    int[] dimensions = new int[4];
    int count = 0;
    for(int i=0; i<4; i++ ){
        if(toReduce.size(i) == mask.size(i)){
            dimensions[count++] = i;
        }
    }
    if(count < 4){
        dimensions = Arrays.copyOfRange(dimensions, 0, count);
    }

    switch (poolingType) {
        case MAX:
            //TODO This is ugly - replace it with something better... Need something like a Broadcast CAS op
            INDArray negInfMask;
            if(mask.dataType() == DataType.BOOL){
                negInfMask = Transforms.not(mask).castTo(dataType);
            } else {
                negInfMask = mask.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, dimensions));
            //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, 3);
        case AVG:
        case SUM:
            INDArray masked = Nd4j.createUninitialized(dataType, toReduce.shape());
            Nd4j.getExecutioner().exec(new BroadcastMulOp(toReduce, mask, masked, dimensions));

            INDArray summed = masked.sum(2, 3);
            if (poolingType == PoolingType.SUM) {
                return summed;
            }
            INDArray maskCounts = mask.sum(1,2,3);
            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, dimensions));

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

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

@Override
public void release(@NonNull INDArray array) {
    //Check for multiple releases of the array
    long id = array.getId();
    Preconditions.checkState(!lruCache.contains(id), "Array was released multiple times: id=%s, shape=%ndShape", id, array);


    DataType dt = array.dataType();
    long thisBytes = array.data().length() * dt.width();
    if(array.dataType() == DataType.UTF8) {
        //Don't cache string arrays due to variable length buffers
        if(array.closeable())
            array.close();
    } else if (currentCacheSize + thisBytes > maxCacheBytes) {
        if(thisBytes > maxCacheBytes){
            //Can't store even if we clear everything - too large
            if(array.closeable())
                array.close();
            return;
        }

        //Need to deallocate some arrays to stay under limit - do in "oldest first" order
        Iterator<Long> iter = lruCache.iterator();
        while(currentCacheSize + thisBytes > maxCacheBytes){
            long next = iter.next();
            iter.remove();
            INDArray nextOldest = lruCacheValues.remove(next);
            DataType ndt = nextOldest.dataType();
            long nextBytes = ndt.width() * nextOldest.data().length();
            arrayStores.get(ndt).removeObject(nextOldest);
            currentCacheSize -= nextBytes;

            if(nextOldest.closeable())
                nextOldest.close();
        }

        //After clearing space - can now cache
        cacheArray(array);
    } else {
        //OK to cache
        cacheArray(array);
    }

    //Store in LRU cache for "last used" removal if we exceed cache size
    lruCache.add(array.getId());
    lruCacheValues.put(array.getId(), array);
}
 

@Override
protected Map<String, INDArray> preprocessPlaceholders(Map<String, INDArray> placeholders, At at) {
    arrayUseTracker.clear();

    //We'll also use this method as a "pre execution" hook-in, to mark variables as something we should never deallocate
    //This occurs by never marking these "ConstantDep" and "VariableDep" instances as satisfied, so there's always
    // an unsatisfied dependency for them in the array use tracker
    //TODO we shouldn't be clearing this on every single iteration, in 99.5% of cases variables will be same as last iteration...
    for (SDVariable v : sameDiff.variables()) {
        if (v.getVariableType() == VariableType.CONSTANT) {
            arrayUseTracker.addDependency(v.getArr(), new ConstantDep(v.name()));
        } else if (v.getVariableType() == VariableType.VARIABLE) {
            arrayUseTracker.addDependency(v.getArr(), new VariableDep(v.name()));
        }
    }

    //Workaround for some TF/Keras based models that require explicit train/test as a placeholder
    boolean kerasWorkaround = false;
    List<String> phs = sameDiff.inputs();
    if (phs != null && !phs.isEmpty()) {
        for (String s : phs) {
            if (s.endsWith(KERAS_TRAIN_TEST) && !placeholders.containsKey(s)) {
                // The behaviour of some Keras layers (like GRU) differs depending on whether the model is training.
                // We provide this value directly, unless the user has provided this manually
                INDArray scalar = mmgr.allocate(false, DataType.BOOL).assign(at.operation().isTrainingPhase());
                placeholders = new HashMap<>(placeholders); //Array might be singleton, or otherwise unmodifiable
                placeholders.put(s, scalar);
                kerasWorkaround = true;
            }
        }
    }


    if (placeholders == null || placeholders.isEmpty()) {
        return placeholders;
    }

    //Handle casting of the input array automatically.
    //The idea here is to avoid unexpected errors if the user (for example) tries to perform inference with a double
    // array for a float placeholder
    //TODO eventually we might have ops that support multiple input types, and hence won't need this casting
    Map<String, INDArray> out = new HashMap<>();
    for (Map.Entry<String, INDArray> e : placeholders.entrySet()) {
        Preconditions.checkState(sameDiff.hasVariable(e.getKey()), "Invalid placeholder passed for execution: " +
                "No variable/placeholder with name %s exists", e.getKey());
        INDArray arr = e.getValue();
        //First: check workspaces
        if (arr.isAttached()) {
            MemoryWorkspace ws = arr.data() == null ? null : arr.data().getParentWorkspace();
            if (ws != null && ws.getWorkspaceType() != MemoryWorkspace.Type.CIRCULAR) {
                if (!ws.isScopeActive()) {
                    throw new ND4JIllegalStateException("Placeholder \"" + e.getKey() + "\" array uses leaked workspace pointer from workspace ["
                            + ws.getId() + "]: Workspace the array was defined in is no longer open.\nAll open workspaces: " + DefaultOpExecutioner.allOpenWorkspaces()
                            + "\n" + SCOPE_PANIC_MSG);
                }

                if (ws.getGenerationId() != arr.data().getGenerationId())
                    throw new ND4JIllegalStateException("Placeholder \"" + e.getKey() + "\" array uses outdated workspace pointer from workspace ["
                            + ws.getId() + "]: Workspace array was defined in has been closed and reopened at least once since array creation. Array WS iteration: " +
                            arr.data().getGenerationId() + ". Workspace current iteration: " +
                            ws.getGenerationId() + "\nAll open workspaces: " + DefaultOpExecutioner.allOpenWorkspaces() + "\n" + SCOPE_PANIC_MSG);
            }
        }


        //Second: cast the input to the required type
        //TODO For the casting case, we SHOULD actually deallocate this when we're done with it, which is usually sooner than "exec done"
        DataType dt = sameDiff.getVariable(e.getKey()).dataType();
        if (kerasWorkaround && e.getKey().endsWith(KERAS_TRAIN_TEST)) {
            arrayUseTracker.addDependency(arr, new ExecDoneDep());
        } else if (arr.dataType() == dt) {
            //Mark as a placeholder array in the array use tracker, so we never deallocate this array...
            arrayUseTracker.addDependency(e.getValue(), new PlaceholderDep(e.getKey()));
        } else {
            INDArray cast = mmgr.allocate(false, dt, arr.shape());
            cast.assign(arr);
            arr = cast;
            //This array CAN be deallocated once consumed, because of the cast
            //TODO we can likely close this sooner
            arrayUseTracker.addDependency(arr, new ExecDoneDep());
        }
        out.put(e.getKey(), arr);
    }

    return out;
}
 

@Override
public INDArray preOutput(INDArray input, INDArray weights, INDArray bias, int[] kernel, int[] strides, int[] pad,
                          ConvolutionLayer.AlgoMode mode, ConvolutionLayer.FwdAlgo fwdAlgo, ConvolutionMode convolutionMode,
                          int[] dilation, CNN2DFormat format, LayerWorkspaceMgr workspaceMgr) {
    if(input.dataType() != DataType.FLOAT || weights.dataType() != DataType.FLOAT)
        return null;    //MKL-DNN only supports floating point dtype


    int hDim = 2;
    int wDim = 3;
    if(format == CNN2DFormat.NHWC){
        hDim = 1;
        wDim = 2;
    }

    int inH = (int)input.size(hDim);
    int inW = (int)input.size(wDim);
    int[] outSize;
    if (convolutionMode == ConvolutionMode.Same) {
        outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, null, convolutionMode, dilation, format); //Also performs validation
        pad = ConvolutionUtils.getSameModeTopLeftPadding(outSize, new int[] {inH, inW}, kernel, strides, dilation);
    } else {
        outSize = ConvolutionUtils.getOutputSize(input, kernel, strides, pad, convolutionMode, dilation, format); //Also performs validation
    }

    if(context == null ){
        context = Nd4j.getExecutioner().buildContext();
        context.setIArguments(kernel[0], kernel[1],
                strides[0], strides[1],
                pad[0], pad[1],
                dilation[0], dilation[1],
                ArrayUtil.fromBoolean(convolutionMode == ConvolutionMode.Same),
                format == CNN2DFormat.NCHW ? 0 : 1,  //0=NCHW, 1=NHWC
                1   //Weight format: 1 - [oC, iC, kH, kW]
        );
    };

    int outDepth = (int) weights.size(0);
    long[] outShape = (format == CNN2DFormat.NCHW) ? new long[]{input.size(0), outDepth, outSize[0], outSize[1]} : new long[]{input.size(0), outSize[0], outSize[1], outDepth};
    INDArray out = workspaceMgr.createUninitialized(ArrayType.ACTIVATIONS, input.dataType(), outShape);

    INDArray[] inputsArr = bias == null ? new INDArray[]{input, weights} : new INDArray[]{input, weights, bias};
    context.purge();
    for( int i=0; i<inputsArr.length; i++ ){
        context.setInputArray(i, inputsArr[i]);
    }

    context.setOutputArray(0, out);
    Conv2D op = new Conv2D();
    Nd4j.exec(op, context);

    return out;
}
 

@Override
public Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray weights, INDArray bias, INDArray delta, int[] kernel, int[] strides, int[] pad,
                                                 INDArray biasGradView, INDArray weightGradView, IActivation afn, ConvolutionLayer.AlgoMode mode,
                                                 ConvolutionLayer.BwdFilterAlgo bwdFilterAlgo, ConvolutionLayer.BwdDataAlgo bwdDataAlgo, ConvolutionMode convolutionMode,
                                                 int[] dilation, CNN2DFormat format, LayerWorkspaceMgr workspaceMgr) {
    if(input.dataType() != DataType.FLOAT || weights.dataType() != DataType.FLOAT)
        return null;    //MKL-DNN only supports floating point dtype

    int hDim = 2;
    int wDim = 3;
    if(format == CNN2DFormat.NHWC){
        hDim = 1;
        wDim = 2;
    }

    if (convolutionMode == ConvolutionMode.Same) {
        pad = ConvolutionUtils.getSameModeTopLeftPadding(new int[]{(int)delta.size(hDim), (int)delta.size(wDim)}, new int[] {(int) input.size(hDim), (int) input.size(wDim)},
                kernel, strides, dilation);
    }

    if(contextBwd == null){
        contextBwd = Nd4j.getExecutioner().buildContext();
        contextBwd.setIArguments(kernel[0], kernel[1],
                strides[0], strides[1],
                pad[0], pad[1],
                dilation[0], dilation[1],
                ArrayUtil.fromBoolean(convolutionMode == ConvolutionMode.Same),
                format == CNN2DFormat.NCHW ? 0 : 1,  //0=NCHW, 1=NHWC
                1   //Weight format: 1 - [oC, iC, kH, kW]
        );
    };

    INDArray gradAtInput = workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, input.dataType(), input.shape());

    INDArray[] inputsArr = biasGradView == null ? new INDArray[]{input, weights, delta} : new INDArray[]{input, weights, bias, delta};
    INDArray[] outputArr = biasGradView == null ? new INDArray[]{gradAtInput, weightGradView} : new INDArray[]{gradAtInput, weightGradView, biasGradView};
    contextBwd.purge();
    for( int i=0; i<inputsArr.length; i++ ){
        contextBwd.setInputArray(i, inputsArr[i]);
    }
    for( int i=0; i<outputArr.length; i++ ){
        contextBwd.setOutputArray(i, outputArr[i]);
    }

    Conv2DDerivative op = new Conv2DDerivative();
    Nd4j.exec(op, contextBwd);

    Gradient g = new DefaultGradient();
    if(biasGradView != null) {
        g.gradientForVariable().put(ConvolutionParamInitializer.BIAS_KEY, biasGradView);
    }
    g.gradientForVariable().put(ConvolutionParamInitializer.WEIGHT_KEY, weightGradView);

    return new Pair<>(g, gradAtInput);
}
 

@Override
public INDArray preOutput(INDArray x, boolean training, long[] shape, INDArray gamma, INDArray beta, INDArray mean, INDArray var,
                          double decay, double eps, CNN2DFormat format, LayerWorkspaceMgr workspaceMgr) {
    if(x.dataType() != DataType.FLOAT)
        return null;    //MKL-DNN only supports float

    int axis = (x.rank() != 4 || format == CNN2DFormat.NCHW) ? 1 : 3;

    if(context == null){
        context = Nd4j.getExecutioner().buildContext();
        context.setIArguments(
                ArrayUtil.fromBoolean(gamma != null),
                ArrayUtil.fromBoolean(beta != null),
                axis);   //Axis - 1 = NCHW, 3 = NHWC
        context.setTArguments(eps);
    }

    //Mean and variance: args here are *global*. Depending on train/test mode we might need to use batch mean/var
    INDArray m, v;
    if(training){
        if(meanCache == null){
            try(MemoryWorkspace ws = Nd4j.getMemoryManager().scopeOutOfWorkspaces()) {
                meanCache = Nd4j.createUninitialized(x.dataType(), x.size(axis));
                varCache = Nd4j.createUninitialized(x.dataType(), x.size(axis));
            }
        }

        int[] dims;
        if(x.rank() == 2){
            dims = RANK2_DIMS;
        } else if(format == CNN2DFormat.NCHW){
            dims = RANK4_DIMS_NCHW;
        } else {
            dims = RANK4_DIMS_NHWC;
        }

        x.mean(meanCache, dims);
        Nd4j.exec(new Variance(x, varCache, false, dims));

        m = meanCache;
        v = varCache;
    } else {
        m = mean.reshape(mean.length());
        v = var.reshape(var.length());
    }

    //Note: batchnorm op expects rank 1 inputs for mean/var etc, not rank 2 shape [1,x]
    context.purge();
    context.setInputArray(0, x);
    context.setInputArray(1, m);
    context.setInputArray(2, v);
    if(gamma != null && beta != null) {
        context.setInputArray(3, gamma.rank() == 2 ? gamma.reshape(gamma.length()) : gamma);
        context.setInputArray(4, beta.rank() == 2 ? beta.reshape(beta.length()) : beta);
    }

    INDArray out = workspaceMgr.createUninitialized(ArrayType.ACTIVATIONS, x.dataType(), x.shape());
    context.setOutputArray(0, out);

    BatchNorm bn = new BatchNorm();
    Nd4j.exec(bn, context);
    return out;
}
 

@Override
public Pair<Gradient, INDArray> backpropGradient(INDArray input, INDArray epsilon, long[] shape, INDArray gamma, INDArray beta,
                                                 INDArray dGammaView, INDArray dBetaView, double eps, CNN2DFormat format, LayerWorkspaceMgr layerWorkspaceMgr) {

    boolean nchw = format == CNN2DFormat.NCHW;

    this.eps = eps;

    int cudnnTensorFormat = nchw ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
    int chIdx = nchw ? 1 : 3;
    int hIdx = nchw ? 2 : 1;
    int wIdx = nchw ? 3 : 2;

    val miniBatch = (int) input.size(0);
    val depth = (int) input.size(chIdx);
    val inH = (int) input.size(hIdx);
    val inW = (int) input.size(wIdx);

    final boolean isHalf = (input.dataType() == DataType.HALF);
    INDArray gammaOrig = null;
    INDArray dGammaViewOrig = null;
    INDArray dBetaViewOrig = null;
    if(isHalf) {    //Convert FP16 to FP32 if required (CuDNN BN doesn't support FP16 for these params, only for input/output)
        gammaOrig = gamma;
        dGammaViewOrig = dGammaView;
        dBetaViewOrig = dBetaView;
        /*
        From CuDNN docs: bnScale, resultBnScaleDiff, resultBnBiasDiff, savedMean, savedInvVariance
        "Note: The data type of this tensor descriptor must be 'float' for FP16 and FP32 input tensors, and 'double'
        for FP64 input tensors."
        >> Last 2 are the meanCache and varCache; first 3 are below
         */
        gamma = gamma.castTo(DataType.FLOAT);
        dGammaView = dGammaView.castTo(DataType.FLOAT);
        dBetaView = dBetaView.castTo(DataType.FLOAT);
    }

    Gradient retGradient = new DefaultGradient();

    if (!Shape.hasDefaultStridesForShape(epsilon)) {
        // apparently not supported by cuDNN
        epsilon = epsilon.dup('c');
    }

    val srcStride = ArrayUtil.toInts(input.stride());
    val deltaStride = ArrayUtil.toInts(epsilon.stride());

    if (Nd4j.getExecutioner() instanceof GridExecutioner)
        ((GridExecutioner) Nd4j.getExecutioner()).flushQueue();

    checkCudnn(cudnnSetTensor4dDescriptorEx(cudnnContext.srcTensorDesc, dataType, (int) miniBatch, (int) depth, (int) inH, (int) inW,
            (int) srcStride[0], (int) srcStride[chIdx], (int) srcStride[hIdx], (int) srcStride[wIdx]));
    checkCudnn(cudnnSetTensor4dDescriptorEx(cudnnContext.deltaTensorDesc, dataType, (int) miniBatch, (int) depth, (int) inH, (int) inW,
            (int) deltaStride[0], (int) deltaStride[chIdx], (int) deltaStride[hIdx], (int) deltaStride[wIdx]));

    long[] nextEpsShape = nchw ? new long[] {miniBatch, depth, inH, inW} : new long[] {miniBatch, inH, inW, depth};
    INDArray nextEpsilon = layerWorkspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, input.dataType(), nextEpsShape, 'c');
    val dstStride = ArrayUtil.toInts(nextEpsilon.stride());

    checkCudnn(cudnnSetTensor4dDescriptorEx(cudnnContext.dstTensorDesc, dataType, miniBatch, depth, inH, inW,
                    dstStride[0], dstStride[chIdx], dstStride[hIdx], dstStride[wIdx]));
    checkCudnn(cudnnSetTensor4dDescriptor(cudnnContext.gammaBetaTensorDesc, cudnnTensorFormat, toCudnnDataType(gamma.data().dataType()), (int)shape[0],
            (int)shape[1], shape.length > 2 ? (int)shape[2] : 1, shape.length > 3 ? (int)shape[3] : 1));

    Allocator allocator = AtomicAllocator.getInstance();
    CudaContext context = allocator.getFlowController().prepareActionAllWrite(input, epsilon, nextEpsilon, gamma,
                    dGammaView, dBetaView);
    Pointer srcData = allocator.getPointer(input, context);
    Pointer epsData = allocator.getPointer(epsilon, context);
    Pointer dstData = allocator.getPointer(nextEpsilon, context);
    Pointer gammaData = allocator.getPointer(gamma, context);
    Pointer dGammaData = allocator.getPointer(dGammaView, context);
    Pointer dBetaData = allocator.getPointer(dBetaView, context);
    Pointer meanCacheData = allocator.getPointer(meanCache, context);
    Pointer varCacheData = allocator.getPointer(varCache, context);

    checkCudnn(cudnnSetStream(cudnnContext, new CUstream_st(context.getCublasStream())));
    checkCudnn(cudnnBatchNormalizationBackward(cudnnContext, batchNormMode, alpha, this.beta, alpha, alpha,
                    cudnnContext.srcTensorDesc, srcData, cudnnContext.deltaTensorDesc, epsData,
                    cudnnContext.dstTensorDesc, dstData, cudnnContext.gammaBetaTensorDesc, gammaData, dGammaData,
                    dBetaData, eps, meanCacheData, varCacheData));

    allocator.getFlowController().registerActionAllWrite(context, input, epsilon, nextEpsilon, gamma, dGammaView,
                    dBetaView);

    retGradient.setGradientFor(BatchNormalizationParamInitializer.GAMMA, dGammaView);
    retGradient.setGradientFor(BatchNormalizationParamInitializer.BETA, dBetaView);

    context.syncOldStream();

    //Convert back and assign, if required:
    if(isHalf){
        gammaOrig.assign(gamma.castTo(DataType.HALF));
        dGammaViewOrig.assign(dGammaView.castTo(DataType.HALF));
        dBetaViewOrig.assign(dBetaView.castTo(DataType.HALF));
    }

    return new Pair<>(retGradient, nextEpsilon);
}
 

/**
 * Validate that the operation is being applied on a numerical INDArray (not boolean or utf8).
 * Some operations (such as sum, norm2, add(Number) etc) don't make sense when applied to boolean/utf8 arrays
 *
 * @param opName Operation name to print in the exception
 * @param v      Variable to perform operation on
 */
public static void validateNumerical(String opName, INDArray v) {
    if (v == null)
        return;
    if (v.dataType() == DataType.BOOL || v.dataType() == DataType.UTF8)
        throw new IllegalStateException("Cannot apply operation \"" + opName + "\" to array with non-numerical data type " + v.dataType());
}