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

Allow tf.resource_name argument attributes in tf saved model dialect

Browse Source 
PiperOrigin-RevId: 315595041
Change-Id: Ib8dd9ba31f179532095e540ad279c5e71ebbe20c
This commit is contained in:
Jaesung Chung 2020-06-09 17:29:15 -07:00 committed by TensorFlower Gardener
parent c9c8b624e3
commit ec2fb44030
3 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.cc
View File

@ -254,9 +254,12 @@ LogicalResult VerifyExportedFunc(FuncOp func) {
} }
continue; continue;
} }
if (func.getArgAttr(i, "tf.resource_name")) {
continue;
}
return func.emitError() return func.emitError()
<< "all arguments should have 'tf_saved_model.index_path' or " << "all arguments should have 'tf_saved_model.index_path', "
"'tf_saved_model.bound_input' attributes"; "'tf_saved_model.bound_input' or 'tf.resource_name' attributes";
} }
llvm::SmallDenseSet<StringRef, 8> unique_bound_inputs; llvm::SmallDenseSet<StringRef, 8> unique_bound_inputs;
for (int i = 0, e = func.getNumArguments(); i < e; i++) { for (int i = 0, e = func.getNumArguments(); i < e; i++) {

13
tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops.mlir
View File

@ -40,3 +40,16 @@ module attributes {tf_saved_model.semantics} {
} }
} }
// -----
module attributes {tf_saved_model.semantics} {
// CHECK: func @f
func @f(
%arg0: tensor<f32> {tf.resource_name = "resource"}
) attributes { tf_saved_model.exported_names = ["foo.some_func"] } {
return
}
}

2
tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_ops_invalid.mlir
View File

@ -120,7 +120,7 @@ module attributes {tf_saved_model.semantics} {
module attributes {tf_saved_model.semantics} { module attributes {tf_saved_model.semantics} {
// expected-error@+1 {{all arguments should have 'tf_saved_model.index_path' or 'tf_saved_model.bound_input' attributes}} // expected-error@+1 {{all arguments should have 'tf_saved_model.index_path', 'tf_saved_model.bound_input' or 'tf.resource_name' attributes}}
func @f( func @f(
%arg0: tensor<f32> %arg0: tensor<f32>
) attributes { tf_saved_model.exported_names = ["f"] } { ) attributes { tf_saved_model.exported_names = ["f"] } {