license import { GPGPUProgram } from './gpgpu_math' ; import { getCoordsDataType } from './shader_compiler' ; export class ReverseProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] ; userCode : string ; constructor ( xShape , axis ) { const rank = xShape . length ; if ( rank > 4 ) { throw new Error ( `WebGL backend: Reverse of rank-${ rank } tensor is not yet supported` ) ; } this . outputShape = xShape ; if ( rank === 1 ) { this . userCode = `
        void main() {
          int coord = getOutputCoords();
          setOutput(getX(${ xShape [ 0 ] } - coord - 1));
        }
      ` ; return ; } const getInCoord = ( i ) => { if ( axis . indexOf ( i ) !== - 1 && xShape [ i ] !== 1 ) { return `${ xShape [ i ] } - coords[${ i }] - 1` ; } return `coords[${ i }]` ; } ; const inCoords = xShape . map ( ( _ , i ) => getInCoord ( i ) ) . join ( ',' ) ; const type = getCoordsDataType ( rank ) ; this . userCode = `
      void main() {
        ${ type } coords = getOutputCoords();
        setOutput(getX(${ inCoords }));
      }
    ` ; } }