license import * as tf from '@tensorflow/tfjs-core' ; import { engine , Tensor , TensorInfo , test_util } from '@tensorflow/tfjs-core' ; import { describeWithFlags } from '@tensorflow/tfjs-core/dist/jasmine_util' ; import { WEBGL_ENVS } from './backend_webgl_test_registry' ; import { GPGPUProgram , MathBackendWebGL } from './webgl' ; describeWithFlags ( 'custom-op webgl' , WEBGL_ENVS , ( ) => { class SquareAndAddKernel implements GPGPUProgram { variableNames = [ 'X' ] ; outputShape : number [ ] ; userCode : string ; constructor ( inputShape ) { this . outputShape = inputShape . slice ( ) ; this . userCode = `
          void main() {
            float x = getXAtOutCoords();
            float value = x * x + x;
            setOutput(value);
          }
        ` ; } } class SquareAndAddBackpropKernel implements GPGPUProgram { variableNames = [ 'X' ] ; outputShape : number [ ] ; userCode : string ; constructor ( inputShape ) { this . outputShape = inputShape . slice ( ) ; this . userCode = `
          void main() {
            float x = getXAtOutCoords();
            float value = 2.0 * x + 1.0;
            setOutput(value);
          }
        ` ; } } function squareAndAdd < T extends tf . Tensor > ( x ) { const fn = tf . customGrad ( ( x , save ) => { save ( [ x ] ) ; const webglBackend = tf . backend ( ) as MathBackendWebGL ; const program = new SquareAndAddKernel ( x . shape ) ; const backpropProgram = new SquareAndAddBackpropKernel ( x . shape ) ; const outInfo = webglBackend . compileAndRun ( program , [ x ] ) ; const value = engine ( ) . makeTensorFromDataId ( outInfo . dataId , outInfo . shape , outInfo . dtype ) as T ; const gradFunc = ( dy , saved ) => { const [ x ] = saved ; const backInfo = webglBackend . compileAndRun ( backpropProgram , [ x ] ) ; const back = engine ( ) . makeTensorFromDataId ( backInfo . dataId , backInfo . shape , backInfo . dtype ) as T ; return back . mul ( dy ) ; } ; return { value , gradFunc } ; } ) ; return fn ( x ) ; } it ( 'lets users use custom operations' , async ( ) => { const inputArr = [ 1 , 2 , 3 , 4 ] ; const input = tf . tensor ( inputArr ) ; const output = squareAndAdd ( input ) ; test_util . expectArraysClose ( await output . data ( ) , inputArr . map ( x => x * x + x ) ) ; } ) ; it ( 'lets users define gradients for operations' , async ( ) => { const inputArr = [ 1 , 2 , 3 , 4 ] ; const input = tf . tensor ( inputArr ) ; const grads = tf . valueAndGrad ( x => squareAndAdd ( x ) ) ; const { value , grad } = grads ( input ) ; test_util . expectArraysClose ( await value . data ( ) , inputArr . map ( x => x * x + x ) ) ; test_util . expectArraysClose ( await grad . data ( ) , inputArr . map ( x => 2 * x + 1 ) ) ; } ) ; it ( 'multiplies by dy parameter when it is passed' , async ( ) => { const inputArr = [ 1 , 2 , 3 , 4 ] ; const input = tf . tensor ( inputArr ) ; const grads = tf . valueAndGrad ( x => squareAndAdd ( x ) ) ; const { value , grad } = grads ( input , tf . zerosLike ( input ) ) ; test_util . expectArraysClose ( await value . data ( ) , inputArr . map ( x => x * x + x ) ) ; test_util . expectArraysClose ( await grad . data ( ) , inputArr . map ( ( ) => 0.0 ) ) ; } ) ; } ) ;