// CHECK: return %0#1 : tensor<f32>
  func.return %0#1 : tensor<f32>
}

// -----

// Test case to test bug due to checking
// `while_op.getResult(arg_index).use_empty()` instead of
// `while_op.getResult(while_index).use_empty()` in the tfl.while
// canonicalization.
// arg0 is a pass through. After first iteration, arg_index = 0

    auto new_while = rewriter.create<TF::WhileRegionOp>(
        while_op.getLoc(), while_op->getResultTypes(), while_op->getOperands(),
        /*parallel_iterations=*/10,
        /*is_stateless=*/false, /*shape_invariant=*/false);
    new_while.getCond().takeBody(while_op.getCond());
    new_while.getBody().takeBody(while_op.getBody());
    ReplaceReturnOp(new_while.getCond(), rewriter);

    NameAttrList cond_fn;
    cond_fn.set_name("cond_fn");
    NameAttrList body_fn;
    body_fn.set_name("body_fn");
    auto while_op =
        ops::While(s.WithOpName("while"), std::initializer_list<Input>{input},
                   cond_fn, body_fn);
    ops::_Retval retval(s.WithOpName("retval"), while_op.output[0], 0);
    std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));
    TF_CHECK_OK(s.ToGraph(g.get()));

                                TF::WhileOp& while_op) {
  // Given a while loop body function 'body_func', find the tf.While Op that
  // uses it.
  auto uses_optional = body_func.getSymbolUses(module);
  if (!uses_optional.has_value()) {
    body_func.emitOpError() << "no use of while loop body";
    return LogicalResult::failure();
  }
  while_op = nullptr;
  for (auto& use : uses_optional.value()) {

      AddRegionSideEffectsForOp(if_op.then_function().getBody(), op);
      AddRegionSideEffectsForOp(if_op.else_function().getBody(), op);
    } else if (auto while_op = dyn_cast<WhileOp>(op)) {
      AddRegionSideEffectsForOp(while_op.body_function().getBody(), op);
    } else if (auto while_region_op = dyn_cast<WhileRegionOp>(op)) {
      AddRegionSideEffectsForOp(while_region_op.getBody(), op);

    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileExpr.kt");
  }

  @Test
  @TestMetadata("whileLoop.kt")
  public void testWhileLoop() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileLoop.kt");
  }

  @Test
  @TestMetadata("whileLoopBlock.kt")
  public void testWhileLoopBlock() {

    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileExpr.kt");
  }

  @Test
  @TestMetadata("whileLoop.kt")
  public void testWhileLoop() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileLoop.kt");
  }

  @Test
  @TestMetadata("whileLoopBlock.kt")
  public void testWhileLoopBlock() {

  // CHECK-NOT: tf.EmptyTensorList
  %tl = "tf.EmptyTensorList"(%elem_shape, %size) : (tensor<0xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<f32>>>
  %while_op:2 = "tf.WhileRegion"(%tl, %size) <{is_stateless = false}> ({
  // CHECK: ^bb0(%[[CARG0:.*]]: tensor<10xf32>, %[[CARG1:.*]]: tensor<i32>, %[[CARG2:.*]]: tensor<1xi32>):

    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileExpr.kt");
  }

  @Test
  @TestMetadata("whileLoop.kt")
  public void testWhileLoop() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileLoop.kt");
  }

  @Test
  @TestMetadata("whileLoopBlock.kt")
  public void testWhileLoopBlock() {

    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileExpr.kt");
  }

  @Test
  @TestMetadata("whileLoop.kt")
  public void testWhileLoop() {
    runTest("analysis/analysis-api/testData/components/expressionInfoProvider/isUsedAsExpression/whileLoop.kt");
  }

  @Test
  @TestMetadata("whileLoopBlock.kt")
  public void testWhileLoopBlock() {