return absl::OkStatus();
}

// Converts arg_shapes to xla::Shape's and store into xla_input_shapes.
Status GetXlaInputShapes(
    mlir::ModuleOp module, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
    bool use_tuple_args,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    std::vector<xla::Shape>* xla_input_shapes) {

// Legalizes TF/XLA communication ops (TF dialect) to HLO dialect communication
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();

// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();

// Verifies that the TF/XLA ops have all been lowered to MHLO.

    }
  }
  return success();
}

LogicalResult SessionInitializerOp::verify() {
  SessionInitializerOp session_initializer = *this;
  mlir::SymbolTable symbol_table(
      session_initializer->getParentOfType<ModuleOp>());

  for (auto sym_ref : session_initializer.getInitializers()) {
    auto init_func_op = symbol_table.lookup<mlir::func::FuncOp>(
        mlir::cast<FlatSymbolRefAttr>(sym_ref).getValue());

    if (!init_func_op)

    mlir::ModuleOp module, const GraphExportConfig& configs);

// Converts an MLIR module to TensorFlow graph and FunctionLibraryDefinition.
// The "main" function of the module is stored in the graph and the rest of
// functions are stored in the library. Control ret nodes are stored separately
// in `control_ret_nodes`.
Status ConvertMlirToGraph(mlir::ModuleOp module,

#include "mlir/IR/OpDefinition.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/core/platform/errors.h"

namespace tensorflow {

void PopulateTfVersions(mlir::ModuleOp module, const VersionDef& versions) {
  mlir::Builder b(module.getContext());
  auto producer =
      b.getNamedAttr("producer", b.getI32IntegerAttr(versions.producer()));
  auto min_consumer = b.getNamedAttr(

  static constexpr absl::string_view kName = "quant_ptq_weight_only";

  explicit WeightOnlyPtqComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  absl::Nonnull<MLIRContext*> ctx_;
};

// Runs weight-only quantization on a SavedModel at

// quantization summary from a quantized `ModuleOp` and emit it in a human- and
// machine-readable format.
class QuantizationReport {
 public:
  QuantizationReport() = default;

  // Initializes `QuantizationReport` by collecting `QuantizationResults` from
  // `module_op`.
  explicit QuantizationReport(ModuleOp module_op);

  // Adds a `QuantizationResult` to the report.

  func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;
  context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}

TEST(HasTPUPartitionedCallOpInModule, HasTPUPartitionedCallModule) {

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"

namespace stablehlo::quantization {

void DisableDebugging(mlir::ModuleOp module_op) {
  module_op.walk(
      [](mlir::TF::DumpTensorOp dump_op) { dump_op.setEnabled(false); });
}

void ChangeToQuantizedFilename(mlir::ModuleOp module_op) {
  module_op.walk([](mlir::TF::DumpTensorOp dump_op) {
    dump_op.setFileName("quantized_tensor_data.pb");
  });
}

                                       /*RequiresNullTerminator=*/false);

  llvm::SourceMgr source_mgr;
  source_mgr.AddNewSourceBuffer(std::move(mem_buffer), llvm::SMLoc());
  OwningOpRef<ModuleOp> module_ref =
      parseSourceFile<ModuleOp>(source_mgr, module.getContext());
  // Inline and optimize loaded functions.
  MLIRContext* context = &getContext();
  PassManager pm(context);
  pm.addPass(createInlinerPass());