license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class ConcatProgram implements GPGPUProgram { variableNames : string [ ] ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( shapes ) { this . outputShape = backend_util . computeOutShape ( shapes , 1 ) ; this . variableNames = shapes . map ( ( _ , i ) => `T${ i }` ) ; const offsets = new Array ( shapes . length - 1 ) ; offsets [ 0 ] = shapes [ 0 ] [ 1 ] ; for ( let i = 1 ; i < offsets . length ; i ++ ) { offsets [ i ] = offsets [ i - 1 ] + shapes [ i ] [ 1 ] ; } const snippets = [ `if (yC < ${ offsets [ 0 ] }) setOutput(getT0(yR, yC));` ] ; for ( let i = 1 ; i < offsets . length ; i ++ ) { const shift = offsets [ i - 1 ] ; snippets . push ( `else if (yC < ${ offsets [ i ] }) ` + `setOutput(getT${ i }(yR, yC-${ shift }));` ) ; } const lastIndex = offsets . length ; const lastShift = offsets [ offsets . length - 1 ] ; snippets . push ( `else setOutput(getT${ lastIndex }(yR, yC-${ lastShift }));` ) ; this . userCode = `
      void main() {
        ivec2 coords = getOutputCoords();
        int yR = coords.x;
        int yC = coords.y;

        ${ snippets . join ( '\n        ' ) }
      }
    ` ; } }