weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);
    op->setOperand(weight_operand_idx, reshape_op);

    // Create a new function with preprocessed types.
    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    OperandRange func_args = op.getArgs();

        %mul = "tf.Mul"(%cast, %scale) : (tensor<1024x24x24x3xf32>, tensor<f32>) -> tensor<1024x24x24x3xf32>
        func.return %mul : tensor<1024x24x24x3xf32>
      }
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
  const auto test_func =
      module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
  ASSERT_THAT(test_func, NotNull());

// the value `checkpoint_path`.
// Also adds SessionInitializer op.
absl::Status AddSessionInitializerAndInlineCheckpoint(
    ModuleOp module, absl::string_view checkpoint_path);

// Creates a pass that sinks SavedModel asset filenames to constants.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateAssetSinkingPass(
    llvm::StringRef saved_model_dir);

}  // namespace tf_saved_model
}  // namespace mlir

  }
};

}  // namespace tensorflow

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
tensorflow::CreateTensorFlowGraphOptimizationPass(
    std::vector<tensorflow::GraphOptimizationPass*> tf_passes) {
  return std::make_unique<GraphOptPass>(std::move(tf_passes));
}

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
tensorflow::CreateTensorFlowGraphOptimizationPass(
    const std::vector<std::string>& pass_names) {

    // and use it for `shape_inference_graph` attribute on XlaHostCompute.
    func::FuncOp cloned_func;
    SymbolTable manager(op->getParentOfType<ModuleOp>());
    StringRef host_module = op.getHostMlirModule();
    if (!host_module.empty()) {
      mlir::OwningOpRef<mlir::ModuleOp> module_for_func;

      func::FuncOp func = op.GetHostFunc(&module_for_func);

      OpBuilder::InsertionGuard guard(rewriter);

                    mlir::ModuleOp module);

// Lower TF to MHLO and insert HLO into the XlaBuilder. xla_params are HLO-level
// inputs to module_op that have already been added to the XlaBuilder. returns
// are the returned XlaOps.
ABSL_DEPRECATED("Use v2/legalize_tf.h::LegalizeMlirToHlo instead.")
Status BuildHloFromTf(mlir::ModuleOp module_op, xla::XlaBuilder& builder,

    Pass<"tf-resource-op-lifting-for-main-function", "ModuleOp"> {
  let summary = "Lifting resource operations out of control flow statements "
    "for the main function";
  let constructor = "TFDevice::CreateResourceOpLiftingForMainFunctionPass()";
}

def AnnotateParameterReplicationPass :
    Pass<"tf-annotate-parameter-replication", "ModuleOp"> {

// rewrite the graph to use pipelining for better device utilization.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingSequencingPass();

// This is a strictly sequential and formally correct fallback option for the
// embedding pipelining pass intended for debugging during pipelining
// development.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingPipeliningPass();

}

void ClusterOutliningPass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symbol_table(module);
  OpBuilder builder(module.getContext());
  module.walk([&](tf_device::ClusterOp cluster) {
    OutlineCluster(cluster, &symbol_table, &builder);
  });
}

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

    mlir::ModuleOp module_op, mlir::MLIRContext* context,
    std::optional<Session*> session, bool run_tf_to_stablehlo,
    bool deserialize_xla_call_module,
    llvm::ArrayRef<llvm::ArrayRef<int64_t>> input_arg_shapes = {});

// Overload of `PreprocessAndFreezeGraph` that uses the default MLIR dump file
// prefix.
inline absl::Status PreprocessAndFreezeGraph(mlir::ModuleOp module_op,