if (while_op != nullptr) {
        use.getUser()->emitOpError() << "multiple users of function.";
        return LogicalResult::failure();
      } else {
        while_op = llvm::cast<TF::WhileOp>(use.getUser());
      }
    } else {
      use.getUser()->emitOpError() << "non while use of function.";
      return LogicalResult::failure();
    }
  }

      return;
    }
    if (auto while_op = dyn_cast<TF::WhileOp>(op)) {
      for (auto callee : {while_op.cond_function(), while_op.body_function()}) {
        PropagatePotentiallyWrittenUpFromCallee(callee.getRegion(),
                                                while_op.getInput());
      }
      return;
    }
    if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {

void RemoveUnusedWhileResultsPass::runOnOperation() {
  func::FuncOp func = getOperation();
  // Try to prune defining ops of unused results.
  func.walk([&](TF::WhileRegionOp while_op) {
    for (OpResult result : while_op.getResults()) {
      TryPruneResultDefiningOp(while_op, result);
    }
  });
  // Now eliminate passthrough operands/results with existing canonicalization
  // pattern.
  MLIRContext* context = &getContext();

                       /*flatten_tuple*/ true);

  // Create new `mhlo.while` op with extra token operand and result.
  auto new_while = builder.create<WhileOp>(region_while.getLoc(),
                                           new_result_types, new_val_operands);

  // Move all regions from the old `mhlo.while` op to its replacement.
  new_while.getCond().takeBody(region_while.getCond());

  }
}

// -----
// Verify one while body function per while loop op.
module {
  func.func @main() {
    %cst_main = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %0 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    // expected-error @+1 {{'tf.While' op multiple users of function.}}

  return success();
}

namespace {
// Canonicalize While op so that results and operands match and external values
// are via implicit capture rather than via block args.
struct WhileResultOperandsMatchAndImplicitCapture
    : public OpRewritePattern<WhileOp> {
  using OpRewritePattern<WhileOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(WhileOp while_op,

// -----
// Verify error for multiple WhileOp use of while_body function.
module {
  func.func @main() {
    %cst_main = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %0 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    // expected-error @+1 {{'tf.While' op multiple users of function.}}

          while_op.getInput().getTypes(), while_op.getOutput().getTypes(),
          while_op.getBody().getArgumentTypes());
      return PropagateShapeToRegions(compatible_types,
                                     {&while_op.getCond(), &while_op.getBody()},
                                     max_iterations);
    }
    return PropagateShapeToRegions(while_op.getInput().getTypes(),

    ValueRange opInput = op.getInput();
    TypeRange scf_block_arguments_type = opInput.getType();

    // Create the `scf.while` op.
    auto scf_while_op = rewriter.create<scf::WhileOp>(
        op.getLoc(), op.getResultTypes(), opInput);

    // Create the `before` block of the `scf.while` op (with an `scf.condition`
    // op as the terminator). Note that the arguments' type of this block is

        void>::getResultSegmentSizeAttr();
    attrs_to_ignore.insert(attr_name.data());
  }

  if (llvm::isa<mlir::TF::CaseOp, mlir::TF::IfOp, mlir::TF::WhileOp>(inst))
    attrs_to_ignore.insert("is_stateless");

  if (llvm::isa<mlir::TF::WhileOp>(inst))
    attrs_to_ignore.insert("shape_invariant");

  return attrs_to_ignore;
}

// Populates all derived attributes of a MLIR operation in a proto