Fix table recreation bug

tensorflow · Oct 5, 2022 · da7ed9c · da7ed9c
1 parent e33267c
commit da7ed9c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/e2e/integration_tests/convert_predict.py b/e2e/integration_tests/convert_predict.py
@@ -457,7 +457,7 @@ def lookup(input):
   return {
       "async": False,
       "inputs": {
-          "Placeholder:0": {
+          "input:0": {
               "value": ["a", "b", "c"], "shape": [3], "dtype": "string"
           }
       },

diff --git a/tfjs-converter/src/operations/executors/hash_table_executor.ts b/tfjs-converter/src/operations/executors/hash_table_executor.ts
@@ -31,14 +31,21 @@ export const executeOp: InternalOpAsyncExecutor = async(
   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 existingTableHandle =
+          resourceManager.getHashTableHandleByName(node.name);
+      // Table is shared with initializer.
+      if (existingTableHandle != null) {
+        return [existingTableHandle];
+      } else {
+        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];
+        const hashTable = new HashTable(keyDType, valueDType);
+        resourceManager.addHashTable(node.name, hashTable);
+        return [hashTable.handle];
+      }
     }
     case 'LookupTableImport':
     case 'LookupTableImportV2': {