license import { TensorInfo } from '../kernel_registry' ; import * as util from '../util' ; const NEW_AXIS = - 2 ; const SHRINK_AXIS = - 1 ; interface StridedSliceSparseSpec { dims : number ; numAddAxisAfterEllipsis : number ; begin : number [ ] ; end : number [ ] ; strides : number [ ] ; beginMask : number ; endMask : number ; ellipsisMask : number ; newAxisMask : number ; shrinkAxisMask : number ; } interface StridedSliceDenseSpec { dims : number ; beginMask ? : number ; endMask ? : number ; beginValid : boolean ; endValid : boolean ; begin ? : number [ ] ; end ? : number [ ] ; strides ? : number [ ] ; finalShapeGatherIndices ? : number [ ] ; finalShapeGatherIndicesSparse ? : number [ ] ; inputShapeGatherIndicesSparse ? : number [ ] ; shrinkAxisMask ? : number ; } export type SliceInfo = { finalShapeSparse : number [ ] ; finalShape : number [ ] ; isIdentity : boolean ; sliceDim0 : boolean ; isSimpleSlice : boolean ; begin : number [ ] ; end : number [ ] ; strides : number [ ] ; } ; export function assertParamsValid ( input , begin , size ) { const inputRank = input . shape . length ; util . assert ( inputRank === begin . length , ( ) => `Error in slice${ inputRank }D: Length of begin ${ begin } must ` + `match the rank of the array (${ inputRank }).` ) ; util . assert ( inputRank === size . length , ( ) => `Error in slice${ inputRank }D: Length of size ${ size } must ` + `match the rank of the array (${ inputRank }).` ) ; for ( let i = 0 ; i < inputRank ; ++ i ) { util . assert ( begin [ i ] + size [ i ] <= input . shape [ i ] , ( ) => `Error in slice${ inputRank }D: begin[${ i }] + size[${ i }] ` + `(${ begin [ i ] + size [ i ] }) would overflow input.shape[${ i }] (${ input . shape [ i ] })` ) ; } } /** Converts a binary mask to an array of axes. Used in stridedSlice(). */ export function maskToAxes ( mask ) { const axes = [ ] ; let axis = 0 ; while ( mask > 0 ) { if ( mask & 1 ) { axes . push ( axis ) ; } mask /= 2 ; axis ++ ; } return axes ; } /** Computes the output shape given the strided slice params. */ export function computeOutShape ( begin , end , strides ) { const size = [ ] ; for ( let axis = 0 ; axis < begin . length ; axis ++ ) { size [ axis ] = Math . ceil ( ( end [ axis ] - begin [ axis ] ) / strides [ axis ] ) ; } return size ; } export function stridesWithElidedDims ( strides , ellipsisInsertionIndex , numElidedAxes , inputShape ) { const newStrides = [ ... strides ] ; for ( let i = newStrides . length ; i < inputShape . length ; i ++ ) { newStrides . push ( 1 ) ; } for ( let i = 0 ; i < numElidedAxes ; i ++ ) { if ( i === 0 ) { newStrides [ ellipsisInsertionIndex ] = 1 ; } else { newStrides . splice ( ellipsisInsertionIndex , 0 , 1 ) ; newStrides . pop ( ) ; } } return newStrides ; } function unnormalizeAxis ( ellipsisInsertionIndex , numElidedAxes , normalizedAxis ) { if ( normalizedAxis <= ellipsisInsertionIndex ) { return normalizedAxis ; } return normalizedAxis - ( numElidedAxes - 1 ) ; } function getElidedAxes ( numElidedAxes , ellipsisInsertionIndex ) { const elidedAxes = [ ] ; for ( let i = 0 ; i < numElidedAxes ; i ++ ) { elidedAxes . push ( ellipsisInsertionIndex + i ) ; } return elidedAxes ; } export function getNormalizedAxes ( inputShape , ellipsisAxes , numInterpolatedAxes , begin , end , strides , beginMask , endMask , ellipsisMask ) { const inputRank = inputShape . length ; let normalizedBegin = new Array ( inputRank ) , normalizedEnd = new Array ( inputRank ) , normalizedStrides = new Array ( inputRank ) ; if ( ellipsisAxes . length && numInterpolatedAxes > 0 ) { const fullIndex = ellipsisAxes [ 0 ] ; const numElidedAxes = numInterpolatedAxes + 1 ; normalizedBegin = startIndicesWithElidedDims ( beginMask , fullIndex , numElidedAxes , begin , inputShape ) ; normalizedEnd = stopIndicesWithElidedDims ( endMask , fullIndex , numElidedAxes , end , inputShape ) ; normalizedStrides = stridesWithElidedDims ( strides , fullIndex , numElidedAxes , inputShape ) ; } else { for ( let axis = 0 ; axis < inputRank ; axis ++ ) { normalizedBegin [ axis ] = startForAxis ( beginMask , begin , strides , inputShape , axis , ellipsisMask ) ; normalizedEnd [ axis ] = stopForAxis ( endMask , end , strides , inputShape , axis , ellipsisMask ) ; normalizedStrides [ axis ] = stridesForAxis ( strides , axis , ellipsisMask ) ; } } return { begin : normalizedBegin , end : normalizedEnd , strides : normalizedStrides } ; } export function startIndicesWithElidedDims ( beginMask , ellipsisInsertionIndex , numElidedAxes , originalBegin , inputShape ) { const newIndices = [ ... inputShape ] ; const elidedAxes = getElidedAxes ( numElidedAxes , ellipsisInsertionIndex ) ; for ( let axis = 0 ; axis < newIndices . length ; axis ++ ) { if ( elidedAxes . indexOf ( axis ) > - 1 ) { newIndices [ axis ] = 0 ; } else { const originalAxis = unnormalizeAxis ( ellipsisInsertionIndex , numElidedAxes , axis ) ; let originalValue = originalBegin [ originalAxis ] ; if ( beginMask & 1 << originalAxis ) { originalValue = 0 ; } newIndices [ axis ] = originalValue ; } } return newIndices ; } export function stopIndicesWithElidedDims ( endMask , ellipsisInsertionIndex , numElidedAxes , originalEnd , inputShape ) { const newIndices = [ ... inputShape ] ; const elidedAxes = getElidedAxes ( numElidedAxes , ellipsisInsertionIndex ) ; for ( let axis = 0 ; axis < newIndices . length ; axis ++ ) { if ( elidedAxes . indexOf ( axis ) > - 1 ) { newIndices [ axis ] = Number . MAX_SAFE_INTEGER ; } else { const originalAxis = unnormalizeAxis ( ellipsisInsertionIndex , numElidedAxes , axis ) ; let originalValue = originalEnd [ originalAxis ] ; if ( endMask & 1 << originalAxis ) { originalValue = Number . MAX_SAFE_INTEGER ; } newIndices [ axis ] = originalValue ; } } for ( let i = 0 ; i < newIndices . length ; i ++ ) { const axisSize = inputShape [ i ] ; if ( newIndices [ i ] < 0 ) { newIndices [ i ] += axisSize ; } newIndices [ i ] = util . clamp ( 0 , newIndices [ i ] , inputShape [ i ] ) ; } return newIndices ; } export function stridesForAxis ( strides , axis , ellipsisMask ) { let stride = strides [ axis ] ; if ( ellipsisMask & ( 1 << axis ) || stride == null ) { stride = 1 ; } return stride ; } export function startForAxis ( beginMask , startIndices , strides , inputShape , axis , ellipsisMask ) { let start = startIndices [ axis ] ; const stride = strides [ axis ] || 1 ; if ( beginMask & 1 << axis || ellipsisMask & 1 << axis || start == null ) { if ( stride > 0 ) { start = Number . MIN_SAFE_INTEGER ; } else { start = Number . MAX_SAFE_INTEGER ; } } const axisSize = inputShape [ axis ] ; if ( start < 0 ) { start += axisSize ; } start = util . clamp ( 0 , start , axisSize - 1 ) ; return start ; } export function stopForAxis ( endMask , stopIndices , strides , inputShape , axis , ellipsisMask ) { let stop = stopIndices [ axis ] ; const stride = strides [ axis ] || 1 ; if ( endMask & ( 1 << axis ) || ellipsisMask & ( 1 << axis ) || stop == null ) { if ( stride > 0 ) { stop = Number . MAX_SAFE_INTEGER ; } else { stop = Number . MIN_SAFE_INTEGER ; } } const axisSize = inputShape [ axis ] ; if ( stop < 0 ) { stop += axisSize ; } if ( stride > 0 ) { stop = util . clamp ( 0 , stop , axisSize ) ; } else { stop = util . clamp ( - 1 , stop , axisSize - 1 ) ; } return stop ; } /**
 * Returns true if the slice occupies a continous set of elements in the
 * 'flat' space.
 */ export function isSliceContinous ( shape , begin , size ) { let firstNonOneAxis = size . length ; for ( let i = 0 ; i < size . length ; i ++ ) { if ( size [ i ] > 1 ) { firstNonOneAxis = i ; break ; } } for ( let i = firstNonOneAxis + 1 ; i < size . length ; i ++ ) { if ( begin [ i ] > 0 || size [ i ] !== shape [ i ] ) { return false ; } } return true ; } export function computeFlatOffset ( begin , strides ) { let flatOffset = begin . length > 0 ? begin [ begin . length - 1 ] : 1 ; for ( let i = 0 ; i < begin . length - 1 ; i ++ ) { flatOffset += begin [ i ] * strides [ i ] ; } return flatOffset ; } export function parseSliceParams ( x , begin , size ? ) { let begin_ ; const xRank = x . shape . length ; if ( typeof begin === 'number' ) { begin_ = [ begin , ... new Array ( xRank - 1 ) . fill ( 0 ) ] ; } else if ( begin . length < xRank ) { begin_ = begin . concat ( new Array ( xRank - begin . length ) . fill ( 0 ) ) ; } else { begin_ = begin . slice ( ) ; } begin_ . forEach ( d => { util . assert ( d !== - 1 , ( ) => 'slice() does not support negative begin indexing.' ) ; } ) ; let size_ ; if ( size == null ) { size_ = new Array ( xRank ) . fill ( - 1 ) ; } else if ( typeof size === 'number' ) { size_ = [ size , ... new Array ( xRank - 1 ) . fill ( - 1 ) ] ; } else if ( size . length < xRank ) { size_ = size . concat ( new Array ( xRank - size . length ) . fill ( - 1 ) ) ; } else { size_ = size ; } size_ = size_ . map ( ( d , i ) => { if ( d >= 0 ) { return d ; } else { util . assert ( d === - 1 , ( ) => `Negative size values should be exactly -1 but got ` + `${ d } for the slice() size at index ${ i }.` ) ; return x . shape [ i ] - begin_ [ i ] ; } } ) ; return [ begin_ , size_ ] ; } export function sliceInfo ( xShape , begin , end , strides , beginMask , endMask , ellipsisMask , newAxisMask , shrinkAxisMask ) { let stridesNonNull ; if ( strides == null ) { stridesNonNull = new Array ( begin . length ) ; stridesNonNull . fill ( 1 ) ; } else { stridesNonNull = strides ; } if ( ellipsisMask != null && ( ellipsisMask & ( ellipsisMask - 1 ) ) !== 0 ) { throw new Error ( 'Multiple ellipses in slice is not allowed.' ) ; } let ellipsisSeen = false ; const sparseSpec = { dims : stridesNonNull . length , numAddAxisAfterEllipsis : 0 , begin : begin . slice ( ) , end : end . slice ( ) , strides : stridesNonNull . slice ( ) , beginMask , endMask , ellipsisMask , newAxisMask , shrinkAxisMask } ; for ( let i = 0 ; i < sparseSpec . dims ; i ++ ) { if ( ellipsisSeen && ( ( 1 << i ) & newAxisMask ) !== 0 ) { sparseSpec . numAddAxisAfterEllipsis ++ ; } if ( ( 1 << i ) & ellipsisMask ) { ellipsisSeen = true ; } } if ( ! ellipsisSeen ) { sparseSpec . ellipsisMask |= ( 1 << sparseSpec . dims ) ; sparseSpec . dims ++ ; } const denseSpec = { dims : xShape . length , beginMask : 0 , endMask : 0 , beginValid : false , endValid : false } ; buildDenseSpec ( sparseSpec , denseSpec ) ; let isIdentity = true ; let sliceDim0 = true ; let isSimpleSlice = true ; const processingShape = [ ] ; const finalShape = [ ] ; for ( let i = 0 ; i < xShape . length ; ++ i ) { if ( denseSpec . strides [ i ] === 0 ) { throw Error ( `strides[${ i }] must be non-zero` ) ; } const shrinkI = ! ! ( denseSpec . shrinkAxisMask & ( 1 << i ) ) ; const dimI = xShape [ i ] ; if ( dimI === - 1 ) { processingShape . push ( shrinkI ? 1 : - 1 ) ; continue ; } const masks = [ denseSpec . beginMask & ( 1 << i ) , denseSpec . endMask & ( 1 << i ) ] ; const validRange = [ denseSpec . strides [ i ] > 0 ? 0 : - 1 , denseSpec . strides [ i ] > 0 ? dimI : dimI - 1 ] ; if ( shrinkI && denseSpec . strides [ i ] <= 0 ) { throw Error ( 'only stride 1 allowed on non-range indexing.' ) ; } isSimpleSlice = isSimpleSlice && ( denseSpec . strides [ i ] === 1 ) ; const beginAndEndMasked = ! ! ( ( denseSpec . beginMask & ( 1 << i ) ) && ( denseSpec . endMask & ( 1 << i ) ) ) ; if ( denseSpec . beginValid && denseSpec . endValid ) { if ( shrinkI ) { const xFwd = denseSpec . begin [ i ] < 0 ? dimI + denseSpec . begin [ i ] : denseSpec . begin [ i ] ; denseSpec . begin [ i ] = xFwd ; denseSpec . end [ i ] = denseSpec . begin [ i ] + 1 ; if ( xFwd < 0 || xFwd >= dimI ) { throw Error ( `slice index ${ denseSpec . begin [ i ] } of dimension ${ i } out of bounds.` ) ; } } else { denseSpec . begin [ i ] = canonical ( denseSpec . begin [ i ] , 0 , denseSpec . strides [ i ] , dimI , masks , validRange ) ; denseSpec . end [ i ] = canonical ( denseSpec . end [ i ] , 1 , denseSpec . strides [ i ] , dimI , masks , validRange ) ; } const takeAllInDimension = denseSpec . strides [ i ] === 1 && denseSpec . begin [ i ] === 0 && denseSpec . end [ i ] === dimI ; isIdentity = isIdentity && takeAllInDimension ; sliceDim0 = sliceDim0 && ( ( i === 0 && denseSpec . strides [ i ] === 1 ) || takeAllInDimension ) ; } else { isIdentity = isIdentity && ( ( denseSpec . strides [ i ] === 1 ) && beginAndEndMasked ) ; sliceDim0 = sliceDim0 && ( ( i === 0 && denseSpec . strides [ i ] === 1 ) || beginAndEndMasked ) ; } let intervalLength ; let knownInterval = false ; if ( denseSpec . beginValid && denseSpec . endValid ) { intervalLength = denseSpec . end [ i ] - denseSpec . begin [ i ] ; knownInterval = true ; } else if ( shrinkI ) { intervalLength = 1 ; knownInterval = true ; } else if ( beginAndEndMasked ) { if ( dimI >= 0 ) { if ( denseSpec . strides [ i ] < 0 ) { intervalLength = - dimI ; } else { intervalLength = dimI ; } knownInterval = true ; } } if ( knownInterval ) { let sizeI ; if ( intervalLength === 0 || ( ( intervalLength < 0 ) !== ( denseSpec . strides [ i ] < 0 ) ) ) { sizeI = 0 ; } else { sizeI = Math . trunc ( intervalLength / denseSpec . strides [ i ] ) + ( intervalLength % denseSpec . strides [ i ] !== 0 ? 1 : 0 ) ; } processingShape . push ( sizeI ) ; } else { processingShape . push ( - 1 ) ; } } for ( let denseDim = 0 ; denseDim < denseSpec . finalShapeGatherIndices . length ; ++ denseDim ) { const gatherIndex = denseSpec . finalShapeGatherIndices [ denseDim ] ; if ( gatherIndex >= 0 ) { finalShape . push ( processingShape [ gatherIndex ] ) ; } else if ( gatherIndex === NEW_AXIS ) { finalShape . push ( 1 ) ; } } const finalShapeSparse = finalShape . filter ( ( dim , i ) => denseSpec . finalShapeGatherIndices [ i ] !== NEW_AXIS ) ; return { finalShapeSparse , finalShape , isIdentity , sliceDim0 , isSimpleSlice , begin : denseSpec . begin , end : denseSpec . end , strides : denseSpec . strides } ; } function buildDenseSpec ( sparse , dense ) { dense . beginMask = 0 ; dense . endMask = 0 ; dense . shrinkAxisMask = 0 ; let fullIndex = 0 ; dense . beginValid = sparse . begin != null ; dense . endValid = sparse . end != null ; dense . begin = new Array ( dense . dims ) ; dense . end = new Array ( dense . dims ) ; dense . strides = new Array ( dense . dims ) ; dense . finalShapeGatherIndices = [ ] ; dense . finalShapeGatherIndicesSparse = [ ] ; dense . inputShapeGatherIndicesSparse = new Array ( dense . dims ) ; for ( let i = 0 ; i < sparse . dims ; i ++ ) { if ( ( 1 << i ) & sparse . ellipsisMask ) { const nextIndex = Math . min ( dense . dims - ( sparse . dims - i ) + 1 + sparse . numAddAxisAfterEllipsis , dense . dims ) ; for ( ; fullIndex < nextIndex ; fullIndex ++ ) { dense . begin [ fullIndex ] = 0 ; dense . end [ fullIndex ] = 0 ; dense . strides [ fullIndex ] = 1 ; dense . beginMask |= ( 1 << fullIndex ) ; dense . endMask |= ( 1 << fullIndex ) ; dense . finalShapeGatherIndices . push ( fullIndex ) ; dense . finalShapeGatherIndicesSparse . push ( - 1 ) ; dense . inputShapeGatherIndicesSparse [ fullIndex ] = i ; } } else if ( ( 1 << i ) & sparse . newAxisMask ) { dense . finalShapeGatherIndices . push ( NEW_AXIS ) ; dense . finalShapeGatherIndicesSparse . push ( - 1 ) ; } else { if ( fullIndex === dense . begin . length ) { throw Error ( `Index out of range using input dim ${ fullIndex }; input ` + `has only ${ dense . dims } dims, ${ dense . begin . length }.` ) ; } if ( sparse . begin != null ) { dense . begin [ fullIndex ] = sparse . begin [ i ] ; } if ( sparse . end != null ) { dense . end [ fullIndex ] = sparse . end [ i ] ; } dense . strides [ fullIndex ] = sparse . strides [ i ] ; if ( sparse . beginMask & ( 1 << i ) ) { dense . beginMask |= ( 1 << fullIndex ) ; } if ( sparse . endMask & ( 1 << i ) ) { dense . endMask |= ( 1 << fullIndex ) ; } if ( sparse . shrinkAxisMask & ( 1 << i ) ) { dense . finalShapeGatherIndices . push ( SHRINK_AXIS ) ; dense . finalShapeGatherIndicesSparse . push ( - 1 ) ; dense . shrinkAxisMask |= ( 1 << fullIndex ) ; } else { dense . finalShapeGatherIndices . push ( fullIndex ) ; dense . finalShapeGatherIndicesSparse . push ( i ) ; } dense . inputShapeGatherIndicesSparse [ fullIndex ] = i ; fullIndex ++ ; } } } function canonical ( x , c , strideI , dimI , masks , validRange ) { if ( masks [ c ] ) { return strideI > 0 ? validRange [ c ] : validRange [ ( c + 1 ) & 1 ] ; } else { const xFwd = x < 0 ? dimI + x : x ; return xFwd < validRange [ 0 ] ? validRange [ 0 ] : xFwd > validRange [ 1 ] ? validRange [ 1 ] : xFwd ; } }