license import { DataType , env } from '@tensorflow/tfjs-core' ; import * as tensorflow from '../data/compiled_api' ; import { getRegisteredOp } from './custom_op/register' ; import { getNodeNameAndIndex } from './executors/utils' ; import * as arithmetic from './op_list/arithmetic' ; import * as basicMath from './op_list/basic_math' ; import * as control from './op_list/control' ; import * as convolution from './op_list/convolution' ; import * as creation from './op_list/creation' ; import * as dynamic from './op_list/dynamic' ; import * as evaluation from './op_list/evaluation' ; import * as graph from './op_list/graph' ; import * as hashTable from './op_list/hash_table' ; import * as image from './op_list/image' ; import * as logical from './op_list/logical' ; import * as matrices from './op_list/matrices' ; import * as normalization from './op_list/normalization' ; import * as reduction from './op_list/reduction' ; import * as sliceJoin from './op_list/slice_join' ; import * as sparse from './op_list/sparse' ; import * as spectral from './op_list/spectral' ; import * as string from './op_list/string' ; import * as transformation from './op_list/transformation' ; import { Graph , InputParamValue , Node , OpMapper , ParamValue } from './types' ; export class OperationMapper { private static _instance : OperationMapper ; private opMappers : { [ key ] : OpMapper ; } ; public static get Instance ( ) { return this . _instance || ( this . _instance = new this ( ) ) ; } private constructor ( ) { const ops = [ arithmetic , basicMath , control , convolution , creation , dynamic , evaluation , graph , hashTable , image , logical , matrices , normalization , reduction , sliceJoin , sparse , spectral , string , transformation ] ; const mappersJson = [ ] . concat ( ... ops . map ( op => op . json ) ) ; this . opMappers = mappersJson . reduce < { [ key ] : OpMapper ; } > ( ( map , mapper ) => { map [ mapper . tfOpName ] = mapper ; return map ; } , { } ) ; } transformGraph ( graph , signature = { } ) { const tfNodes = graph . node ; const placeholders = [ ] ; const weights = [ ] ; const initNodes = [ ] ; const nodes = tfNodes . reduce < { [ key ] : Node ; } > ( ( map , node ) => { map [ node . name ] = this . mapNode ( node ) ; if ( node . op . startsWith ( 'Placeholder' ) ) { placeholders . push ( map [ node . name ] ) ; } else if ( node . op === 'Const' ) { weights . push ( map [ node . name ] ) ; } else if ( node . input == null || node . input . length === 0 ) { initNodes . push ( map [ node . name ] ) ; } return map ; } , { } ) ; let inputs = [ ] ; const outputs = [ ] ; let inputNodeNameToKey = { } ; let outputNodeNameToKey = { } ; if ( signature != null ) { inputNodeNameToKey = this . mapSignatureEntries ( signature . inputs ) ; outputNodeNameToKey = this . mapSignatureEntries ( signature . outputs ) ; } const allNodes = Object . keys ( nodes ) ; allNodes . forEach ( key => { const node = nodes [ key ] ; node . inputNames . forEach ( ( name , index ) => { const [ nodeName , , outputName ] = getNodeNameAndIndex ( name ) ; const inputNode = nodes [ nodeName ] ; if ( inputNode . outputs != null ) { const outputIndex = inputNode . outputs . indexOf ( outputName ) ; if ( outputIndex !== - 1 ) { const inputName = `${ nodeName }:${ outputIndex }` ; node . inputNames [ index ] = inputName ; } } node . inputs . push ( inputNode ) ; inputNode . children . push ( node ) ; } ) ; } ) ; if ( Object . keys ( outputNodeNameToKey ) . length === 0 ) { allNodes . forEach ( key => { const node = nodes [ key ] ; if ( node . children . length === 0 ) { outputs . push ( node ) ; } } ) ; } else { Object . keys ( outputNodeNameToKey ) . forEach ( name => { const [ nodeName , ] = getNodeNameAndIndex ( name ) ; const node = nodes [ nodeName ] ; if ( node != null ) { node . signatureKey = outputNodeNameToKey [ name ] ; outputs . push ( node ) ; } } ) ; } if ( Object . keys ( inputNodeNameToKey ) . length > 0 ) { Object . keys ( inputNodeNameToKey ) . forEach ( name => { const [ nodeName , ] = getNodeNameAndIndex ( name ) ; const node = nodes [ nodeName ] ; if ( node ) { node . signatureKey = inputNodeNameToKey [ name ] ; inputs . push ( node ) ; } } ) ; } else { inputs = placeholders ; } let functions = { } ; if ( graph . library != null && graph . library . function != null ) { functions = graph . library . function . reduce ( ( functions , func ) => { functions [ func . signature . name ] = this . mapFunction ( func ) ; return functions ; } , { } as { [ key ] : Graph ; } ) ; } const result = { nodes , inputs , outputs , weights , placeholders , signature , functions } ; if ( initNodes . length > 0 ) { result . initNodes = initNodes ; } return result ; } private mapSignatureEntries ( entries ) { return Object . keys ( entries || { } ) . reduce < { [ key ] : string ; } > ( ( prev , curr ) => { prev [ entries [ curr ] . name ] = curr ; return prev ; } , { } ) ; } private mapNode ( node ) { const mapper = getRegisteredOp ( node . op ) || this . opMappers [ node . op ] || { } as OpMapper ; if ( node . attr == null ) { node . attr = { } ; } const newNode = { name : node . name , op : node . op , category : mapper . category , inputNames : ( node . input || [ ] ) . map ( input => input . startsWith ( '^' ) ? input . substr ( 1 ) : input ) , inputs : [ ] , children : [ ] , inputParams : { } , attrParams : { } , rawAttrs : node . attr , outputs : mapper . outputs } ; if ( mapper . inputs != null ) { newNode . inputParams = mapper . inputs . reduce < { [ key ] : InputParamValue ; } > ( ( map , param ) => { map [ param . name ] = { type : param . type , inputIndexStart : param . start , inputIndexEnd : param . end } ; return map ; } , { } ) ; } if ( mapper . attrs != null ) { newNode . attrParams = mapper . attrs . reduce < { [ key ] : ParamValue ; } > ( ( map , param ) => { const type = param . type ; let value = undefined ; switch ( param . type ) { case 'string' : value = getStringParam ( node . attr , param . tfName , param . defaultValue as string ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getStringParam ( node . attr , param . tfDeprecatedName , param . defaultValue as string ) ; } break ; case 'string[]' : value = getStringArrayParam ( node . attr , param . tfName , param . defaultValue as string [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getStringArrayParam ( node . attr , param . tfDeprecatedName , param . defaultValue as string [ ] ) ; } break ; case 'number' : value = getNumberParam ( node . attr , param . tfName , ( param . defaultValue || 0 ) as number ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getNumberParam ( node . attr , param . tfDeprecatedName , param . defaultValue as number ) ; } break ; case 'number[]' : value = getNumericArrayParam ( node . attr , param . tfName , param . defaultValue as number [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getNumericArrayParam ( node . attr , param . tfDeprecatedName , param . defaultValue as number [ ] ) ; } break ; case 'bool' : value = getBoolParam ( node . attr , param . tfName , param . defaultValue as boolean ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getBoolParam ( node . attr , param . tfDeprecatedName , param . defaultValue as boolean ) ; } break ; case 'bool[]' : value = getBoolArrayParam ( node . attr , param . tfName , param . defaultValue as boolean [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getBoolArrayParam ( node . attr , param . tfDeprecatedName , param . defaultValue as boolean [ ] ) ; } break ; case 'shape' : value = getTensorShapeParam ( node . attr , param . tfName , param . defaultValue as number [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getTensorShapeParam ( node . attr , param . tfDeprecatedName , param . defaultValue as number [ ] ) ; } break ; case 'shape[]' : value = getTensorShapeArrayParam ( node . attr , param . tfName , param . defaultValue as number [ ] [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getTensorShapeArrayParam ( node . attr , param . tfDeprecatedName , param . defaultValue as number [ ] [ ] ) ; } break ; case 'dtype' : value = getDtypeParam ( node . attr , param . tfName , param . defaultValue as DataType ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getDtypeParam ( node . attr , param . tfDeprecatedName , param . defaultValue as DataType ) ; } break ; case 'dtype[]' : value = getDtypeArrayParam ( node . attr , param . tfName , param . defaultValue as DataType [ ] ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getDtypeArrayParam ( node . attr , param . tfDeprecatedName , param . defaultValue as DataType [ ] ) ; } break ; case 'func' : value = getFuncParam ( node . attr , param . tfName , param . defaultValue as string ) ; if ( value === undefined && ! ! param . tfDeprecatedName ) { value = getFuncParam ( node . attr , param . tfDeprecatedName , param . defaultValue as string ) ; } break ; case 'tensor' : case 'tensors' : break ; default : throw new Error ( `Unsupported param type: ${ param . type } for op: ${ node . op }` ) ; } map [ param . name ] = { value , type } ; return map ; } , { } ) ; } return newNode ; } private mapFunction ( functionDef ) { const tfNodes = functionDef . nodeDef ; const placeholders = [ ] ; const weights = [ ] ; let nodes = { } ; if ( tfNodes != null ) { nodes = tfNodes . reduce < { [ key ] : Node ; } > ( ( map , node ) => { map [ node . name ] = this . mapNode ( node ) ; if ( node . op === 'Const' ) { weights . push ( map [ node . name ] ) ; } return map ; } , { } ) ; } const inputs = [ ] ; const outputs = [ ] ; functionDef . signature . inputArg . forEach ( arg => { const [ nodeName , ] = getNodeNameAndIndex ( arg . name ) ; const node = { name : nodeName , op : 'Placeholder' , inputs : [ ] , inputNames : [ ] , category : 'graph' , inputParams : { } , attrParams : { dtype : { value : parseDtypeParam ( arg . type ) , type : 'dtype' } } , children : [ ] } ; node . signatureKey = arg . name ; inputs . push ( node ) ; nodes [ nodeName ] = node ; } ) ; const allNodes = Object . keys ( nodes ) ; allNodes . forEach ( key => { const node = nodes [ key ] ; node . inputNames . forEach ( ( name , index ) => { const [ nodeName , , outputName ] = getNodeNameAndIndex ( name ) ; const inputNode = nodes [ nodeName ] ; if ( inputNode . outputs != null ) { const outputIndex = inputNode . outputs . indexOf ( outputName ) ; if ( outputIndex !== - 1 ) { const inputName = `${ nodeName }:${ outputIndex }` ; node . inputNames [ index ] = inputName ; } } node . inputs . push ( inputNode ) ; inputNode . children . push ( node ) ; } ) ; } ) ; const returnNodeMap = functionDef . ret ; functionDef . signature . outputArg . forEach ( output => { const [ nodeName , index ] = getNodeNameAndIndex ( returnNodeMap [ output . name ] ) ; const node = nodes [ nodeName ] ; if ( node != null ) { node . defaultOutput = index ; outputs . push ( node ) ; } } ) ; const signature = this . mapArgsToSignature ( functionDef ) ; return { nodes , inputs , outputs , weights , placeholders , signature } ; } private mapArgsToSignature ( functionDef ) { return { methodName : functionDef . signature . name , inputs : functionDef . signature . inputArg . reduce ( ( map , arg ) => { map [ arg . name ] = this . mapArgToTensorInfo ( arg ) ; return map ; } , { } as { [ key ] : tensorflow . ITensorInfo ; } ) , outputs : functionDef . signature . outputArg . reduce ( ( map , arg ) => { map [ arg . name ] = this . mapArgToTensorInfo ( arg , functionDef . ret ) ; return map ; } , { } as { [ key ] : tensorflow . ITensorInfo ; } ) , } ; } private mapArgToTensorInfo ( arg , nameMap ? ) { let name = arg . name ; if ( nameMap != null ) { name = nameMap [ name ] ; } return { name , dtype : arg . type } ; } } export function decodeBase64 ( text ) { const global = env ( ) . global ; if ( typeof global . atob !== 'undefined' ) { return global . atob ( text ) ; } else if ( typeof Buffer !== 'undefined' ) { return new Buffer ( text , 'base64' ) . toString ( ) ; } else { throw new Error ( 'Unable to decode base64 in this environment. ' + 'Missing built-in atob() or Buffer()' ) ; } } export function parseStringParam ( s , keepCase ) { const value = Array . isArray ( s ) ? String . fromCharCode . apply ( null , s ) : decodeBase64 ( s ) ; return keepCase ? value : value . toLowerCase ( ) ; } export function getStringParam ( attrs , name , def , keepCase = false ) { const param = attrs [ name ] ; if ( param != null ) { return parseStringParam ( param . s , keepCase ) ; } return def ; } export function getBoolParam ( attrs , name , def ) { const param = attrs [ name ] ; return param ? param . b : def ; } export function getNumberParam ( attrs , name , def ) { const param = attrs [ name ] || { } ; const value = param [ 'i' ] != null ? param [ 'i' ] : ( param [ 'f' ] != null ? param [ 'f' ] : def ) ; return ( typeof value === 'number' ) ? value : parseInt ( value , 10 ) ; } export function parseDtypeParam ( value ) { if ( typeof ( value ) === 'string' ) { value = tensorflow . DataType [ value as any ] ; } switch ( value ) { case tensorflow . DataType . DT_FLOAT : case tensorflow . DataType . DT_HALF : return 'float32' ; case tensorflow . DataType . DT_INT32 : case tensorflow . DataType . DT_INT64 : case tensorflow . DataType . DT_INT8 : case tensorflow . DataType . DT_UINT8 : return 'int32' ; case tensorflow . DataType . DT_BOOL : return 'bool' ; case tensorflow . DataType . DT_DOUBLE : return 'float32' ; case tensorflow . DataType . DT_STRING : return 'string' ; default : return null ; } } export function getFuncParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param && param . func ) { return param . func . name ; } return def ; } export function getDtypeParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param && param . type ) { return parseDtypeParam ( param . type ) ; } return def ; } export function getDtypeArrayParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param && param . list && param . list . type ) { return param . list . type . map ( v => parseDtypeParam ( v ) ) ; } return def ; } export function parseTensorShapeParam ( shape ) { if ( shape . unknownRank ) { return undefined ; } if ( shape . dim != null ) { return shape . dim . map ( dim => ( typeof dim . size === 'number' ) ? dim . size : parseInt ( dim . size , 10 ) ) ; } return [ ] ; } export function getTensorShapeParam ( attrs , name , def ? ) { const param = attrs [ name ] ; if ( param && param . shape ) { return parseTensorShapeParam ( param . shape ) ; } return def ; } export function getNumericArrayParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param ) { return ( ( param . list . f && param . list . f . length ? param . list . f : param . list . i ) || [ ] ) . map ( v => ( typeof v === 'number' ) ? v : parseInt ( v , 10 ) ) ; } return def ; } export function getStringArrayParam ( attrs , name , def , keepCase = false ) { const param = attrs [ name ] ; if ( param && param . list && param . list . s ) { return param . list . s . map ( ( v ) => { return parseStringParam ( v , keepCase ) ; } ) ; } return def ; } export function getTensorShapeArrayParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param && param . list && param . list . shape ) { return param . list . shape . map ( ( v ) => { return parseTensorShapeParam ( v ) ; } ) ; } return def ; } export function getBoolArrayParam ( attrs , name , def ) { const param = attrs [ name ] ; if ( param && param . list && param . list . b ) { return param . list . b ; } return def ; }