const TypeRange result_types,
                           const SetVector<Operation*>& reverse_subgraph,
                           const func::FuncOp stablehlo_func_op,
                           ModuleOp module_op) {
  MLIRContext* ctx = module_op.getContext();
  OpBuilder builder(ctx);
  Operation* last_subgraph_op = reverse_subgraph.front();
  builder.setInsertionPointAfter(last_subgraph_op);

  while (auto op = v.getDefiningOp()) {
    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;

TEST_F(PreCalibrationComponentTest,
       HasCustomAggregatorOpAndQuantizableFuncForSimpleDotGeneral) {
  PreCalibrationComponent component(ctx_.get());
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {} {
      func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> attributes {} {
        %0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>

          signature_def_map,
      absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases);

  // Runs the static-range post-training quantization (PTQ) on `module_op`.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // A non-owning `MLIRContext`. This `MLIRContext` should exceed the lifetime
  // of `StaticRangePtqComponent`.

#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {
namespace {

using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::tensorflow::monitoring::testing::CellReader;
using ::testing::Test;

static constexpr char kMetricsName[] =
    "/tensorflow/core/tf2xla/tf_quant_op_count";

class DecomposeTFOpsPass
    : public PassWrapper<DecomposeTFOpsPass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DecomposeTFOpsPass)

  explicit DecomposeTFOpsPass(std::optional<ModuleOp> external_tfr_module)
      : external_tfr_module_(external_tfr_module) {}

  StringRef getArgument() const final { return "tfr-decompose"; }

  StringRef getDescription() const final {

  if (Operation *nested_module = symbol_table.lookup(kNestedModule)) {
    nested_module->emitOpError("unexpected already present outlined module.");
    return signalPassFailure();
  }
  ModuleOp outlined_module = ModuleOp::create(getOperation().getLoc());
  outlined_module->setAttrs(getOperation()->getAttrDictionary());
  outlined_module->setAttr(SymbolTable::getSymbolAttrName(),

using ::tensorflow::quantization::RunPasses;

WeightOnlyPtqComponent::WeightOnlyPtqComponent(absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> WeightOnlyPtqComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(
      kName, /*add_passes_func=*/
      [&config](PassManager& pm) {

      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
      return %0 : tensor<1x3xf32>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kModuleWithCompositeDotGeneral);
  ASSERT_TRUE(module_op);

  // Create a pass manager with `SaveQuantizationReportInstrumentation` and

  return IsReplicatedGraph(graph, function_library);
}

bool IsSupportedByReplicatedBridge(mlir::ModuleOp module) {
  return IsReplicatedGraph(module);
}

bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
  bool has_tpu_partitioned_call = false;
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {