@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
63 lines • 2.45 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
var shader_compiler_1 = require("./shader_compiler");
var SliceProgram = (function () {
function SliceProgram(destSize) {
this.variableNames = ['source'];
this.outputShape = destSize;
this.rank = destSize.length;
var dtype = shader_compiler_1.getCoordsDataType(this.rank);
var sourceCoords = getCoords(this.rank);
this.userCode = "\n uniform " + dtype + " start;\n\n void main() {\n " + dtype + " sourceLoc = start + getOutputCoords();\n setOutput(getSource(" + sourceCoords + "));\n }\n ";
}
SliceProgram.prototype.getCustomSetupFunc = function (start) {
var _this = this;
if (start.length !== this.rank) {
throw Error("The rank (" + this.rank + ") of the program must match the " +
("length of start (" + start.length + ")"));
}
return function (gpgpu, webGLProgram) {
if (_this.startLoc == null) {
_this.startLoc = gpgpu.getUniformLocationNoThrow(webGLProgram, 'start');
if (_this.startLoc == null) {
return;
}
}
if (_this.rank === 1) {
gpgpu.gl.uniform1i(_this.startLoc, start[0]);
}
else if (_this.rank === 2) {
gpgpu.gl.uniform2i(_this.startLoc, start[0], start[1]);
}
else if (_this.rank === 3) {
gpgpu.gl.uniform3i(_this.startLoc, start[0], start[1], start[2]);
}
else if (_this.rank === 4) {
gpgpu.gl.uniform4i(_this.startLoc, start[0], start[1], start[2], start[3]);
}
else {
throw Error("Slicing for rank " + _this.rank + " is not yet supported");
}
};
};
return SliceProgram;
}());
exports.SliceProgram = SliceProgram;
function getCoords(rank) {
if (rank === 1) {
return 'sourceLoc';
}
else if (rank === 2) {
return 'sourceLoc.x, sourceLoc.y';
}
else if (rank === 3) {
return 'sourceLoc.x, sourceLoc.y, sourceLoc.z';
}
else if (rank === 4) {
return 'sourceLoc.x, sourceLoc.y, sourceLoc.z, sourceLoc.w';
}
else {
throw Error("Slicing for rank " + rank + " is not yet supported");
}
}
//# sourceMappingURL=slice_gpu.js.map