statistics_map.insert(single_map.statistics().begin(),
                          single_map.statistics().end());
  }
  return statistics_map;
}

absl::Status AddCalibrationStatistics(
    mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
    const CalibrationOptions& calibration_options,
    const PyFunctionLibrary& py_function_library) {
  TF_ASSIGN_OR_RETURN(const CalibrationStatisticsFlatMap statistics_map,

  for (Region& region : op->getRegions()) {
    if (!region.getOps<TF::CalibrationStatisticsSaverOp>().empty()) {
      return true;
    }
  }

  SymbolTable symbol_table(op->getParentOfType<ModuleOp>());
  // Check the functions associated to CaseOp, IfOp and WhileOp.
  for (const NamedAttribute& attr : op->getAttrs()) {
    FlatSymbolRefAttr symbol_attr =
        dyn_cast_or_null<FlatSymbolRefAttr>(attr.getValue());

  }

  void removeAttrMapAttribute(TF::PartitionedCallOp call_op,
                              StringRef function_name,
                              StringRef error_message) const {
    ModuleOp module = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    mlir::func::FuncOp composite_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    if (!composite_func) return;

    return success();
  }
};

template <typename T>
std::vector<T> GetAllOps(mlir::ModuleOp* module) {
  std::vector<T> ops;
  module->walk([&](T op) { ops.emplace_back(op); });
  return ops;
}

bool checkWhetherGraphHasValidStaticLookupTables(ModuleOp module) {
  auto hashtables = GetAllOps<TF::HashTableV2Op>(&module);
  // No needs to run the legalization patterns.

// Adds the first portion of StableHLO->TF passes happening before quantization.
// The `pass_manager` that runs on a `mlir::ModuleOp` expects a graph containing
// a `mlir::TF::XlaCallModuleOp` with serialized StableHLO module. The resulting
// `mlir::ModuleOp` after running these passes will be an MHLO module, or a
// StableHLO module if `pass_config.enable_stablehlo_quantizer` is `true`. This

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {
namespace {

class CastBf16OpsToF32Pass
    : public PassWrapper<CastBf16OpsToF32Pass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CastBf16OpsToF32Pass)
  explicit CastBf16OpsToF32Pass() = default;

  StringRef getArgument() const final {

limitations under the License.
==============================================================================*/

include "mlir/Pass/PassBase.td"

def EmbeddingPipeliningPass : Pass<"tf-embedding-pipelining", "mlir::ModuleOp"> {
  let summary = "Rewrite graph for embedding pipelining";
  let constructor = "TFDevice::CreateEmbeddingPipeliningPass()";
    let description = [{

                           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);

// tf.AssignVariableOp(tf.VarHandleOp, tf.Const) patterns in the initializer
// function and replaces tf.Consts with the results of RestoreV2.
class InsertRestoreOpPass
    : public PassWrapper<InsertRestoreOpPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertRestoreOpPass)

  explicit InsertRestoreOpPass() = default;

  // The argument used to refer to the pass in the textual format (e.g. on the

  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;