while_op, while_op.getBodyAttr());
      for (auto &arg : while_op->getOpOperands()) {
        BlockArgument cond_arg = cond.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(cond_arg), *getLatticeElement(arg.get()));
        BlockArgument body_arg = body.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(body_arg), *getLatticeElement(arg.get()));
      }

      auto type_from_then = GetElementTypeFromAccess(
          then_branch.getArgument(use.getOperandNumber() - 1), module,
          infer_from_op);
      if (type_from_then.has_value()) return type_from_then;
      auto type_from_else = GetElementTypeFromAccess(
          else_branch.getArgument(use.getOperandNumber() - 1), module,
          infer_from_op);
      if (type_from_else.has_value()) return type_from_else;

    for (Operation& op : block.getOperations()) {
      if (!op.hasTrait<OpTrait::ReturnLike>()) continue;
      for (OpOperand& operand : op.getOpOperands()) {
        int i = operand.getOperandNumber();
        if (used_argument[i] && operand.get() != used_argument[i]) {
          assert(i < number_of_results);
          invariant.reset(i);
        }
        used_argument[i] = operand.get();
      }
    }

      if (llvm::isa<tf_device::LaunchOp>(parent)) {
        new_value = parent->getResult(operand.getOperandNumber());
      } else if (auto replicate =
                     llvm::dyn_cast<tf_device::ReplicateOp>(parent)) {
        int n = replicate.getN();
        new_value = parent->getResult(operand.getOperandNumber() * n);
      } else {
        parent->emitOpError("Unsupported wrapper op.");
        return nullptr;

        LLVM_DEBUG(llvm::dbgs()
                   << "Not duplicating constant operand since it has only one "
                      "usage. Op: "
                   << curr_operand->getOperandNumber()
                   << ", operand idx: " << curr_operand->getOperandNumber()
                   << ", loc: " << owning_op->getLoc() << "\n");
        continue;
      }

      mlir::OpBuilder builder{owning_op->getContext()};

  if (parent) {
    if (llvm::isa<tf_device::ParallelExecuteOp>(parent)) {
      if (auto execute =
              llvm::dyn_cast_or_null<tf_device::ParallelExecuteOp>(parent)) {
        int operand_idx = use->getOperandNumber();
        auto execute_results = execute.GetRegionOutputs(
            device_return->getParentRegion()->getRegionNumber());
        if (operand_idx >= execute_results.size()) return nullptr;

      auto device_attr = func.getResultAttrOfType<StringAttr>(
          operand.getOperandNumber(), kFuncDeviceAttr);
      if (device_attr && !device_attr.getValue().empty()) continue;

      // Update the existing attribute named `kFuncDeviceAttr` if found.
      // Otherwise introduce a new attribute.
      auto& resultAttrForOp = result_attrs[operand.getOperandNumber()];
      bool found = false;

        arguments_to_erase) {
  auto arg = region->getArgument(argument_index);
  for (auto& usage : arg.getUses()) {
    auto* user_op = usage.getOwner();
    int operand_index = usage.getOperandNumber();
    PropagateUsage(user_op, operand_index, value, work_list,
                   arguments_to_erase);
  }
}

// Traces usage of 'var_handle_op' and replaces it's usage with constant value
// 'value'.

      }
      // Update the user operation to use the result of the cloned ConstOp.
      mlir::Operation *new_op = cloned_op_by_device[device_attr.getValue()];
      user_op->setOperand(use.getOperandNumber(), new_op->getResult(0));
    }
    // Erase the original ConstOp if all its uses have been replaced.
    if (all_uses_replaced) {
      op.erase();
    }
    return WalkResult::advance();
  });
}

  for (OpOperand& operand : return_op->getOpOperands()) {
    if (non_tpu_args.contains(operand.getOperandNumber())) {
      if (BlockArgument argument =
              llvm::dyn_cast<BlockArgument>(operand.get())) {
        if (argument.getArgNumber() != operand.getOperandNumber()) {
          return_op.emitOpError()
              << "non TPU ops do not divide state into two pieces.";