constexpr StringRef kPlatformTpu = "TPU";

class ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass
    : public impl::
          ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPassBase<
              ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass)

                           .getValue();
  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);

  // We should not unwrap if the function is not from
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass.
  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {
    return;
  }

  MLIRContext* context = call_op.getContext();
  OpBuilder builder(context);

void AddCallModuleSerializationPasses(OpPassManager& pm) {
  AddShapeLegalizationPasses(pm);
  pm.addPass(createReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass());
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass may create
  // duplicate constants. Add canonicalizer to deduplicate.
  pm.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
  pm.addPass(TF::CreateXlaCallModuleSerializationPass());
}

  }];
  let dependentDialects = [
      "mlir::func::FuncDialect",
      "mlir::stablehlo::StablehloDialect",
      "TF::TensorFlowDialect",
  ];
}

def ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass : Pass<"stablehlo-replace-stablehlo-ops-in-main-function-with-xla-call-module-ops", "mlir::ModuleOp"> {
  let summary = "Replaces the StableHLO ops with a separate XlaCallModuleOps.";
  let description = [{