license import { KernelConfig , KernelFunc , slice_util , StridedSlice , StridedSliceAttrs , StridedSliceInputs , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; import { reshape } from './Reshape' ; import { slice } from './Slice' ; import { stridedSliceImpl } from './StridedSlice_impl' ; export function stridedSlice ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { begin , end , strides , beginMask , endMask , ellipsisMask , newAxisMask , shrinkAxisMask } = attrs ; assertNotComplex ( x , 'stridedSlice' ) ; 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 xBuf = backend . bufferSync ( x ) ; const outBuf = stridedSliceImpl ( finalShapeSparse , xBuf , $strides , $begin ) ; result = backend . makeTensorInfo ( finalShape , outBuf . dtype , outBuf . values ) ; } return result ; } export const stridedSliceConfig = { kernelName : StridedSlice , backendName : 'cpu' , kernelFunc : stridedSlice as { } as KernelFunc } ;