license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class StridedSliceProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] ; userCode : string ; constructor ( begin , strides , size ) { this . outputShape = size ; const rank = size . length ; const inputDtype = getCoordsDataType ( size . length ) ; const dtype = getCoordsDataType ( size . length ) ; let newCoords = '' ; if ( rank === 1 ) { newCoords = 'coords * strides + begin' ; } else { let outputAxis = 0 ; newCoords = size . map ( ( _ , i ) => { outputAxis ++ ; return size . length === 1 ? `coords * strides[${ i }] + begin[${ i }]` : `coords[${ outputAxis - 1 }] * strides[${ i }] + begin[${ i }]` ; } ) . join ( ',' ) ; } this . userCode = `
      ${ inputDtype } begin = ${ inputDtype }(${ begin });
      ${ inputDtype } strides = ${ inputDtype }(${ strides });

      void main() {
        ${ dtype } coords = getOutputCoords();
        setOutput(getX(${ newCoords }));
      }
    ` ; } }