license import * as tf from '@tensorflow/tfjs-core' ; import { test_util } from '@tensorflow/tfjs-core' ; const { expectArraysClose } = test_util ; import { describeWithFlags } from '@tensorflow/tfjs-core/dist/jasmine_util' ; import { PACKED_ENVS , WEBGL_ENVS } from './backend_webgl_test_registry' ; describeWithFlags ( 'batchNorm' , WEBGL_ENVS , ( ) => { it ( 'should work for broadcasted inputs' , async ( ) => { const x = tf . tensor4d ( [ 2 , 4 , 9 , 23 ] , [ 2 , 1 , 1 , 2 ] ) ; const mean = tf . tensor4d ( [ 1 ] , [ 1 , 1 , 1 , 1 ] ) ; const variance = tf . tensor4d ( [ 1 ] , [ 1 , 1 , 1 , 1 ] ) ; const result = tf . batchNorm4d ( x , mean , variance ) ; expectArraysClose ( await result . data ( ) , [ 0.9995003 , 2.9985011 , 7.9960027 , 21.9890079 ] ) ; } ) ; it ( 'should work when squarification results in zero padding' , async ( ) => { const maxTextureSize = tf . env ( ) . getNumber ( 'WEBGL_MAX_TEXTURE_SIZE' ) ; tf . env ( ) . set ( 'WEBGL_MAX_TEXTURE_SIZE' , 5 ) ; const x = tf . tensor3d ( [ 0.49955603 , 0.04158615 , - 1.09440524 , 2.03854165 , - 0.61578344 , 2.87533573 , 1.18105987 , 0.807462 , 1.87888837 , 2.26563962 , - 0.37040935 , 1.35848753 , - 0.75347094 , 0.15683117 , 0.91925946 , 0.34121279 , 0.92717143 , 1.89683965 ] , [ 2 , 3 , 3 ] ) ; const mean = tf . tensor1d ( [ 0.39745062 , - 0.48062894 , 0.4847822 ] ) ; const variance = tf . tensor1d ( [ 0.32375343 , 0.67117643 , 1.08334653 ] ) ; const offset = tf . tensor1d ( [ 0.69398749 , - 1.29056387 , 0.9429723 ] ) ; const scale = tf . tensor1d ( [ - 0.5607271 , 0.9878457 , 0.25181573 ] ) ; const varianceEpsilon = .001 ; const result = tf . batchNorm3d ( x , mean , variance , offset , scale , varianceEpsilon ) ; tf . env ( ) . set ( 'WEBGL_MAX_TEXTURE_SIZE' , maxTextureSize ) ; expectArraysClose ( await result . data ( ) , [ 0.59352049 , - 0.66135202 , 0.5610874 , - 0.92077015 , - 1.45341019 , 1.52106473 , - 0.07704776 , 0.26144429 , 1.28010017 , - 1.14422404 , - 1.15776136 , 1.15425493 , 1.82644104 , - 0.52249442 , 1.04803919 , 0.74932291 , 0.40568101 , 1.2844412 ] ) ; } ) ; } ) ; describeWithFlags ( 'batchnorm packed' , PACKED_ENVS , ( ) => { it ( 'should not leak memory' , ( ) => { const x = tf . tensor4d ( [ 2 , 4 , 9 , 23 ] , [ 2 , 1 , 1 , 2 ] ) ; const mean = tf . tensor1d ( [ 1 , 2 ] ) ; const variance = tf . tensor1d ( [ 2 , 3 ] ) ; const varianceEpsilon = .001 ; const startNumBytes = tf . memory ( ) . numBytes ; const startNumTensors = tf . memory ( ) . numTensors ; tf . batchNorm4d ( x , mean , variance , undefined , undefined , varianceEpsilon ) ; const endNumBytes = tf . memory ( ) . numBytes ; const endNumTensors = tf . memory ( ) . numTensors ; expect ( endNumBytes - startNumBytes ) . toEqual ( 16 ) ; expect ( endNumTensors - startNumTensors ) . toEqual ( 1 ) ; } ) ; } ) ;