CreateXlaCallModuleOp(inputs, outputs, result_types, reverse_subgraph,
                        stablehlo_func_op, module_op);

  // 3) Erase the replaced ops.
  for (Operation* subgraph_op : reverse_subgraph) {
    subgraph_op->erase();
  }
}

// Contains the actual logic for updating states and replacing StableHLO ops
// with tf.XlaCallModuleOps.
void UpdateStatesAndReplaceStablehloOps(

    if (argnum_num_users[arg_num] != conv2d_and_block_sizes.size() ||
        conv2d_and_block_sizes.empty()) {
      argnum_and_convolutions.erase(arg_num);
      continue;
    }
    int64_t block_size = conv2d_and_block_sizes[0].second;
    if (block_size < 2) {
      argnum_and_convolutions.erase(arg_num);
      continue;
    }
    // Continue if not all the block sizes for space to depth transform are the
    // same.

    }
    op_state.addTypes(output_tys);
    op_state.addAttributes(attr);
    auto stablehlo_op = builder.create(op_state);
    custom_op.replaceAllUsesWith(stablehlo_op);
    custom_op.erase();
  });
}

std::unique_ptr<OperationPass<func::FuncOp>> CreateTflToStablehloPass() {
  return std::make_unique<TflToStablehloPass>();
}

static PassRegistration<TflToStablehloPass> pass;

     * The compiler is happy with the casting that allows to hide the checked exception.
     * The runtime is happy with the casting because the checked exception type information is captured in a generic type parameter which gets erased.
     */
    private static <T, E extends Exception> Factory<T> toUncheckedThrowingFactory(final ThrowingFactory<T, E> throwingFactory) {
        return new Factory<T>() {
            @Nullable
            @Override

    op_property->inputs.erase(
        std::remove_if(
            op_property->inputs.begin(), op_property->inputs.end(),
            [&](std::pair<int, operator_property::TensorProperty> input) {
              return cifg_non_inputs.find(input.first) != cifg_non_inputs.end();
            }),
        op_property->inputs.end());
    op_property->intermediates.erase(

void PrepareQuantizeDRQPass::removeAllStatsOp(func::FuncOp func) {
  func.walk([&](quantfork::StatisticsOp stats_op) {
    stats_op.replaceAllUsesWith(stats_op.getArg());
    stats_op.erase();
  });
}

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.inc"

void PrepareQuantizeDRQPass::runOnOperation() {
  MLIRContext* ctx = &getContext();

  // structured lowering.
  old_parallel_execute.walk(
      [&](TF::_XlaCompileMlirPlaceholderProgramKeyOp key_op) {
        key_op.replaceAllUsesWith(compile_op->getResult(1));
        key_op.erase();
      });

  // After rewrite, if there is a TPUCompilationResultOp from the same cluster,
  // replace it with the result of the compile op. The TPUCompilationResultOp is

// if the "tf_saved_model.index_path" attribute has "__tf_file_prefix", will be
// reused if it already exist in @main. Otherwise a new file prefix argument
// will be created. @tf_quant__save function will be erased.
//
// Running this pass essentially has the effect of inlining the @tf_quant__save
// into the main graph. This is beneficial when we wish to find and fetch

#include "tensorflow/core/framework/types.pb.h"

namespace mlir {
namespace quant {

// A unit attribute can be attached to the quantize/dequantize ops which are
// added by the quantization passes. These ops can be removed erased without
// losing accuracy.
inline constexpr char kVolatileOpAttrName[] = "volatile";

// Following attributes are used to mark ops that are not quantizable during

  auto flex_op = builder.create<TFL::CustomOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(), flex_op_name,
      CustomOptionForFlexOp(&builder, custom_option_buffer));
  op->replaceAllUsesWith(flex_op);
  op->erase();
  return true;
}

// Sets the "no_fallback" attribute.
Value SetNoFallbackAttr(PatternRewriter &rewriter, Value val) {
  val.getDefiningOp()->setAttr(kNoFallbackAttr, rewriter.getUnitAttr());
  return val;