experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix handling of tf_saved_model.asset in tf-saved-model-optimize-global-tensors

Browse Source 
PiperOrigin-RevId: 351259338
Change-Id: Idfa548a3a20b1b5f9e77e5c59010060d85d4578d
This commit is contained in:
Sean Silva 2021-01-11 16:22:25 -08:00 committed by TensorFlower Gardener
parent 5e90f547a2
commit d7f96c409b
2 changed files with 19 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_optimize_global_tensors.mlir
View File

@ -228,3 +228,19 @@ module attributes {tf_saved_model.semantics} {
return
}
}
// -----
// Test presence of tf_saved_model.asset's as bound inputs.
// It should not crash.
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@legacy_init_op]} : () -> ()
"tf_saved_model.asset"() {filename = "assets/foo.txt", sym_name = "asset"} : () -> ()
// CHECK: @legacy_init_op
func @legacy_init_op(%arg0: tensor<!tf.string> {tf_saved_model.bound_input = @asset}) attributes {tf_saved_model.exported_names = ["f"]} {
%0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf.string, shared_name = "hash_table_shared_name", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf.resource>
"tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, value_index = -1 : i64, vocab_size = 205 : i64} : (tensor<!tf.resource>, tensor<!tf.string>) -> ()
return
}
}

3
tensorflow/compiler/mlir/tensorflow/transforms/optimize_global_tensors.cc
View File

@ -212,6 +212,9 @@ GlobalTensorUsesMap CreateGlobalTensorUsesMap(ModuleOp module) {
}
auto global_tensor = symbol_table.lookup<GlobalTensorOp>(
sym.cast<FlatSymbolRefAttr>().getValue());
if (!global_tensor) {
continue;
}
global_tensor_uses[global_tensor].push_back({func, i});
}
}