license import { Tensor } from './tensor' ; import { NamedTensorMap } from './tensor_types' ; import * as util from './util' ; export interface TapeNode { id : number ; kernelName : string ; outputs : Tensor [ ] ; inputs : NamedTensorMap ; gradient ? : ( dys : Tensor [ ] ) => NamedGradientMap ; saved ? : Tensor [ ] ; } export type NamedGradientMap = { [ inputName ] : ( ) => Tensor ; } ; param tape param xs param y export function getFilteredNodesXToY ( tape , xs , y ) { const tensorsFromX = { } ; const nodesFromX = { } ; for ( let i = 0 ; i < xs . length ; i ++ ) { tensorsFromX [ xs [ i ] . id ] = true ; } for ( let i = 0 ; i < tape . length ; i ++ ) { const node = tape [ i ] ; const nodeInputs = node . inputs ; for ( const inputName in nodeInputs ) { const input = nodeInputs [ inputName ] ; let anyInputFromX = false ; for ( let j = 0 ; j < xs . length ; j ++ ) { if ( tensorsFromX [ input . id ] ) { node . outputs . forEach ( output => tensorsFromX [ output . id ] = true ) ; anyInputFromX = true ; nodesFromX [ node . id ] = true ; break ; } } if ( anyInputFromX ) { break ; } } } const tensorsLeadToY = { } ; tensorsLeadToY [ y . id ] = true ; const nodesToY = { } ; for ( let i = tape . length - 1 ; i >= 0 ; i -- ) { const node = tape [ i ] ; const nodeInputs = node . inputs ; for ( let j = 0 ; j < node . outputs . length ; j ++ ) { if ( tensorsLeadToY [ node . outputs [ j ] . id ] ) { for ( const inputName in nodeInputs ) { tensorsLeadToY [ nodeInputs [ inputName ] . id ] = true ; nodesToY [ node . id ] = true ; } break ; } } } const filteredTape = [ ] ; for ( let i = 0 ; i < tape . length ; i ++ ) { const node = tape [ i ] ; if ( nodesFromX [ node . id ] && nodesToY [ node . id ] ) { const prunedInputs = { } ; for ( const inputName in node . inputs ) { const nodeInput = node . inputs [ inputName ] ; if ( tensorsFromX [ nodeInput . id ] ) { prunedInputs [ inputName ] = nodeInput ; } } const prunedNode = Object . assign ( { } , node ) ; prunedNode . inputs = prunedInputs ; prunedNode . outputs = node . outputs ; filteredTape . push ( prunedNode ) ; } } return filteredTape ; } param tensorAccumulatedGradientMap param filteredTape export function backpropagateGradients ( tensorAccumulatedGradientMap , filteredTape , tidy , add ) { for ( let i = filteredTape . length - 1 ; i >= 0 ; i -- ) { const node = filteredTape [ i ] ; const dys = [ ] ; node . outputs . forEach ( o => { const gradTensor = tensorAccumulatedGradientMap [ o . id ] ; if ( gradTensor != null ) { dys . push ( gradTensor ) ; } else { dys . push ( null ) ; } } ) ; if ( node . gradient == null ) { throw new Error ( `Cannot compute gradient: gradient function not found ` + `for ${ node . kernelName }.` ) ; } const inputGradients = node . gradient ( dys ) ; for ( const inputName in node . inputs ) { if ( ! ( inputName in inputGradients ) ) { throw new Error ( `Cannot backprop through input ${ inputName }. ` + `Available gradients found: ${ Object . keys ( inputGradients ) }.` ) ; } const dx = tidy ( ( ) => inputGradients [ inputName ] ( ) ) ; if ( dx . dtype !== 'float32' ) { throw new Error ( `Error in gradient for op ${ node . kernelName }. The gradient of input ` + `${ inputName } must have 'float32' dtype, but has '${ dx . dtype }'` ) ; } const x = node . inputs [ inputName ] ; if ( ! util . arraysEqual ( dx . shape , x . shape ) ) { throw new Error ( `Error in gradient for op ${ node . kernelName }. The gradient of input ` + `'${ inputName }' has shape '${ dx . shape }', which does not match ` + `the shape of the input '${ x . shape }'` ) ; } if ( tensorAccumulatedGradientMap [ x . id ] == null ) { tensorAccumulatedGradientMap [ x . id ] = dx ; } else { const curGradient = tensorAccumulatedGradientMap [ x . id ] ; tensorAccumulatedGradientMap [ x . id ] = add ( curGradient , dx ) ; curGradient . dispose ( ) ; } } } }