import { getCoordsDataType } from './shader_compiler'; export class GatherNDProgram { constructor(sliceDim, strides, shape, paramsShape) { this.sliceDim = sliceDim; this.strides = strides; this.paramsShape = paramsShape; this.variableNames = ['x', 'indices']; this.outputShape = shape; const dtype = getCoordsDataType(shape.length); let mainLoop = ` int index;`; for (let j = 0; j < this.sliceDim; j++) { mainLoop += ` index = round(getIndices(coords[0], ${j})); out_of_bounds = out_of_bounds || index < 0; out_of_bounds = out_of_bounds || index >= ${this.paramsShape[j]}; flattenIndex += index * ${this.strides[j]};`; } this.userCode = ` void main() { ${dtype} coords = getOutputCoords(); int flattenIndex = 0; bool out_of_bounds = false; ${mainLoop} setOutput(out_of_bounds ? 0.0 : getX(flattenIndex, coords[1])); } `; } } //# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiZ2F0aGVyX25kX2dwdS5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13ZWJnbC9zcmMvZ2F0aGVyX25kX2dwdS50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiQUFpQkEsT0FBTyxFQUFDLGlCQUFpQixFQUFDLE1BQU0sbUJBQW1CLENBQUM7QUFFcEQsTUFBTSxPQUFPLGVBQWU7SUFJMUIsWUFDWSxRQUFnQixFQUFVLE9BQWlCLEVBQUUsS0FBZSxFQUM1RCxXQUFxQjtRQURyQixhQUFRLEdBQVIsUUFBUSxDQUFRO1FBQVUsWUFBTyxHQUFQLE9BQU8sQ0FBVTtRQUMzQyxnQkFBVyxHQUFYLFdBQVcsQ0FBVTtRQUxqQyxrQkFBYSxHQUFHLENBQUMsR0FBRyxFQUFFLFNBQVMsQ0FBQyxDQUFDO1FBTS9CLElBQUksQ0FBQyxXQUFXLEdBQUcsS0FBSyxDQUFDO1FBQ3pCLE1BQU0sS0FBSyxHQUFHLGlCQUFpQixDQUFDLEtBQUssQ0FBQyxNQUFNLENBQUMsQ0FBQztRQUU5QyxJQUFJLFFBQVEsR0FBRztlQUNKLENBQUM7UUFDWixLQUFLLElBQUksQ0FBQyxHQUFHLENBQUMsRUFBRSxDQUFDLEdBQUcsSUFBSSxDQUFDLFFBQVEsRUFBRSxDQUFDLEVBQUUsRUFBRTtZQUN0QyxRQUFRLElBQUk7Z0RBQzhCLENBQUM7O3NEQUVLLElBQUksQ0FBQyxXQUFXLENBQUMsQ0FBQyxDQUFDO29DQUNyQyxJQUFJLENBQUMsT0FBTyxDQUFDLENBQUMsQ0FBQyxHQUFHLENBQUM7U0FDbEQ7UUFFRCxJQUFJLENBQUMsUUFBUSxHQUFHOztZQUVSLEtBQUs7Ozs7WUFJTCxRQUFROzs7O09BSWIsQ0FBQztJQUNOLENBQUM7Q0FDRiIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDE4IEdvb2dsZSBMTEMuIEFsbCBSaWdodHMgUmVzZXJ2ZWQuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cbmltcG9ydCB7R1BHUFVQcm9ncmFtfSBmcm9tICcuL2dwZ3B1X21hdGgnO1xuaW1wb3J0IHtnZXRDb29yZHNEYXRhVHlwZX0gZnJvbSAnLi9zaGFkZXJfY29tcGlsZXInO1xuXG5leHBvcnQgY2xhc3MgR2F0aGVyTkRQcm9ncmFtIGltcGxlbWVudHMgR1BHUFVQcm9ncmFtIHtcbiAgdmFyaWFibGVOYW1lcyA9IFsneCcsICdpbmRpY2VzJ107XG4gIG91dHB1dFNoYXBlOiBudW1iZXJbXTtcbiAgdXNlckNvZGU6IHN0cmluZztcbiAgY29uc3RydWN0b3IoXG4gICAgICBwcml2YXRlIHNsaWNlRGltOiBudW1iZXIsIHByaXZhdGUgc3RyaWRlczogbnVtYmVyW10sIHNoYXBlOiBudW1iZXJbXSxcbiAgICAgIHByaXZhdGUgcGFyYW1zU2hhcGU6IG51bWJlcltdKSB7XG4gICAgdGhpcy5vdXRwdXRTaGFwZSA9IHNoYXBlO1xuICAgIGNvbnN0IGR0eXBlID0gZ2V0Q29vcmRzRGF0YVR5cGUoc2hhcGUubGVuZ3RoKTtcblxuICAgIGxldCBtYWluTG9vcCA9IGBcbiAgICBpbnQgaW5kZXg7YDtcbiAgICBmb3IgKGxldCBqID0gMDsgaiA8IHRoaXMuc2xpY2VEaW07IGorKykge1xuICAgICAgbWFpbkxvb3AgKz0gYFxuICAgICAgICAgIGluZGV4ID0gcm91bmQoZ2V0SW5kaWNlcyhjb29yZHNbMF0sICR7an0pKTtcbiAgICAgICAgICBvdXRfb2ZfYm91bmRzID0gb3V0X29mX2JvdW5kcyB8fCBpbmRleCA8IDA7XG4gICAgICAgICAgb3V0X29mX2JvdW5kcyA9IG91dF9vZl9ib3VuZHMgfHwgaW5kZXggPj0gJHt0aGlzLnBhcmFtc1NoYXBlW2pdfTtcbiAgICAgICAgICBmbGF0dGVuSW5kZXggKz0gaW5kZXggKiAke3RoaXMuc3RyaWRlc1tqXX07YDtcbiAgICB9XG5cbiAgICB0aGlzLnVzZXJDb2RlID0gYFxuICAgICAgICAgdm9pZCBtYWluKCkge1xuICAgICAgICAgICR7ZHR5cGV9IGNvb3JkcyA9IGdldE91dHB1dENvb3JkcygpO1xuICAgICAgICAgIGludCBmbGF0dGVuSW5kZXggPSAwO1xuICAgICAgICAgIGJvb2wgb3V0X29mX2JvdW5kcyA9IGZhbHNlO1xuXG4gICAgICAgICAgJHttYWluTG9vcH1cblxuICAgICAgICAgIHNldE91dHB1dChvdXRfb2ZfYm91bmRzID8gMC4wIDogZ2V0WChmbGF0dGVuSW5kZXgsIGNvb3Jkc1sxXSkpO1xuICAgICAgICB9XG4gICAgICBgO1xuICB9XG59XG4iXX0=