private void invoke ( ScalarOp op , int [ ] dimension ) { dimension = Shape . normalizeAxis ( op . x ( ) . rank ( ) , dimension ) ; // do tad magic /** * Returns the {@link Shape#createShapeInformation(int[], int[], int, int, char)} * and the associated offsets for each {@link INDArray#tensorAlongDimension(int, int...)} * The first item is the shape information. The second one is the offsets. */ Pair < DataBuffer , DataBuffer > tadBuffers = tadManager . getTADOnlyShapeInfo ( op . x ( ) , dimension ) ; Pointer hostTadShapeInfo = tadBuffers . getFirst ( ) . addressPointer ( ) ; Pointer hostTadOffsets = tadBuffers . getSecond ( ) . addressPointer ( ) ; Pointer devTadShapeInfoZ = null ; Pointer devTadOffsetsZ = null ; /** * Returns the {@link Shape#createShapeInformation(int[], int[], int, int, char)} * and the associated offsets for each {@link INDArray#tensorAlongDimension(int, int...)} * The first item is the shape information. The second one is the offsets. * * Note that this is the *result* TAD information. An op is always input (x) and output (z) * for result. * This is for assigning the result to of the operation along * the proper dimension. */ Pair < DataBuffer , DataBuffer > tadBuffersZ = tadManager . getTADOnlyShapeInfo ( op . z ( ) , dimension ) ; devTadShapeInfoZ = tadBuffersZ . getFirst ( ) . addressPointer ( ) ; devTadOffsetsZ = tadBuffersZ . getSecond ( ) . addressPointer ( ) ; if ( extraz . get ( ) == null ) extraz . set ( new PointerPointer ( 32 ) ) ; PointerPointer dummy = extraz . get ( ) . put ( hostTadShapeInfo , hostTadOffsets , devTadShapeInfoZ , devTadOffsetsZ ) ; if ( op . x ( ) . data ( ) . dataType ( ) == DataBuffer . Type . FLOAT ) { loop . execScalarFloat ( dummy , op . opNum ( ) , ( FloatPointer ) op . x ( ) . data ( ) . addressPointer ( ) , ( LongPointer ) op . x ( ) . shapeInfoDataBuffer ( ) . addressPointer ( ) , ( FloatPointer ) op . z ( ) . data ( ) . addressPointer ( ) , ( LongPointer ) op . z ( ) . shapeInfoDataBuffer ( ) . addressPointer ( ) , ( FloatPointer ) op . y ( ) . data ( ) . addressPointer ( ) , ( FloatPointer ) getPointerForExtraArgs ( op ) , ( IntPointer ) Nd4j . getConstantHandler ( ) . getConstantBuffer ( dimension ) . addressPointer ( ) , dimension . length ) ; } else if ( op . x ( ) . data ( ) . dataType ( ) == DataBuffer . Type . DOUBLE ) { loop . execScalarDouble ( dummy , op . opNum ( ) , ( DoublePointer ) op . x ( ) . data ( ) . addressPointer ( ) , ( LongPointer ) op . x ( ) . shapeInfoDataBuffer ( ) . addressPointer ( ) , ( DoublePointer ) op . z ( ) . data ( ) . addressPointer ( ) , ( LongPointer ) op . z ( ) . shapeInfoDataBuffer ( ) . addressPointer ( ) , ( DoublePointer ) op . y ( ) . data ( ) . addressPointer ( ) , ( DoublePointer ) getPointerForExtraArgs ( op ) , ( IntPointer ) Nd4j . getConstantHandler ( ) . getConstantBuffer ( dimension ) . addressPointer ( ) , dimension . length ) ; } }
ScalarOp along dimension