license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class ArgMinMaxProgram implements GPGPUProgram { variableNames = [ 'A' ] ; outputShape : number [ ] ; userCode : string ; constructor ( reduceInfo , op , firstPass ) { const { windowSize , batchSize , outSize } = reduceInfo ; if ( ! firstPass ) { this . variableNames . push ( 'bestIndicesA' ) ; } this . outputShape = [ batchSize , outSize ] ; const compOp = ( op === 'max' ) ? '>' : '<' ; const indexSnippet = firstPass ? 'inOffset + i;' : 'round(getBestIndicesA(batch, inOffset + i));' ; this . userCode = `
      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];
        int outIdx = coords[1];
        int inOffset = outIdx * ${ windowSize };

        int bestIndex = inOffset;
        float bestValue = getA(batch, bestIndex);

        for (int i = 0; i < ${ windowSize }; i++) {
          int inIdx = ${ indexSnippet };
          float candidate = getA(batch, inIdx);
          if (candidate ${ compOp } bestValue) {
            bestValue = candidate;
            bestIndex = inIdx;
          }
        }
        setOutput(float(bestIndex));
      }
    ` ; } }