TPUHostComputationExpansionPass> {
  void runOnOperation() override;
};

void TPUHostComputationExpansionPass::runOnOperation() {
  OpBuilder builder(&getContext());
  getOperation().walk([&](tf_device::ClusterOp cluster) {
    ExpandHeadOutsideCompiledOps(cluster, &builder);
  });
}

}  // anonymous namespace

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

          UnfoldLargeSplatConstantPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnfoldLargeSplatConstantPass)

  void runOnOperation() override {
    auto module = getOperation();

    mlir::OpBuilder op_builder(&module.getBodyRegion());
    module.walk([&](mlir::arith::ConstantOp const_op) {
      MaybeUnfoldLargeSplatConstant(&op_builder, const_op);
    });
  }

 private:

    return success();
  }
};

void ConvertTpuModelToCpuPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  ModuleOp module_op = getOperation();

  patterns.add<ReplaceTpuPartitionedCallOpWithPartitionedCallOp,
               ReplaceBatchFunctionOpToPartitionedCallOp>(ctx);
  patterns.add<RemoveTpuOp>(ctx);
  patterns.add<RemoveIdentity>(ctx);

void ConvertCustomAggregationOpToQuantStatsPass::runOnOperation() {
  MLIRContext *ctx = &getContext();
  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<ConvertCustomAggregationOpToQuantStats>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError()

  void runOnOperation() override;
};

void OptimizeOpOrderPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();
  auto* ctx = func.getContext();
  patterns.add<PushDownDequantize>(ctx);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    signalPassFailure();
  }
}
}  // namespace

    }

    llvm::SourceMgr source_mgr;
    source_mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    auto module = parseSourceFile<ModuleOp>(source_mgr, context);
    if (module->getOperation() == nullptr) {
      return absl::AbortedError("Failed to parse input MLIR model.");
    }

    return quant::stablehlo::ImportedMlirModuleOp(std::move(module), nullptr);
  }

    }
    return_op->setOperand(0, new_result);
    rewriter.eraseOp(req_op);

    return success();
  }
};

void MergeFusionWithDequantizePass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext* ctx = module_op.getContext();
  RewritePatternSet patterns(ctx);
  patterns.add<MergeFusionWithUniformDequantizePattern>(ctx);
  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {

    emitError(op->getLoc()) << "Could not legalize op: " << op->getName();
    return false;
  }
  return true;
}

void VerifyTFXLALegalization::runOnOperation() {
  Operation* func_op = getOperation();
  ConversionTarget default_conversion_target =
      GetDefaultLegalConversionTargets(getContext(), legalize_chlo_);

  bool has_invalid_ops = false;
  func_op->walk([&](Operation* op) {

}

inline bool IsJaxRandomNormal(mlir::func::FuncOp func) {
  return func.getName().contains("tfl_wrapped_jax_random_normal");
}

void LegalizeJaxRandomPass::runOnOperation() {
  auto func = getOperation();
  if (!IsJaxRandomUniform(func) && !IsJaxRandomNormal(func)) return;
  auto result_tuple_ty =
      mlir::dyn_cast_or_null<TupleType>(func.getFunctionType().getResult(0));
  if (!result_tuple_ty) return;

                                expressedType, fqOp.getIsSigned());
  }
};

void ConvertSimulatedQuantPass::runOnOperation() {
  bool hadFailure = false;
  auto func = getOperation();
  RewritePatternSet patterns(func.getContext());
  auto *ctx = func.getContext();
  patterns.add<ConstFakeQuantRewrite, ConstFakeQuantPerAxisRewrite>(
      ctx, &hadFailure);