// Erases the ops.
  for (auto op : erase_list) op->erase();
}

void RemoveVariablesInSessionInitializerPass::runOnOperation() {
  ModuleOp module_op = getOperation();

  for (auto init_func_op : GetInitializerFunctions(module_op)) {
    if (!init_func_op) return;

    if (init_func_op.getBlocks().size() != 1) {

                                             newConstOp);
  return success();
}

void ConvertConstPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();
  auto *context = &getContext();
  patterns.add<QuantizedConstRewrite>(context);
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

std::unique_ptr<OperationPass<func::FuncOp>>

};

// Folds tf.IdentityOp and tf.IdentityNOp if op device and the argument devices
// from the defining ops match.
void TensorDeviceCopyConversionPass::runOnOperation() {
  func::FuncOp func_op = getOperation();

  auto should_fold_op_func = [&func_op](const Value &arg,
                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and

    : public impl::BatchMatMulToEinsumPassBase<BatchMatMulToEinsumPass> {
  void runOnOperation() override;
};

void BatchMatMulToEinsumPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();

  patterns.add<ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulOp>,
               ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulV2Op>>(
      &getContext());

    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
  }

 private:
  void runOnOperation() override;
};

void QuantizePass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<StableHloQuantization, StableHloQuantizationReverse>(&ctx);

  PopulateCommonQuantizationPatterns(ctx, patterns,

}

bool IsMhloUniformQuantizedOp(Operation& op) {
  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
}

void VerifyQuantLegalization::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    // Verify all uq and qint types are lowered.
    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||

void ReduceTypePrecisionPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  patterns.add<CheckRangeAndConvertI8ToI4, SanitizeGatherOpOutputToI4>(
      &getContext());
  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateReduceTypePrecisionPass() {
  return std::make_unique<ReduceTypePrecisionPass>();
}

}  // namespace TFL

    if (!isa<TF::IdentityOp, TF::ReadVariableOp>(op)) break;
    v = op->getOperand(0);
  }
  return v;
}

void AnnotateParameterReplicationPass::runOnOperation() {
  ModuleOp m = getOperation();
  OpBuilder builder(m.getContext());
  m.walk([&](tf_device::ClusterFuncOp cluster_func) {
    auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>();
    if (!replicate) return;

          return op.getCallTargetName() != "stablehlo.composite";
        });

    RewritePatternSet patterns(context);
    patterns.add<ReplaceCustomCallWithComposite>(context);

    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

static PassRegistration<LegalizeStablehloCustomCallToCompositePass>

 public:
  explicit AnalyzeVariablesPass() {}
  void runOnOperation() override;
};

void AnalyzeVariablesPass::runOnOperation() {
  auto* context = &getContext();
  auto module = getOperation();
  bool legalize_to_tfl = true;

  module.walk([&](Operation* op) {
    // Skip ops that are supported natively by TFLite.
    if (IsSupportedTFLiteResourceOp(op)) return WalkResult::advance();