@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
21 lines • 1.32 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
var MultinomialProgram = (function () {
function MultinomialProgram(batchSize, numOutcomes, numSamples) {
this.variableNames = ['probs'];
this.outputShape = [batchSize, numSamples];
this.userCode = "\n uniform float seed;\n\n void main() {\n ivec2 coords = getOutputCoords();\n int batch = coords[0];\n\n float r = random(seed);\n float cdf = 0.0;\n\n for (int i = 0; i < " + (numOutcomes - 1) + "; i++) {\n cdf += getProbs(batch, i);\n\n if (r < cdf) {\n setOutput(float(i));\n return;\n }\n }\n\n // If no other event happened, last event happened.\n setOutput(float(" + (numOutcomes - 1) + "));\n }\n ";
}
MultinomialProgram.prototype.getCustomSetupFunc = function (seed) {
var _this = this;
return function (gpgpu, webGLProgram) {
if (_this.seedLoc == null) {
_this.seedLoc = gpgpu.getUniformLocation(webGLProgram, 'seed');
}
gpgpu.gl.uniform1f(_this.seedLoc, seed);
};
};
return MultinomialProgram;
}());
exports.MultinomialProgram = MultinomialProgram;
//# sourceMappingURL=multinomial_gpu.js.map