From 227e5f1e864f9647badfea2b3cea45e7274fc9bd Mon Sep 17 00:00:00 2001 From: George Karpenkov Date: Wed, 27 Jan 2021 20:13:27 -0800 Subject: [PATCH] [TF2XLA] Print a helpful log message when the argument which has to be constant is not a loop invariant PiperOrigin-RevId: 354230507 Change-Id: I8010486d560e85f11bc57780c590d08a09b65120 --- tensorflow/compiler/tf2xla/const_analysis.cc | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tensorflow/compiler/tf2xla/const_analysis.cc b/tensorflow/compiler/tf2xla/const_analysis.cc index 694aa342aac..5fb393093d0 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -113,20 +113,21 @@ Status GetCompileTimeConstInputs(const NodeDef& node, const OpKernel* op_kernel, for (int i = 0; i < num_inputs; i++) { if (compile_time_const_arg_indices[i]) { // Check that this input is actually a loop invariant. - // NOTE(srbs): Ideally this should raise an error if the loop body - // requires the input at this index to be a compile time const but it is - // not a loop invariant. However, that causes problems because const - // analysis is performed for the entire graph (in the - // MarkForCompilationPass for example) and not just for the ops - // that will actually be run using XLA kernels. So we silently return - // here and let the error be raised during the actual compilation of the - // XLA graph. Node* arg_i = fbody->arg_nodes[i]; Node* ret_i = fbody->ret_nodes[i]; const Node* ret_i_input_0; TF_RETURN_IF_ERROR(ret_i->input_node(0, &ret_i_input_0)); if (ret_i_input_0->id() == arg_i->id()) { const_input_idxs->push_back(i); + } else { + // TODO(b/178546817): Verify that it's OK and raise an error if we are + // using this branch from jit_compile=True. + VLOG(1) << "Argument " << i << " to while-loop " + << node.ShortDebugString() + << " has to be constant, but it's not a loop invariant, " + "cluster compilation likely to fail at compile time: " + << arg_i->def().ShortDebugString() << " vs. " + << ret_i->def().ShortDebugString(); } } }