license import { KernelConfig , Multinomial , MultinomialAttrs , MultinomialInputs , scalar } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const multinomialConfig = { kernelName : Multinomial , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { logits } = args . inputs as MultinomialInputs ; const backend = args . backend as NodeJSKernelBackend ; const { numSamples , seed , normalized } = args . attrs as { } as MultinomialAttrs ; if ( normalized ) { throw new Error ( 'TF Node backend does not support normalized logits ' + 'passed to multinomial' ) ; } const opAttrs = [ createTensorsTypeOpAttr ( 'T' , logits . dtype ) , createTensorsTypeOpAttr ( 'output_dtype' , 'int32' ) , { name : 'seed' , type : backend . binding . TF_ATTR_INT , value : seed } , { name : 'seed2' , type : backend . binding . TF_ATTR_INT , value : seed * seed } , ] ; const numSamplesTensor = scalar ( numSamples , 'int32' ) ; const res = backend . executeSingleOutput ( Multinomial , opAttrs , [ logits , numSamplesTensor ] ) ; numSamplesTensor . dispose ( ) ; return res ; } } ;