license import { DataType , Tensor } from '@tensorflow/tfjs-core' ; import { NamedTensorsMap } from '../../data/types' ; import { ExecutionContext } from '../../executor/execution_context' ; import { HashTable } from '../../executor/hash_table' ; import { ResourceManager } from '../../executor/resource_manager' ; import { InternalOpAsyncExecutor , Node } from '../types' ; import { getParamValue } from './utils' ; export const executeOp = async ( node , tensorMap , context , resourceManager ) : Promise < Tensor [ ] > => { switch ( node . op ) { case 'HashTable' : case 'HashTableV2' : { const keyDType = getParamValue ( 'keyDType' , node , tensorMap , context ) as DataType ; const valueDType = getParamValue ( 'valueDType' , node , tensorMap , context ) as DataType ; const hashTable = new HashTable ( keyDType , valueDType ) ; resourceManager . addHashTable ( node . name , hashTable ) ; return [ hashTable . handle ] ; } case 'LookupTableImport' : case 'LookupTableImportV2' : { const handle = getParamValue ( 'tableHandle' , node , tensorMap , context , resourceManager ) as Tensor ; const keys = getParamValue ( 'keys' , node , tensorMap , context ) as Tensor ; const values = getParamValue ( 'values' , node , tensorMap , context ) as Tensor ; const hashTable = resourceManager . getHashTableById ( handle . id ) ; return [ await hashTable . import ( keys , values ) ] ; } case 'LookupTableFind' : case 'LookupTableFindV2' : { const handle = getParamValue ( 'tableHandle' , node , tensorMap , context , resourceManager ) as Tensor ; const keys = getParamValue ( 'keys' , node , tensorMap , context ) as Tensor ; const defaultValue = getParamValue ( 'defaultValue' , node , tensorMap , context ) as Tensor ; const hashTable = resourceManager . getHashTableById ( handle . id ) ; return [ await hashTable . find ( keys , defaultValue ) ] ; } case 'LookupTableSize' : case 'LookupTableSizeV2' : { const handle = getParamValue ( 'tableHandle' , node , tensorMap , context , resourceManager ) as Tensor ; const hashTable = resourceManager . getHashTableById ( handle . id ) ; return [ hashTable . tensorSize ( ) ] ; } default : throw TypeError ( `Node type ${ node . op } is not implemented` ) ; } } ; export const CATEGORY = 'hash_table' ;