license import { buffer , KernelConfig , KernelFunc , Rank , slice_util , StridedSlice , StridedSliceAttrs , StridedSliceInputs , TensorBuffer , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { stridedSliceImplCPU } from '../kernel_utils/shared' ; import { StridedSliceProgram } from '../strided_slice_gpu' ; import { reshape } from './Reshape' ; import { slice } from './Slice' ; export function stridedSlice ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { begin , end , strides , beginMask , endMask , ellipsisMask , newAxisMask , shrinkAxisMask } = attrs ; const { finalShapeSparse , finalShape , isIdentity , sliceDim0 , isSimpleSlice , begin : $begin , end : $end , strides : $strides } = slice_util . sliceInfo ( x . shape , begin , end , strides , beginMask , endMask , ellipsisMask , newAxisMask , shrinkAxisMask ) ; let result ; if ( isIdentity ) { result = reshape ( { inputs : { x } , backend , attrs : { shape : finalShape } } ) ; } else if ( sliceDim0 || isSimpleSlice ) { util . assert ( x . shape . length >= 1 , ( ) => `Input must have rank at least 1, got: ${ x . shape . length }` ) ; const size = slice_util . computeOutShape ( $begin , $end , $strides ) ; const sliced = slice ( { inputs : { x } , backend , attrs : { begin : $begin , size } } ) ; result = reshape ( { inputs : { x : sliced } , backend , attrs : { shape : finalShape } } ) ; backend . disposeIntermediateTensorInfo ( sliced ) ; } else { const shouldExecuteOnCPU = backend . shouldExecuteOnCPU ( [ x ] ) ; if ( shouldExecuteOnCPU ) { const values = backend . readSync ( x . dataId ) as TypedArray ; const xBuf = buffer ( x . shape , x . dtype , values ) as TensorBuffer < Rank > ; const resultValues = stridedSliceImplCPU ( finalShapeSparse , xBuf , $strides , $begin ) ; result = backend . makeTensorInfo ( finalShape , x . dtype , resultValues . values ) ; } else { const program = new StridedSliceProgram ( $begin , $strides , finalShapeSparse ) ; result = backend . runWebGLProgram ( program , [ x ] , x . dtype ) ; } } const resultReshaped = reshape ( { inputs : { x : result } , backend , attrs : { shape : finalShape } } ) ; backend . disposeIntermediateTensorInfo ( result ) ; return resultReshaped ; } export const stridedSliceConfig = { kernelName : StridedSlice , backendName : 'webgl' , kernelFunc : stridedSlice as { } as KernelFunc } ;