if (output_shape.value()) {
      changed = RefineResultType(op.getOperation(), op.getResult(),
                                 output_shape.value());
      return changed;
    }
  }
  return changed;
}

bool ShapeInference::RefineWithInferTypeOpInterface(
    InferTypeOpInterface infer_ti) {
  Operation* op = infer_ti.getOperation();
  if (none_of(op->getResultTypes(), CanBeRefined)) {

  for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {
    lifted_op->erase();
  }

  return success();
}

void LiftHashTableOpsAsArgsPass::runOnOperation() {
  auto module_op = getOperation();

  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (failed(LiftHashTableOpsToArguments(module_op, func_op))) {
      signalPassFailure();
      return;
    }
  }
}

    this->allow_mutable_tensors = allow_mutable_tensors;
  }
  void runOnOperation() override;
};

void FreezeGlobalTensorsPass::runOnOperation() {
  auto module = getOperation();
  if (!tf_saved_model::HasTfSavedModelSemantics(module)) return;

  DataFlowSolver solver;
  solver.load<dataflow::DeadCodeAnalysis>();
  solver.load<dataflow::SparseConstantPropagation>();

      auto new_constant_op = builder.clone(*constant_op.getOperation());
      constant_op.getResult().getUses().begin()->assign(
          dyn_cast<mlir::stablehlo::ConstantOp>(new_constant_op));
    }
  }
}

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

  func::FuncOp main_func = FindMainFuncOp(module_op);

  ModuleOp module = getOperation();
  bool success = true;
  int num_metadata = 0;
  TF::TPUReplicateMetadataOp metadata;
  MetadataMap metadata_map;
  module.walk([&](TF::TPUReplicateMetadataOp meta) {
    ++num_metadata;
    metadata = meta;
    metadata_map[meta->getAttrOfType<StringAttr>(TF::kTpuReplicateAttr).str()] =
        meta;
  });

  getOperation().walk([&](mlir::Operation* op) {

  op_with_dialect_count_.clear();
  dialect_count_.clear();
  dialect_name_of_.clear();

  // Compute the operation statistics for the currently visited operation.
  total_ops_ = 0;

  getOperation()->walk([&](Operation *op) {
    auto op_with_dialect_name = op->getName().getStringRef();
    auto op_name = op->getName().stripDialect();
    auto dialect_name = op->getDialect()->getNamespace();

    return it->second;
  }
  return llvm::StringRef();
}

static PassRegistration<InsertQuantizedFunctionsPass> pass;

void InsertQuantizedFunctionsPass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symbol_table(module);

  std::unique_ptr<llvm::MemoryBuffer> mem_buffer;
  llvm::StringRef quantized_function_library =
      GetFunctionLibrary(quantization_method_, op_set_);

    pattern_list.add<SimplifyBroadcastReshape>(context);
    patterns = std::move(pattern_list);
    return success();
  }

  void runOnOperation() override {
    auto func = getOperation();
    if (failed(applyPatternsAndFoldGreedily(func, patterns)))
      signalPassFailure();
  }

  FrozenRewritePatternSet patterns;
};

}  // namespace

  MLIRContext context_;
  OwningOpRef<ModuleOp> module_;
};

TEST_F(LegalizationOpConfigTest, FailsWithExpectsLegalizationWithMlir) {
  TF_EXPECT_OK(CreateMlirModule());
  EXPECT_FALSE(IsOpLegalizedWithMlir(*module_->getOperation()));
}

TEST_F(LegalizationOpConfigTest, ExpectsFalseForNonMlirOps) {
  TF_EXPECT_OK(CreateMlirModule());
  TF_ASSERT_OK_AND_ASSIGN(FuncOp main, GetMain());

        /*op=*/op.getOperation(), /*indices=*/ArrayRef<unsigned>{0, 1},
        /*max_bcast_rank=*/6);
  }

  // Allows QI8 output when the operands have valid shapes, which are
  // broadcastable shapes up to four dimension or have same shapes.
  if (IsQI8Type(element_type)) {
    return VerifyOperandsHaveSameShapesOrBroadcastableShape(
        /*op=*/op.getOperation(), /*indices=*/ArrayRef<unsigned>{0, 1},