license import { KernelConfig , KernelFunc , Multinomial , MultinomialAttrs , MultinomialInputs , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { MultinomialProgram } from '../multinomial_gpu' ; import { softmax } from './Softmax' ; export function multinomial ( args ) { const { inputs , backend , attrs } = args ; const { logits } = inputs ; const { numSamples , seed , normalized } = attrs ; const probs = normalized ? logits : softmax ( { inputs : { logits } , backend , attrs : { dim : logits . shape . length - 1 } } ) ; const batchSize = probs . shape [ 0 ] ; const numOutcomes = probs . shape [ 1 ] ; const program = new MultinomialProgram ( batchSize , numOutcomes , numSamples ) ; const customValues = [ [ seed ] ] ; const res = backend . runWebGLProgram ( program , [ probs ] , 'int32' , customValues ) ; if ( ! normalized ) { backend . disposeIntermediateTensorInfo ( probs ) ; } return res ; } export const multinomialConfig = { kernelName : Multinomial , backendName : 'webgl' , kernelFunc : multinomial as { } as KernelFunc } ;