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

[TF2XLA] Do not lose stack traces while doing tf2xla function transformations

Browse Source 
PiperOrigin-RevId: 348732953
Change-Id: I96501f9453e45efe9d85236081afffdf6b1e3a24
This commit is contained in:
George Karpenkov 2020-12-22 20:03:00 -08:00 committed by TensorFlower Gardener
parent 907bde43cc
commit 917ebe0008
2 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/tf2xla/rearrange_function_argument.cc
View File

@ -366,7 +366,8 @@ Status MaybeRewriteWhileNode(
string new_name = string new_name =
fld->UniqueFunctionName(absl::StrCat(attr_value.name(), "_rearrange_")); fld->UniqueFunctionName(absl::StrCat(attr_value.name(), "_rearrange_"));
TF_RETURN_IF_ERROR(GraphToFunctionDef(*fbody->graph, new_name, &new_fdef)); TF_RETURN_IF_ERROR(GraphToFunctionDef(*fbody->graph, new_name, &new_fdef));
TF_RETURN_IF_ERROR(fld->AddFunctionDef(new_fdef)); TF_RETURN_IF_ERROR(
fld->AddFunctionDef(new_fdef, fld->GetStackTraces(attr_value.name())));
// Change node to use rewritten function. // Change node to use rewritten function.
attr_value.set_name(new_name); attr_value.set_name(new_name);
@ -455,7 +456,8 @@ Status MaybeRewriteIfNode(
string new_name = string new_name =
fld->UniqueFunctionName(absl::StrCat(f.name(), "_rearrange_")); fld->UniqueFunctionName(absl::StrCat(f.name(), "_rearrange_"));
TF_RETURN_IF_ERROR(GraphToFunctionDef(*fbody->graph, new_name, &new_fdef)); TF_RETURN_IF_ERROR(GraphToFunctionDef(*fbody->graph, new_name, &new_fdef));
TF_RETURN_IF_ERROR(fld->AddFunctionDef(new_fdef)); TF_RETURN_IF_ERROR(
fld->AddFunctionDef(new_fdef, fld->GetStackTraces(f.name())));
// Change node to use rewritten function. // Change node to use rewritten function.
f.set_name(new_name); f.set_name(new_name);

3
tensorflow/compiler/tf2xla/tf2xla_util.cc
View File

@ -206,6 +206,9 @@ Status PropagateConstIntoFuncAttr(
n->ClearAttr(attr_name); n->ClearAttr(attr_name);
n->AddAttr(attr_name, func_attr); n->AddAttr(attr_name, func_attr);
TF_RETURN_IF_ERROR(fld->AddFunctionDef(
replace_fdef, lookup_fld->GetStackTraces(func_attr.name())));
// Copy associated functions. // Copy associated functions.
TF_RETURN_IF_ERROR(CopyAssociatedFunctions(func_graph, lookup_fld, fld)); TF_RETURN_IF_ERROR(CopyAssociatedFunctions(func_graph, lookup_fld, fld));