license import { env , KernelConfig , KernelFunc , Slice , slice_util , SliceAttrs , SliceInputs , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { sliceImplCPU } from '../kernel_utils/shared' ; import { SliceProgram } from '../slice_gpu' ; import { SlicePackedProgram } from '../slice_packed_gpu' ; function shallowSlice ( x , begin , size , backend ) { const xTexData = backend . texData . get ( x . dataId ) ; const t = backend . makeTensorInfo ( size , x . dtype ) ; const newTexData = backend . texData . get ( t . dataId ) ; Object . assign ( newTexData , xTexData ) ; newTexData . refCount = 1 ; newTexData . shape = size ; newTexData . dtype = x . dtype ; let flatOffset = slice_util . computeFlatOffset ( begin , util . computeStrides ( x . shape ) ) ; if ( xTexData . slice ) { flatOffset += xTexData . slice . flatOffset ; } newTexData . slice = { flatOffset , origDataId : xTexData . slice && xTexData . slice . origDataId || x . dataId } ; const refCount = backend . dataRefCount . get ( newTexData . slice . origDataId ) || 1 ; backend . dataRefCount . set ( newTexData . slice . origDataId , refCount + 1 ) ; return t ; } export function slice ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { begin , size } = attrs ; const [ $begin , $size ] = slice_util . parseSliceParams ( x , begin , size ) ; slice_util . assertParamsValid ( x , $begin , $size ) ; if ( util . sizeFromShape ( $size ) === 0 ) { return backend . makeTensorInfo ( $size , x . dtype , [ ] ) ; } if ( backend . shouldExecuteOnCPU ( [ x ] ) || x . dtype === 'string' ) { const xTexData = backend . texData . get ( x . dataId ) ; const outValues = sliceImplCPU ( xTexData . values as TypedArray , $begin , $size , x . shape , x . dtype ) ; return backend . makeTensorInfo ( $size , x . dtype , outValues ) ; } const { isPacked } = backend . texData . get ( x . dataId ) ; const isContinous = slice_util . isSliceContinous ( x . shape , $begin , $size ) ; if ( isPacked || ! isContinous ) { const program = env ( ) . getBool ( 'WEBGL_PACK_ARRAY_OPERATIONS' ) ? new SlicePackedProgram ( $size ) : new SliceProgram ( $size ) ; const customValues = [ $begin ] ; return backend . runWebGLProgram ( program , [ x ] , x . dtype , customValues ) ; } backend . uploadToGPU ( x . dataId ) ; return shallowSlice ( x , $begin , $size , backend ) ; } export const sliceConfig = { kernelName : Slice , backendName : 'webgl' , kernelFunc : slice as { } as KernelFunc } ;