license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class GatherNDProgram implements GPGPUProgram { variableNames = [ 'x' , 'indices' ] ; outputShape : number [ ] ; userCode : string ; constructor ( private sliceDim , private strides , shape ) { this . outputShape = shape ; const stridesType = getCoordsDataType ( strides . length ) ; const dtype = getCoordsDataType ( shape . length ) ; const strideString = this . sliceDim > 1 ? 'strides[j]' : 'strides' ; this . userCode = `
        ${ stridesType } strides = ${ stridesType }(${ this . strides });
         void main() {
          ${ dtype } coords = getOutputCoords();
          int flattenIndex = 0;
          for (int j = 0; j < ${ this . sliceDim }; j++) {
            int index = round(getIndices(coords[0], j));
            flattenIndex += index * ${ strideString };
          }
          setOutput(getX(flattenIndex, coords[1]));
        }
      ` ; } }