new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {
    if (!while_op.getResult(i).use_empty()) {

                                PatternRewriter& rewriter) const override {
    auto input_op = op.getInput().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<QuantizeOp>(input_op)) {
      if (!q->getAttr(mlir::quant::kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
        // Don't remove leading and trailing QDQ for PTQ workflow, so the io

        new_operands.push_back(build_list_op.getOut());
      } else {
        auto attr_name = compose_func.getArgAttrOfType<StringAttr>(
            arg.index(), kAttrArgumentNameAttr);
        auto attribute = op->getAttr(attr_name.getValue());
        if (!attribute) {
          attribute =
              compose_func.getArgAttr(arg.index(), kAttrArgumentDefaultAttr);
        }

  for (auto extra_operand : extra_operands)
    new_types.push_back(extra_operand.getType());

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_types, operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());
  while_op.replaceAllUsesWith(
      new_while_op.getResults().take_front(while_op.getNumResults()));

  builder.insert(metadata_op.clone());
  for (Operation& op : func.getOps()) {
    if (!IsTPUOp(&op)) continue;
    op.setAttr(TF::kReplicationInfoAttr, new_group);
  }
  TF::TPUCompilationResultOp new_result = compilation_op.clone();
  new_result->setAttr(kTpuCompilationStatus, new_group);
  builder.insert(new_result);
}

template <typename OpType>
LogicalResult FindAndExcludeOp(func::FuncOp func,

      {StringAttr::get(context, "outputs"),
       StringAttr::get(context, absl::StrJoin(output_names, ","))});
  auto dictAttr = DictionaryAttr::get(context, func_attrs);
  main_func->setAttr(StringAttr::get(context, kEntryFunctionAttr), dictAttr);
  main_func->setAttr(
      kTfSavedModelExportedNamesAttr,
      builder.getStrArrayAttr({kImportModelDefaultGraphFuncName}));

  if (input_names.size() != main_func.getNumArguments() ||

      /*disabled_checks=*/empty_array_attr);
  xla_call_module_op->setAttr(TF::kStablehloEntryFunctionAttrName,
                              SymbolRefAttr::get(stablehlo_func_op));
  // Set jax.uses_shape_polymorphism=true to enable shape refinement at runtime.
  // This is needed for native serialization version >= 8.
  xla_call_module_op->setAttr(
      kStablehloModuleAttrsAttrName,

  if (!op->hasAttr(TF::kTpuReplicateAttr)) {
    op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
                    "missing _tpu_replicate attr");
    return false;
  }
  auto opattr = op->getAttr(TF::kTpuReplicateAttr);
  if (opattr != attr) {
    op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
                    "invalid _tpu_replicate attr.")

namespace {

// Creates ConstOp for int32_t value.
ConstOp createI32ConstOp(int32_t value, Location loc,
                         PatternRewriter* rewriter) {
  auto int_attr = IntegerAttr::get(rewriter->getIntegerType(32), value);
  return rewriter->create<ConstOp>(loc, int_attr);
}

// Creates ConstantOp for array of int32_t.
arith::ConstantOp createI32ConstantOp(llvm::ArrayRef<int32_t> values,

    // Get the names of all defined attributes, including both derived and
    // non-derived ones.
    llvm::StringSet<> getDefinedAttributeNames() {
      llvm::StringSet<> all_attrs;
      for (auto& attr : (*this)->getAttrs()) {
        all_attrs.insert(attr.getName().strref());
      }
      for (const auto& operand : llvm::enumerate(getArgumentTypes())) {
        if (auto attr_name = getArgAttrOfType<StringAttr>(