func.func @serving_default(%arg0: tensor<1x32x32x128xf32> {tf_saved_model.index_path = ["a"]}) -> (tensor<1x64x64x128xf32> {tf_saved_model.index_path = ["b"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "c:0", outputs = "d:0"}, tf_saved_model.exported_names = ["serving_default"]} {
      %0  = "tf.Const"() {value = dense<[56, 904]> : tensor<2xi32>} : () -> tensor<2xi32>

  // LINT.IfChange(savedmodel_to_stablehlo)
  m.def(
      "savedmodel_to_stablehlo",
      [](absl::string_view input_path,
         const std::vector<std::string>& exported_model_signatures =
             {"serving_default"},
         const std::vector<std::string>& tag_names = {"serve"},
         absl::string_view input_arg_shapes_str = "") -> py::bytes {
        auto module_bytecode =

// CHECK: return %[[CUSTOM_AGGREGATOR_1]] : tensor<1x3xf32>
// CHECK: }
// CHECK: }

// -----

// Tests that stablehlo op serialization also works for the "serving_default"
// function.

func.func @serving_default(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
  %0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>

  TF_SignatureDefFunction* serving_default =
      TF_GetSavedModelSignatureDefFunction(saved_model, "serving_default",
                                           status);
  EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);

  TF_SignatureDefFunctionMetadata* metadata =
      TF_SignatureDefFunctionGetMetadata(serving_default);

  const TF_SignatureDefParamList* args =

  func.func @serving_default(%arg0: tensor<1x4xf32> {tf_saved_model.index_path = ["x"]}) -> (tensor<1x3xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_x:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {

# ==============================================================================

# LINT.IfChange(savedmodel_to_stablehlo)
def savedmodel_to_stablehlo(
    input_path: str,
    exported_model_signatures: list[str] = ["serving_default"],
    tag_names: list[str] = ["serve"],
    input_arg_shapes_str: str = "",
) -> bytes: ...
# LINT.ThenChange()

# LINT.IfChange(tensorflow_module_to_stablehlo)
def tensorflow_module_to_stablehlo(

        exported_model_signatures=["serving_default"],
        tag_names=["serve"],
        input_arg_shapes_str="1,28,28,3::32"
)

```

#### Arguments:

* `input_path` (required): Path to your SavedModel directory.
* `exported_model_signatures` (optional): List of signature names to convert.
                                          Defaults to ["serving_default"].

      self, output_saved_model_path: str
  ) -> str:
    """Extracts the first XlaCallModule op from output saved model to string."""
    root = load.load(output_saved_model_path)
    tf_graph_def = root.signatures['serving_default'].graph.as_graph_def()
    for function in tf_graph_def.library.function:
      for node_def in function.node_def:
        if node_def.op == 'XlaCallModule':
          with ir.Context() as context:

opt<std::string> exported_model_signatures(
    "exported-model-signatures",
    llvm::cl::desc(
        "Comma-separated list of exported model signatures to convert"),
    llvm::cl::Optional, llvm::cl::init("serving_default"));

// NOLINTNEXTLINE
opt<std::string> tag_names(
    "tags",
    llvm::cl::desc("Comma-separated list of tags for loading SavedModel. "
                   "Ignored for MLIR input"),

// NOLINTNEXTLINE
opt<std::string> exported_model_signatures(
    "exported_model_signatures", llvm::cl::desc("model signature names"),
    llvm::cl::Optional, llvm::cl::init("serving_default"));

namespace mlir {
namespace odml {

absl::StatusOr<OwningOpRef<mlir::ModuleOp>> ImportSavedModelOrMLIR(
    const std::string& input_path, MLIRContext* context,
    llvm::SourceMgr* source_mgr,