nd4j-backends/nd4j-backend-impls/nd4j-cuda/src/main/java/org/nd4j/linalg/jcublas/JCublasNDArrayFactory.java
@Override
    public void shuffle(List<INDArray> arrays, Random rnd, List<int[]> dimensions) {
        // no dimension - no shuffle
        if (dimensions == null || dimensions.size() == 0)
            throw new RuntimeException("Dimension can't be null or 0-length");

        if (arrays == null || arrays.size() == 0)
            throw new RuntimeException("No input arrays provided");

        if (dimensions.size() > 1 && arrays.size() != dimensions.size())
            throw new IllegalStateException("Number of dimensions do not match number of arrays to shuffle");

        Nd4j.getExecutioner().push();

        // first we build TAD for input array and dimensions

        AtomicAllocator allocator = AtomicAllocator.getInstance();

        CudaContext context = null;

        for (int x = 0; x < arrays.size(); x++) {
            context = allocator.getFlowController().prepareAction(arrays.get(x));
        }

        int tadLength = 1;
        for (int i = 0; i < dimensions.get(0).length; i++) {
            tadLength *= arrays.get(0).shape()[dimensions.get(0)[i]];
        }

        val numTads = arrays.get(0).length() / tadLength;

        val map = ArrayUtil.buildInterleavedVector(rnd, (int) numTads);

        val shuffle = new CudaIntDataBuffer(map);

        Pointer shuffleMap = allocator.getPointer(shuffle, context);

        PointerPointer extras = new PointerPointer(null, // not used
                        context.getOldStream(), allocator.getDeviceIdPointer());


        long[] xPointers = new long[arrays.size()];
        long[] xShapes = new long[arrays.size()];
        long[] tadShapes = new long[arrays.size()];
        long[] tadOffsets = new long[arrays.size()];

        for (int i = 0; i < arrays.size(); i++) {
            INDArray array = arrays.get(i);

            Pointer x = AtomicAllocator.getInstance().getPointer(array, context);
            Pointer xShapeInfo = AtomicAllocator.getInstance().getPointer(array.shapeInfoDataBuffer(), context);


            TADManager tadManager = Nd4j.getExecutioner().getTADManager();

            int[] dimension = dimensions.size() > 1 ? dimensions.get(i) : dimensions.get(0);

            Pair<DataBuffer, DataBuffer> tadBuffers = tadManager.getTADOnlyShapeInfo(array, dimension);

//            log.info("Original shape: {}; dimension: {}; TAD shape: {}", array.shapeInfoDataBuffer().asInt(), dimension, tadBuffers.getFirst().asInt());

            Pointer tadShapeInfo = AtomicAllocator.getInstance().getPointer(tadBuffers.getFirst(), context);

            DataBuffer offsets = tadBuffers.getSecond();

            if (offsets.length() != numTads)
                throw new ND4JIllegalStateException("Can't symmetrically shuffle arrays with non-equal number of TADs");

            Pointer tadOffset = AtomicAllocator.getInstance().getPointer(offsets, context);

            xPointers[i] = x.address();
            xShapes[i] = xShapeInfo.address();
            tadShapes[i] = tadShapeInfo.address();
            tadOffsets[i] = tadOffset.address();
        }


        CudaDoubleDataBuffer tempX = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempShapes = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempTAD = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempOffsets = new CudaDoubleDataBuffer(arrays.size());

        AtomicAllocator.getInstance().memcpyBlocking(tempX, new LongPointer(xPointers), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempShapes, new LongPointer(xShapes), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempTAD, new LongPointer(tadShapes), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempOffsets, new LongPointer(tadOffsets), xPointers.length * 8, 0);


        if (Nd4j.dataType() == DataBuffer.Type.DOUBLE) {
            nativeOps.shuffleDouble(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        } else if (Nd4j.dataType() == DataBuffer.Type.FLOAT) {
            nativeOps.shuffleFloat(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        } else {
            // HALFs
            nativeOps.shuffleHalf(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        }


        for (int f = 0; f < arrays.size(); f++) {
            allocator.getFlowController().registerAction(context, arrays.get(f));
        }


        // just to keep reference
        shuffle.address();

        tempX.dataType();
        tempShapes.dataType();
        tempOffsets.dataType();
        tempTAD.dataType();
    }
@Override
    public void shuffle(List<INDArray> arrays, Random rnd, List<int[]> dimensions) {
        // no dimension - no shuffle
        if (dimensions == null || dimensions.size() == 0)
            throw new RuntimeException("Dimension can't be null or 0-length");

        if (arrays == null || arrays.size() == 0)
            throw new RuntimeException("No input arrays provided");

        if (dimensions.size() > 1 && arrays.size() != dimensions.size())
            throw new IllegalStateException("Number of dimensions do not match number of arrays to shuffle");

        Nd4j.getExecutioner().push();

        // first we build TAD for input array and dimensions

        AtomicAllocator allocator = AtomicAllocator.getInstance();

        CudaContext context = null;

        for (int x = 0; x < arrays.size(); x++) {
            context = allocator.getFlowController().prepareAction(arrays.get(x));
        }

        int tadLength = 1;
        for (int i = 0; i < dimensions.get(0).length; i++) {
            tadLength *= arrays.get(0).shape()[dimensions.get(0)[i]];
        }

        val numTads = arrays.get(0).length() / tadLength;

        val map = ArrayUtil.buildInterleavedVector(rnd, (int) numTads);

        val shuffle = new CudaIntDataBuffer(map);

        Pointer shuffleMap = allocator.getPointer(shuffle, context);

        PointerPointer extras = new PointerPointer(null, // not used
                        context.getOldStream(), allocator.getDeviceIdPointer());


        long[] xPointers = new long[arrays.size()];
        long[] xShapes = new long[arrays.size()];
        long[] tadShapes = new long[arrays.size()];
        long[] tadOffsets = new long[arrays.size()];

        for (int i = 0; i < arrays.size(); i++) {
            INDArray array = arrays.get(i);

            Pointer x = AtomicAllocator.getInstance().getPointer(array, context);
            Pointer xShapeInfo = AtomicAllocator.getInstance().getPointer(array.shapeInfoDataBuffer(), context);


            TADManager tadManager = Nd4j.getExecutioner().getTADManager();

            int[] dimension = dimensions.size() > 1 ? dimensions.get(i) : dimensions.get(0);

            Pair<DataBuffer, DataBuffer> tadBuffers = tadManager.getTADOnlyShapeInfo(array, dimension);

//            log.info("Original shape: {}; dimension: {}; TAD shape: {}", array.shapeInfoDataBuffer().asInt(), dimension, tadBuffers.getFirst().asInt());

            Pointer tadShapeInfo = AtomicAllocator.getInstance().getPointer(tadBuffers.getFirst(), context);

            DataBuffer offsets = tadBuffers.getSecond();

            if (offsets.length() != numTads)
                throw new ND4JIllegalStateException("Can't symmetrically shuffle arrays with non-equal number of TADs");

            Pointer tadOffset = AtomicAllocator.getInstance().getPointer(offsets, context);

            xPointers[i] = x.address();
            xShapes[i] = xShapeInfo.address();
            tadShapes[i] = tadShapeInfo.address();
            tadOffsets[i] = tadOffset.address();
        }


        CudaDoubleDataBuffer tempX = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempShapes = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempTAD = new CudaDoubleDataBuffer(arrays.size());
        CudaDoubleDataBuffer tempOffsets = new CudaDoubleDataBuffer(arrays.size());

        AtomicAllocator.getInstance().memcpyBlocking(tempX, new LongPointer(xPointers), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempShapes, new LongPointer(xShapes), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempTAD, new LongPointer(tadShapes), xPointers.length * 8, 0);
        AtomicAllocator.getInstance().memcpyBlocking(tempOffsets, new LongPointer(tadOffsets), xPointers.length * 8, 0);


        if (Nd4j.dataType() == DataBuffer.Type.DOUBLE) {
            nativeOps.shuffleDouble(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        } else if (Nd4j.dataType() == DataBuffer.Type.FLOAT) {
            nativeOps.shuffleFloat(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        } else {
            // HALFs
            nativeOps.shuffleHalf(extras, new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)),
                            new PointerPointer(allocator.getPointer(tempX, context)),
                            new PointerPointer(allocator.getPointer(tempShapes, context)), arrays.size(),
                            (IntPointer) shuffleMap, new PointerPointer(allocator.getPointer(tempTAD, context)),
                            new PointerPointer(allocator.getPointer(tempOffsets, context)));
        }


        for (int f = 0; f < arrays.size(); f++) {
            allocator.getFlowController().registerAction(context, arrays.get(f));
        }


        // just to keep reference
        shuffle.address();

        tempX.dataType();
        tempShapes.dataType();
        tempOffsets.dataType();
        tempTAD.dataType();
    }
@ Override public void shuffle ( List < INDArray > arrays , Random rnd , List < int [ ] > dimensions ) { // no dimension - no shuffle if ( dimensions == null || dimensions . size ( ) == 0 ) throw new RuntimeException ( "Dimension can't be null or 0-length" ) ; if ( arrays == null || arrays . size ( ) == 0 ) throw new RuntimeException ( "No input arrays provided" ) ; if ( dimensions . size ( ) > 1 && arrays . size ( ) != dimensions . size ( ) ) throw new IllegalStateException ( "Number of dimensions do not match number of arrays to shuffle" ) ; Nd4j . getExecutioner ( ) . push ( ) ; // first we build TAD for input array and dimensions AtomicAllocator allocator = AtomicAllocator . getInstance ( ) ; CudaContext context = null ; for ( int x = 0 ; x < arrays . size ( ) ; x ++ ) { context = allocator . getFlowController ( ) . prepareAction ( arrays . get ( x ) ) ; } int tadLength = 1 ; for ( int i = 0 ; i < dimensions . get ( 0 ) . length ; i ++ ) { tadLength *= arrays . get ( 0 ) . shape ( ) [ dimensions . get ( 0 ) [ i ] ] ; } val numTads = arrays . get ( 0 ) . length ( ) / tadLength ; val map = ArrayUtil . buildInterleavedVector ( rnd , ( int ) numTads ) ; val shuffle = new CudaIntDataBuffer ( map ) ; Pointer shuffleMap = allocator . getPointer ( shuffle , context ) ; PointerPointer extras = new PointerPointer ( null , // not used context . getOldStream ( ) , allocator . getDeviceIdPointer ( ) ) ; long [ ] xPointers = new long [ arrays . size ( ) ] ; long [ ] xShapes = new long [ arrays . size ( ) ] ; long [ ] tadShapes = new long [ arrays . size ( ) ] ; long [ ] tadOffsets = new long [ arrays . size ( ) ] ; for ( int i = 0 ; i < arrays . size ( ) ; i ++ ) { INDArray array = arrays . get ( i ) ; Pointer x = AtomicAllocator . getInstance ( ) . getPointer ( array , context ) ; Pointer xShapeInfo = AtomicAllocator . getInstance ( ) . getPointer ( array . shapeInfoDataBuffer ( ) , context ) ; TADManager tadManager = Nd4j . getExecutioner ( ) . getTADManager ( ) ; int [ ] dimension = dimensions . size ( ) > 1 ? dimensions . get ( i ) : dimensions . get ( 0 ) ; Pair < DataBuffer , DataBuffer > tadBuffers = tadManager . getTADOnlyShapeInfo ( array , dimension ) ; //            log.info("Original shape: {}; dimension: {}; TAD shape: {}", array.shapeInfoDataBuffer().asInt(), dimension, tadBuffers.getFirst().asInt()); Pointer tadShapeInfo = AtomicAllocator . getInstance ( ) . getPointer ( tadBuffers . getFirst ( ) , context ) ; DataBuffer offsets = tadBuffers . getSecond ( ) ; if ( offsets . length ( ) != numTads ) throw new ND4JIllegalStateException ( "Can't symmetrically shuffle arrays with non-equal number of TADs" ) ; Pointer tadOffset = AtomicAllocator . getInstance ( ) . getPointer ( offsets , context ) ; xPointers [ i ] = x . address ( ) ; xShapes [ i ] = xShapeInfo . address ( ) ; tadShapes [ i ] = tadShapeInfo . address ( ) ; tadOffsets [ i ] = tadOffset . address ( ) ; } CudaDoubleDataBuffer tempX = new CudaDoubleDataBuffer ( arrays . size ( ) ) ; CudaDoubleDataBuffer tempShapes = new CudaDoubleDataBuffer ( arrays . size ( ) ) ; CudaDoubleDataBuffer tempTAD = new CudaDoubleDataBuffer ( arrays . size ( ) ) ; CudaDoubleDataBuffer tempOffsets = new CudaDoubleDataBuffer ( arrays . size ( ) ) ; AtomicAllocator . getInstance ( ) . memcpyBlocking ( tempX , new LongPointer ( xPointers ) , xPointers . length * 8 , 0 ) ; AtomicAllocator . getInstance ( ) . memcpyBlocking ( tempShapes , new LongPointer ( xShapes ) , xPointers . length * 8 , 0 ) ; AtomicAllocator . getInstance ( ) . memcpyBlocking ( tempTAD , new LongPointer ( tadShapes ) , xPointers . length * 8 , 0 ) ; AtomicAllocator . getInstance ( ) . memcpyBlocking ( tempOffsets , new LongPointer ( tadOffsets ) , xPointers . length * 8 , 0 ) ; if ( Nd4j . dataType ( ) == DataBuffer . Type . DOUBLE ) { nativeOps . shuffleDouble ( extras , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , arrays . size ( ) , ( IntPointer ) shuffleMap , new PointerPointer ( allocator . getPointer ( tempTAD , context ) ) , new PointerPointer ( allocator . getPointer ( tempOffsets , context ) ) ) ; } else if ( Nd4j . dataType ( ) == DataBuffer . Type . FLOAT ) { nativeOps . shuffleFloat ( extras , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , arrays . size ( ) , ( IntPointer ) shuffleMap , new PointerPointer ( allocator . getPointer ( tempTAD , context ) ) , new PointerPointer ( allocator . getPointer ( tempOffsets , context ) ) ) ; } else { // HALFs nativeOps . shuffleHalf ( extras , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , new PointerPointer ( allocator . getPointer ( tempX , context ) ) , new PointerPointer ( allocator . getPointer ( tempShapes , context ) ) , arrays . size ( ) , ( IntPointer ) shuffleMap , new PointerPointer ( allocator . getPointer ( tempTAD , context ) ) , new PointerPointer ( allocator . getPointer ( tempOffsets , context ) ) ) ; } for ( int f = 0 ; f < arrays . size ( ) ; f ++ ) { allocator . getFlowController ( ) . registerAction ( context , arrays . get ( f ) ) ; } // just to keep reference shuffle . address ( ) ; tempX . dataType ( ) ; tempShapes . dataType ( ) ; tempOffsets . dataType ( ) ; tempTAD . dataType ( ) ; }
Symmetric in place shuffle of an ndarray
along a specified set of dimensions. Each array in list should have it's own dimension at the same index of dimensions array

@param arrays      the ndarrays to shuffle
@param dimensions the dimensions to do the shuffle
@return
Symmetric in place shuffle of an ndarray along a specified set of dimensions . Each array in list should have it s own dimension at the same index of dimensions array