diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.cc b/tensorflow/compiler/xla/service/while_loop_analysis.cc
index 40f268f889b..ac52a5fc2a4 100644
--- a/tensorflow/compiler/xla/service/while_loop_analysis.cc
+++ b/tensorflow/compiler/xla/service/while_loop_analysis.cc
@@ -80,7 +80,7 @@ static optional<int64> GetGTEOperandIndex(const HloInstruction* instr,
 
 // Tries to get the tuple index of the induction variable of a while loop.
 //
-// Checks that the loop condition and root both plumb the induction variable
+// Checks that the loop condition and body both plumb the induction variable
 // through the same tuple index, and that they both apply exactly one op to the
 // induction variable before  deciding whether to do another loop iteration (in
 // the loop condition's case) or packing the induction variable into the result
@@ -96,8 +96,7 @@ static optional<int64> GetGTEOperandIndex(const HloInstruction* instr,
 //   root = tuple(..., inc, ...)  // inc is N'th operand of tuple().
 //
 // If so, returns N.  Otherwise, returns nullopt.
-static optional<int64> GetLoopInductionVarTupleIdx(
-    const HloInstruction* while_op) {
+optional<int64> GetLoopInductionVarTupleIdx(const HloInstruction* while_op) {
   CHECK_EQ(while_op->opcode(), HloOpcode::kWhile);
   VLOG(2) << "Finding induction variable for loop "
           << while_op->ToShortString();
diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.h b/tensorflow/compiler/xla/service/while_loop_analysis.h
index 9bb784a544b..10b64459974 100644
--- a/tensorflow/compiler/xla/service/while_loop_analysis.h
+++ b/tensorflow/compiler/xla/service/while_loop_analysis.h
@@ -35,6 +35,11 @@ absl::optional<int64> ComputeWhileLoopTripCount(
 // known, nullopt otherwise.
 absl::optional<int64> ComputeWhileLoopTripCountUpperBound(
     HloInstruction *while_op);
+
+// Returns the tuple index of the loop induction variable if there is such an
+// induction variable detected. Otherwise returns nullopt.
+absl::optional<int64> GetLoopInductionVarTupleIdx(
+    const HloInstruction *while_op);
 }  // namespace xla
 
 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_ANALYSIS_H_