/**
 * Implementation for [Coconet] {@link
 * https://ismir2017.smcnus.org/wp-content/uploads/2017/10/187_Paper.pdf%7D} * models.
 *
 *  license import * as tf from '@tensorflow/tfjs-core' ; import { INoteSequence } from '@magenta/music' ; import * as logging from './magenta/logging' ; import * as sequences from './magenta/sequences' ; import { IS_IOS , NUM_PITCHES , MIN_PITCH , pianorollToSequence , sequenceToPianoroll , } from './coconet_utils' ; import { makeNotesTriadForKey , getHappyTriadsForKey , getSadTriadsForKey , getAllHappyTriads , getAllSadTriads , } from '../tonal-utils' ; import { N_NOTE_DIVISIONS } from '../constants' ; param step param voice interface InfillMask { step : number ; voice : number ; } param discourageNotes param nudgeFactor interface SimilarDifferentConfig { discourageNotes : boolean ; nudgeFactor : number ; } interface param key param mode interface MoodConfig { key : string ; mode : string ; happy : boolean ; } param temperature param numIterations param infillMask param moodConfig interface CoconetConfig { temperature ? : number ; numIterations ? : number ; infillMask ? : InfillMask [ ] ; similarDifferentConfig ? : SimilarDifferentConfig ; moodConfig ? : MoodConfig ; } interface LayerSpec { pooling ? : number [ ] ; filters ? : number [ ] ; activation ? : string ; dilation ? : number [ ] ; activeNoteRGB ? : string ; minPitch ? : number ; maxPitch ? : number ; poolPad ? : 'same' | 'valid' ; convPad ? : string ; convStride ? : number ; } interface ModelSpec { useSoftmaxLoss : boolean ; batchNormVarianceEpsilon : number ; numInstruments : number ; numFilters : number ; numLayers : number ; numRegularConvLayers : number ; dilation ? : number [ ] [ ] ; layers : LayerSpec [ ] ; interleaveSplitEveryNLayers ? : number ; numPointwiseSplits ? : number ; } const DEFAULT_SPEC = { useSoftmaxLoss : true , batchNormVarianceEpsilon : 1.0e-7 , numInstruments : 4 , numFilters : 128 , numLayers : 33 , numRegularConvLayers : 0 , dilation : [ [ 1 , 1 ] , [ 2 , 2 ] , [ 4 , 4 ] , [ 8 , 8 ] , [ 16 , 16 ] , [ 16 , 32 ] , [ 1 , 1 ] , [ 2 , 2 ] , [ 4 , 4 ] , [ 8 , 8 ] , [ 16 , 16 ] , [ 16 , 32 ] , [ 1 , 1 ] , [ 2 , 2 ] , [ 4 , 4 ] , [ 8 , 8 ] , [ 16 , 16 ] , [ 16 , 32 ] , [ 1 , 1 ] , [ 2 , 2 ] , [ 4 , 4 ] , [ 8 , 8 ] , [ 16 , 16 ] , [ 16 , 32 ] , [ 1 , 1 ] , [ 2 , 2 ] , [ 4 , 4 ] , [ 8 , 8 ] , [ 16 , 16 ] , [ 16 , 32 ] , ] , layers : null , interleaveSplitEveryNLayers : 16 , numPointwiseSplits : 4 , } ; class ConvNet { private residualPeriod = 2 ; private outputForResidual : tf . Tensor = null ; private residualCounter = - 1 ; private spec : ModelSpec ; private rawVars : { [ varName ] : tf . Tensor ; } = null ; constructor ( spec , vars ) { this . spec = spec ; this . rawVars = vars ; } dispose ( ) { if ( this . rawVars !== null ) { tf . dispose ( this . rawVars ) ; } if ( this . outputForResidual ) { this . outputForResidual . dispose ( ) ; } } public predictFromPianoroll ( pianoroll , masks ) { return tf . tidy ( ( ) => { let featuremaps = this . getConvnetInput ( pianoroll , masks ) ; const n = this . spec . layers . length ; for ( let i = 0 ; i < n ; i ++ ) { this . residualCounter += 1 ; this . residualSave ( featuremaps ) ; let numPointwiseSplits = null ; if ( this . spec . interleaveSplitEveryNLayers && i > 0 && i < n - 2 && i % ( this . spec . interleaveSplitEveryNLayers + 1 ) === 0 ) { numPointwiseSplits = this . spec . numPointwiseSplits ; } featuremaps = this . applyConvolution ( featuremaps , this . spec . layers [ i ] , i , i >= this . spec . numRegularConvLayers , numPointwiseSplits ) ; featuremaps = this . applyResidual ( featuremaps , i === 0 , i === n - 1 , i ) ; featuremaps = this . applyActivation ( featuremaps , this . spec . layers [ i ] , i ) ; featuremaps = this . applyPooling ( featuremaps , this . spec . layers [ i ] , i ) ; } return this . computePredictions ( featuremaps ) ; } ) ; } private computePredictions ( logits ) { if ( this . spec . useSoftmaxLoss ) { return logits . transpose ( [ 0 , 1 , 3 , 2 ] ) . softmax ( ) . transpose ( [ 0 , 1 , 3 , 2 ] ) ; } return logits . sigmoid ( ) ; } private residualReset ( ) { this . outputForResidual = null ; this . residualCounter = 0 ; } private residualSave ( x ) { if ( this . residualCounter % this . residualPeriod === 1 ) { this . outputForResidual = x ; } } private applyResidual ( x , isFirst , isLast , i ) { if ( this . outputForResidual == null ) { return x ; } if ( this . outputForResidual . shape [ this . outputForResidual . shape . length - 1 ] !== x . shape [ x . shape . length - 1 ] ) { this . residualReset ( ) ; return x ; } if ( this . residualCounter % this . residualPeriod === 0 ) { if ( ! isFirst && ! isLast ) { x = x . add ( this . outputForResidual ) ; } } return x ; } private getVar ( name , layerNum ) { const varname = `model/conv${ layerNum }/${ name }` ; return this . rawVars [ varname ] as tf . Tensor4D ; } private getSepConvVar ( name , layerNum ) { const varname = `model/conv${ layerNum }/SeparableConv2d/${ name }` ; return this . rawVars [ varname ] as tf . Tensor4D ; } private getPointwiseSplitVar ( name , layerNum , splitNum ) { const varname = `model/conv${ layerNum }/split_${ layerNum }_${ splitNum }/${ name }` ; return this . rawVars [ varname ] ; } private applyConvolution ( x , layer , i , depthwise , numPointwiseSplits ? ) { if ( layer . filters == null ) { return x ; } const filterShape = layer . filters ; const stride = layer . convStride || 1 ; const padding = layer . convPad ? ( layer . convPad . toLowerCase ( ) as 'same' | 'valid' ) : 'same' ; let conv = null ; if ( depthwise ) { const dWeights = this . getSepConvVar ( 'depthwise_weights' , i ) ; if ( ! numPointwiseSplits ) { const pWeights = this . getSepConvVar ( 'pointwise_weights' , i ) ; const biases = this . getSepConvVar ( 'biases' , i ) ; const sepConv = tf . separableConv2d ( x , dWeights , pWeights , [ stride , stride ] , padding , layer . dilation as [ number , number ] , 'NHWC' ) ; conv = sepConv . add ( biases ) ; } else { conv = tf . depthwiseConv2d ( x , dWeights , [ stride , stride ] , padding , 'NHWC' , layer . dilation as [ number , number ] ) ; const splits = tf . split ( conv , numPointwiseSplits , conv . rank - 1 ) ; const pointwiseSplits = [ ] ; for ( let splitIdx = 0 ; splitIdx < numPointwiseSplits ; splitIdx ++ ) { const outputShape = filterShape [ 3 ] / numPointwiseSplits ; const weights = this . getPointwiseSplitVar ( 'kernel' , i , splitIdx ) ; const biases = this . getPointwiseSplitVar ( 'bias' , i , splitIdx ) ; const dot = tf . matMul ( splits [ splitIdx ] . reshape ( [ - 1 , outputShape ] ) , weights , false , false ) ; const bias = tf . add ( dot , biases ) ; pointwiseSplits . push ( bias . reshape ( [ splits [ splitIdx ] . shape [ 0 ] , splits [ splitIdx ] . shape [ 1 ] , splits [ splitIdx ] . shape [ 2 ] , outputShape , ] ) ) ; } conv = tf . concat ( pointwiseSplits , conv . rank - 1 ) ; } } else { const weights = this . getVar ( 'weights' , i ) ; const stride = layer . convStride || 1 ; const padding = layer . convPad ? ( layer . convPad . toLowerCase ( ) as 'same' | 'valid' ) : 'same' ; conv = tf . conv2d ( x , weights , [ stride , stride ] , padding , 'NHWC' , [ 1 , 1 ] ) ; } return this . applyBatchnorm ( conv as tf . Tensor4D , i ) as tf . Tensor4D ; } private applyBatchnorm ( x , i ) { const gammas = this . getVar ( 'gamma' , i ) ; const betas = this . getVar ( 'beta' , i ) ; const mean = this . getVar ( 'popmean' , i ) ; const variance = this . getVar ( 'popvariance' , i ) ; if ( IS_IOS ) { const v = variance . arraySync ( ) [ 0 ] [ 0 ] [ 0 ] ; const stdevs = tf . tensor ( v . map ( ( x ) => Math . sqrt ( x + this . spec . batchNormVarianceEpsilon ) ) ) ; return x . sub ( mean ) . mul ( gammas . div ( stdevs ) ) . add ( betas ) ; } return tf . batchNorm ( x , tf . squeeze ( mean ) , tf . squeeze ( variance ) , tf . squeeze ( betas ) , tf . squeeze ( gammas ) , this . spec . batchNormVarianceEpsilon ) ; } private applyActivation ( x , layer , i ) { if ( layer . activation === 'identity' ) { return x ; } return x . relu ( ) ; } private applyPooling ( x , layer , i ) { if ( layer . pooling == null ) { return x ; } const pooling = layer . pooling ; const padding = layer . poolPad ? ( layer . poolPad . toLowerCase ( ) as 'same' | 'valid' ) : 'same' ; return tf . maxPool ( x , [ pooling [ 0 ] , pooling [ 1 ] ] , [ pooling [ 0 ] , pooling [ 1 ] ] , padding ) ; } private getConvnetInput ( pianoroll , masks ) { pianoroll = tf . scalar ( 1 , 'float32' ) . sub ( masks ) . mul ( pianoroll ) ; masks = tf . scalar ( 1 , 'float32' ) . sub ( masks ) ; return pianoroll . concat ( masks , 3 ) ; } } /**
 * Coconet model implementation in TensorflowJS.
 * Thanks to [James Wexler](https://github.com/jameswex) for the original
 * implementation.
 */ class Coconet { private checkpointURL : string ; private spec : ModelSpec = null ; private convnet : ConvNet ; private initialized = false ; param checkpointURL constructor ( checkpointURL ) { this . checkpointURL = checkpointURL ; this . spec = DEFAULT_SPEC ; } /**
     * Loads variables from the checkpoint and instantiates the model.
     */ async initialize ( ) { this . dispose ( ) ; const startTime = performance . now ( ) ; this . instantiateFromSpec ( ) ; const vars = await fetch ( `${ this . checkpointURL }/weights_manifest.json` ) . then ( response => response . json ( ) ) . then ( ( manifest ) => tf . io . loadWeights ( manifest , this . checkpointURL ) ) ; this . convnet = new ConvNet ( this . spec , vars ) ; this . initialized = true ; logging . logWithDuration ( 'Initialized model' , startTime , 'Coconet' ) ; } dispose ( ) { if ( this . convnet ) { this . convnet . dispose ( ) ; } this . initialized = false ; } isInitialized ( ) { return this . initialized ; } /**
     * Sets up layer configuration from params
     */ instantiateFromSpec ( ) { const nonFinalLayerFilterOuterSizes = 3 ; const finalTwoLayersFilterOuterSizes = 2 ; this . spec . layers = [ ] ; this . spec . layers . push ( { filters : [ nonFinalLayerFilterOuterSizes , nonFinalLayerFilterOuterSizes , this . spec . numInstruments * 2 , this . spec . numFilters , ] , } ) ; for ( let i = 0 ; i < this . spec . numLayers - 3 ; i ++ ) { this . spec . layers . push ( { filters : [ nonFinalLayerFilterOuterSizes , nonFinalLayerFilterOuterSizes , this . spec . numFilters , this . spec . numFilters , ] , dilation : this . spec . dilation ? this . spec . dilation [ i ] : null , } ) ; } this . spec . layers . push ( { filters : [ finalTwoLayersFilterOuterSizes , finalTwoLayersFilterOuterSizes , this . spec . numFilters , this . spec . numFilters , ] , } ) ; this . spec . layers . push ( { filters : [ finalTwoLayersFilterOuterSizes , finalTwoLayersFilterOuterSizes , this . spec . numFilters , this . spec . numInstruments , ] , activation : 'identity' , } ) ; } param sequence param config async infill ( sequence , config ? ) { sequences . assertIsRelativeQuantizedSequence ( sequence ) ; const numSteps = sequence . totalQuantizedSteps || sequence . notes [ sequence . notes . length - 1 ] . quantizedEndStep ; const pianoroll = sequenceToPianoroll ( sequence , numSteps ) ; let temperature = 0.99 ; let numIterations ; let outerMasks ; let moodConfig ; let similarDifferentConfig ; if ( config ) { temperature = config . temperature || temperature ; outerMasks = this . getCompletionMaskFromInput ( config . infillMask , pianoroll ) ; numIterations = config . numIterations || ( await this . getNumIterationsFromOuterMasks ( outerMasks ) ) ; moodConfig = config . moodConfig || moodConfig ; similarDifferentConfig = config . similarDifferentConfig || similarDifferentConfig ; } else { outerMasks = this . getCompletionMask ( pianoroll ) ; numIterations = await this . getNumIterationsFromOuterMasks ( outerMasks ) ; } const samples = await this . run ( pianoroll , numIterations , temperature , outerMasks , moodConfig , similarDifferentConfig ) ; const outputSequence = pianorollToSequence ( samples , numSteps ) ; pianoroll . dispose ( ) ; samples . dispose ( ) ; outerMasks . dispose ( ) ; return outputSequence ; } /**
     * Runs sampling on pianorolls.
     */ private async run ( pianorolls , numSteps , temperature , outerMasks , moodConfig , similarDifferentConfig ) { return this . gibbs ( pianorolls , numSteps , temperature , outerMasks , moodConfig , similarDifferentConfig ) ; } private getCompletionMaskFromInput ( masks , pianorolls ) { if ( ! masks ) { return this . getCompletionMask ( pianorolls ) ; } else { const buffer = tf . buffer ( [ pianorolls . shape [ 1 ] , 4 ] ) ; for ( let i = 0 ; i < masks . length ; i ++ ) { buffer . set ( 1 , masks [ i ] . step , masks [ i ] . voice ) ; } return tf . tidy ( ( ) => { return buffer . toTensor ( ) . expandDims ( 1 ) . tile ( [ 1 , NUM_PITCHES , 1 ] ) . expandDims ( 0 ) as tf . Tensor4D ; } ) ; } } private getCompletionMask ( pianorolls ) { return tf . tidy ( ( ) => { const isEmpty = pianorolls . sum ( 2 , true ) . equal ( tf . scalar ( 0 , 'float32' ) ) ; return tf . cast ( isEmpty , 'float32' ) . add ( tf . zerosLike ( pianorolls ) ) ; } ) ; } private async getNumIterationsFromOuterMasks ( outerMasks ) { const defaultNumIterations = 96 ; const numIterations = await tf . sum ( tf . max ( outerMasks , 2 ) ) . array ( ) ; if ( typeof numIterations === 'number' ) { return numIterations ; } else if ( Array . isArray ( numIterations ) ) { const first = numIterations [ 0 ] ; return typeof first === 'number' ? first : defaultNumIterations ; } return defaultNumIterations ; } param { tf . Tensor4D } pianorolls param { bool } discourageNotes param { number } nudgeFactor returns { tf . Tensor4D } private priorOverOriginalNotes ( pianorolls , similarDifferentConfig ) { if ( similarDifferentConfig === undefined || similarDifferentConfig . discourageNotes === undefined || similarDifferentConfig . nudgeFactor === undefined ) { return tf . onesLike ( pianorolls ) ; } return similarDifferentConfig . discourageNotes ? ( tf . mul ( tf . scalar ( similarDifferentConfig . nudgeFactor ) , tf . scalar ( 1.5 ) . sub ( pianorolls ) ) as tf . Tensor4D ) : ( tf . mul ( tf . scalar ( similarDifferentConfig . nudgeFactor ) , tf . scalar ( 0.5 ) . add ( pianorolls ) ) as tf . Tensor4D ) ; } private getTriadPitchBuffer ( pianorolls , key , mode , hardPrior ) { const triadPitches = makeNotesTriadForKey ( key , mode ) ; const nPitches = pianorolls . shape [ 2 ] ; const pitches = tf . buffer ( [ nPitches ] ) ; if ( hardPrior ) { for ( let i = 0 ; i < triadPitches . length ; i ++ ) { pitches . set ( 1 , triadPitches [ i ] . pitch - MIN_PITCH ) ; } } else { const triadPitchSet = new Set < number > ( triadPitches . map ( note => note . pitch ) ) ; for ( let i = 0 ; i < nPitches ; i ++ ) { if ( triadPitchSet . has ( i + MIN_PITCH ) ) { pitches . set ( 1.5 , i ) ; } else { pitches . set ( 0.5 , i ) ; } } } return pitches ; } private mapTemperatureToNHeldSteps ( temperature , nNoteDivisions = N_NOTE_DIVISIONS ) { if ( temperature <= 0.25 ) { return nNoteDivisions / 2 ; } else if ( temperature <= 0.75 ) { return nNoteDivisions / 4 ; } else if ( temperature <= 1.5 ) { return nNoteDivisions / 8 ; } else { return nNoteDivisions / 16 ; } } private getMostLikelyMoodPrior ( predictions , moodConfig , temperature ) { if ( moodConfig === undefined ) { return tf . onesLike ( predictions ) ; } const [ nBatch , nQSteps , nPitches , nVoices ] = predictions . shape ; const moodTriads = moodConfig . happy ? getAllHappyTriads ( ) : getAllSadTriads ( ) ; const nHeldQSteps = 1 ; const nProgressions = nQSteps / nHeldQSteps ; const moodTriadPriorCandidatesJs = moodTriads . map ( triad => { return tf . tidy ( ( ) => { const pitches = this . getTriadPitchBuffer ( predictions , triad . key , triad . quality , true ) ; return pitches . toTensor ( ) . expandDims ( 0 ) . tile ( [ nHeldQSteps , 1 ] ) . expandDims ( 2 ) . tile ( [ 1 , 1 , nVoices ] ) . expandDims ( 0 ) . tile ( [ nBatch , 1 , 1 , 1 ] ) as tf . Tensor4D ; } ) ; } ) ; const moodTriadPriorCandidatesTf = tf . concat ( moodTriadPriorCandidatesJs , 0 ) ; const nCandidates = moodTriadPriorCandidatesTf . shape [ 0 ] ; const triadProgressions = [ ] ; for ( let i = 0 ; i < nProgressions ; i ++ ) { const predictionOverlap = predictions . slice ( [ 0 , i * nHeldQSteps , 0 , 0 ] , [ nBatch , nHeldQSteps , nPitches , nVoices ] ) ; const predsPerCandidate = predictionOverlap . tile ( [ nCandidates , 1 , 1 , 1 ] ) ; const triadLHs = tf . mulStrict ( predsPerCandidate , moodTriadPriorCandidatesTf ) . sum ( [ 1 , 2 , 3 ] ) ; const bestTriadIdx = tf . argMax ( triadLHs ) ; const moodTriadPrior = moodTriadPriorCandidatesTf . gather ( bestTriadIdx , 0 ) . expandDims ( 0 ) as tf . Tensor4D ; triadProgressions . push ( moodTriadPrior ) ; predictionOverlap . dispose ( ) ; predsPerCandidate . dispose ( ) ; triadLHs . dispose ( ) ; bestTriadIdx . dispose ( ) ; } moodTriadPriorCandidatesJs . forEach ( prior => prior . dispose ( ) ) ; moodTriadPriorCandidatesTf . dispose ( ) ; return tf . tidy ( ( ) => { const fullTriadProgressionPrior = tf . concat ( triadProgressions , 1 ) ; triadProgressions . forEach ( prior => prior . dispose ( ) ) ; return fullTriadProgressionPrior ; } ) ; } private getRandomMoodPrior ( predictions , moodConfig , temperature ) { if ( moodConfig === undefined ) { return tf . onesLike ( predictions ) ; } const [ nBatch , nQSteps , nPitches , nVoices ] = predictions . shape ; const moodTriads = moodConfig . happy ? getHappyTriadsForKey ( moodConfig . key , moodConfig . mode ) : getSadTriadsForKey ( moodConfig . key , moodConfig . mode ) ; const nHeldQSteps = this . mapTemperatureToNHeldSteps ( temperature ) ; const nProgressions = nQSteps / nHeldQSteps ; const triadProgressions = [ ] ; for ( let i = 0 ; i < nProgressions ; i ++ ) { const triadIdx = Math . floor ( Math . random ( ) * moodTriads . length ) ; const triad = moodTriads [ triadIdx ] ; const pitches = this . getTriadPitchBuffer ( predictions , triad . key , triad . quality , false ) ; const pitchPriorIthProgression = pitches . toTensor ( ) . expandDims ( 0 ) . tile ( [ nHeldQSteps , 1 ] ) as tf . Tensor2D ; triadProgressions . push ( pitchPriorIthProgression ) ; } return tf . concat ( triadProgressions , 0 ) . expandDims ( 2 ) . tile ( [ 1 , 1 , nVoices ] ) . expandDims ( 0 ) . tile ( [ nBatch , 1 , 1 , 1 ] ) as tf . Tensor4D ; } private async gibbs ( pianorolls , numSteps , temperature , outerMasks , moodConfig , similarDifferentConfig ) { const numStepsTensor = tf . scalar ( numSteps , 'float32' ) ; let pianoroll = pianorolls . clone ( ) ; let softPriors ; if ( ! moodConfig ) { softPriors = this . priorOverOriginalNotes ( pianorolls , similarDifferentConfig ) ; } for ( let s = 0 ; s < numSteps ; s ++ ) { const pm = this . yaoSchedule ( s , numStepsTensor ) ; const innerMasks = this . bernoulliMask ( pianoroll . shape , pm , outerMasks ) ; await tf . nextFrame ( ) ; const predictions = tf . tidy ( ( ) => { return this . convnet . predictFromPianoroll ( pianoroll , innerMasks ) ; } ) as tf . Tensor4D ; await tf . nextFrame ( ) ; if ( moodConfig ) { softPriors = tf . tidy ( ( ) => { const similarDifferentPrior = this . priorOverOriginalNotes ( pianorolls , similarDifferentConfig ) ; const moodPrior = this . getMostLikelyMoodPrior ( this . nudgeWithPrior ( predictions , similarDifferentPrior ) , moodConfig , temperature ) ; return this . nudgeWithPrior ( similarDifferentPrior , moodPrior ) ; } ) ; } const newPredictions = softPriors ? this . nudgeWithPrior ( predictions , softPriors ) : predictions ; await tf . nextFrame ( ) ; pianoroll = tf . tidy ( ( ) => { const samples = this . samplePredictions ( newPredictions , temperature ) as tf . Tensor4D ; const updatedPianorolls = tf . where ( tf . cast ( innerMasks , 'bool' ) , samples , pianoroll ) ; pianoroll . dispose ( ) ; predictions . dispose ( ) ; newPredictions . dispose ( ) ; innerMasks . dispose ( ) ; pm . dispose ( ) ; return updatedPianorolls ; } ) ; await tf . nextFrame ( ) ; } numStepsTensor . dispose ( ) ; return pianoroll ; } private yaoSchedule ( i , n ) { return tf . tidy ( ( ) => { const pmin = tf . scalar ( 0.1 , 'float32' ) ; const pmax = tf . scalar ( 0.9 , 'float32' ) ; const alpha = tf . scalar ( 0.7 , 'float32' ) ; const wat = pmax . sub ( pmin ) . mul ( tf . scalar ( i , 'float32' ) ) . div ( n ) ; const secondArg = pmax . sub ( wat ) . div ( alpha ) ; return pmin . reshape ( [ 1 ] ) . concat ( secondArg . reshape ( [ 1 ] ) ) . max ( ) ; } ) ; } private bernoulliMask ( shape , pm , outerMasks ) { return tf . tidy ( ( ) => { const [ bb , tt , pp , ii ] = shape ; const probs = tf . tile ( tf . randomUniform ( [ bb , tt , 1 , ii ] , 0 , 1 , 'float32' ) , [ 1 , 1 , pp , 1 ] ) ; const masks = probs . less ( pm ) ; return tf . cast ( masks , 'float32' ) . mul ( outerMasks ) ; } ) ; } private nudgeWithPrior ( predictions , softPriors ) { return tf . tidy ( ( ) => { const newPredictions = tf . add ( tf . log ( predictions ) , tf . log ( softPriors ) ) as tf . Tensor4D ; return tf . exp ( newPredictions ) ; } ) ; } private samplePredictions ( predictions , temperature ) { return tf . tidy ( ( ) => { predictions = tf . pow ( predictions , tf . scalar ( 1 / temperature , 'float32' ) ) ; const cmf = tf . cumsum ( predictions , 2 , false , false ) ; const totalMasses = cmf . slice ( [ 0 , 0 , cmf . shape [ 2 ] - 1 , 0 ] , [ cmf . shape [ 0 ] , cmf . shape [ 1 ] , 1 , cmf . shape [ 3 ] ] ) ; const u = tf . randomUniform ( totalMasses . shape , 0 , 1 , 'float32' ) ; const i = u . mul ( totalMasses ) . less ( cmf ) . argMax ( 2 ) ; return tf . oneHot ( i . flatten ( ) , predictions . shape [ 2 ] , 1 , 0 ) . reshape ( [ predictions . shape [ 0 ] , predictions . shape [ 1 ] , predictions . shape [ 3 ] , predictions . shape [ 2 ] , ] ) . transpose ( [ 0 , 1 , 3 , 2 ] ) ; } ) ; } } export { Coconet } ;