license import { NamedTensorMap } from '@tensorflow/tfjs-core' ; import { NamedTensorsMap } from '../data/types' ; import { parseNodeName } from '../operations/executors/utils' ; import { Graph , Node } from '../operations/types' ; export interface ExecutionInfo { inputs : NamedTensorMap ; outputs : Node [ ] ; usedNodes : Set < string > ; missingInputs : string [ ] ; dynamicNode : Node ; syncInputs : string [ ] ; } /**
 * Given graph inputs and desired outputs, find the minimal set of nodes
 * to execute in order to compute the outputs. In addition return other useful
 * info such:
 * - Missing inputs needed to compute the output.
 * - Whether the subgraph contains dynamic ops (control flow, dynamic shape).
 * - Alternative inputs in order to avoid async (dynamic op) execution.
 */ export function getExecutionSubgraph ( inputs , outputs , weightMap , initNodes ? ) { const usedNodes = new Set < string > ( ) ; const missingInputs = [ ] ; let dynamicNode = null ; let syncInputs = null ; const seen = new Set < string > ( ) ; const inputNodeNames = Object . keys ( inputs ) . map ( name => parseNodeName ( name ) [ 0 ] ) ; let initNodeNames = [ ] ; if ( initNodes != null ) { initNodeNames = initNodes . map ( node => parseNodeName ( node . name ) [ 0 ] ) ; } const frontier = [ ... outputs ] ; while ( frontier . length > 0 ) { const node = frontier . pop ( ) ; if ( isControlFlow ( node ) || isDynamicShape ( node ) || isHashTable ( node ) ) { if ( dynamicNode == null ) { dynamicNode = node ; syncInputs = dynamicNode . children . map ( child => child . name ) . filter ( name => usedNodes . has ( name ) ) ; } } usedNodes . add ( node . name ) ; if ( weightMap [ node . name ] != null ) { continue ; } if ( inputNodeNames . indexOf ( node . name ) !== - 1 ) { continue ; } if ( initNodeNames . indexOf ( node . name ) !== - 1 ) { continue ; } if ( node . inputs . length === 0 ) { missingInputs . push ( node . name ) ; continue ; } node . inputs . forEach ( input => { if ( seen . has ( input . name ) ) { return ; } seen . add ( input . name ) ; frontier . push ( input ) ; } ) ; } return { inputs , outputs , usedNodes , missingInputs , dynamicNode , syncInputs } ; } /**
 * Given the execution info, return a list of nodes in topological order that
 * need to be executed to compute the output.
 */ export function getNodesInTopologicalOrder ( graph , weightMap , executionInfo ) { const { usedNodes , inputs } = executionInfo ; const frontier = [ ] ; const inputNodes = Object . keys ( inputs ) . map ( name => parseNodeName ( name ) [ 0 ] ) . map ( name => graph . nodes [ name ] ) ; const initNodes = graph . initNodes ; inputNodes . forEach ( input => { if ( usedNodes . has ( input . name ) ) { frontier . push ( input ) ; } } ) ; graph . weights . forEach ( weight => { if ( usedNodes . has ( weight . name ) ) { frontier . push ( weight ) ; } } ) ; if ( initNodes != null ) { initNodes . forEach ( node => { if ( usedNodes . has ( node . name ) ) { frontier . push ( node ) ; } } ) ; } const seen = new Set < string > ( ) ; const orderedNodes = [ ] ; while ( frontier . length > 0 ) { const node = frontier . pop ( ) ; seen . add ( node . name ) ; if ( ! weightMap [ node . name ] ) { orderedNodes . push ( node ) ; } node . children . forEach ( child => { if ( ! seen . has ( child . name ) && usedNodes . has ( child . name ) && child . inputs . every ( input => seen . has ( input . name ) ) ) { frontier . push ( child ) ; } } ) ; } return orderedNodes ; } const CONTROL_FLOW_OPS = [ 'Switch' , 'Merge' , 'Enter' , 'Exit' , 'NextIteration' , 'StatelessIf' , 'StatelessWhile' , 'if' , 'While' ] ; const DYNAMIC_SHAPE_OPS = [ 'NonMaxSuppressionV2' , 'NonMaxSuppressionV3' , 'NonMaxSuppressionV5' , 'Where' ] ; const HASH_TABLE_OPS = [ 'HashTable' , 'HashTableV2' , 'LookupTableImport' , 'LookupTableImportV2' , 'LookupTableFind' , 'LookupTableFindV2' , 'LookupTableSize' , 'LookupTableSizeV2' ] ; export function isControlFlow ( node ) { return CONTROL_FLOW_OPS . indexOf ( node . op ) >= 0 ; } export function isDynamicShape ( node ) { return DYNAMIC_SHAPE_OPS . indexOf ( node . op ) >= 0 ; } export function isHashTable ( node ) { return HASH_TABLE_OPS . indexOf ( node . op ) >= 0 ; }