license import { GPGPUProgram } from './gpgpu_math' ; import { UniformType } from './shader_compiler' ; export class MultinomialProgram implements GPGPUProgram { variableNames = [ 'probs' ] ; outputShape : number [ ] ; userCode : string ; customUniforms = [ { name : 'seed' , type : 'float' as UniformType } ] ; constructor ( batchSize , numOutcomes , numSamples ) { this . outputShape = [ batchSize , numSamples ] ; this . userCode = `
      void main() {
        ivec2 coords = getOutputCoords();
        int batch = coords[0];

        float r = random(seed);
        float cdf = 0.0;

        for (int i = 0; i < ${ numOutcomes - 1 }; i++) {
          cdf += getProbs(batch, i);

          if (r < cdf) {
            setOutput(float(i));
            return;
          }
        }

        // If no other event happened, last event happened.
        setOutput(float(${ numOutcomes - 1 }));
      }
    ` ; } }