// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

module {
func.func @serving_default(%arg0: tensor<3x2xf32>) -> tensor<3x2xf32> attributes {tf.entry_function = {inputs = "serving_default_x", outputs = "outputs"}} {
// CHECK:       {

// CHECK-LABEL:   version: 3,

// CHECK-LABEL:   operator_codes: [ {
// CHECK:           version: 1
// CHECK:         } ],

// CHECK-LABEL:   subgraphs: [ {

      tf_executor.fetch %out_1 : tensor<*xi64>
    }
    return %0 : tensor<*xi64>
  }
// Sanity check: The contents of @serving_default is untouched.
// CHECK: func.func private @serving_default(%[[ARG_0:.*]]: tensor<?x!tf_type.string>) -> tensor<*xi64>
// CHECK-NEXT: %[[RES:.*]] = tf_executor.graph
// CHECK: %[[OUT:.*]], %[[CTL:.*]] = tf_executor.island wraps "tf.Const"()

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

module {
func.func @serving_default(%arg0: tensor<*xf32>) -> tensor<*xf32> attributes {tf.entry_function = {inputs = "serving_default_x", outputs = "outputs"}} {
// CHECK:     {
// CHECK-NEXT:  version: 3,
// CHECK-NEXT:  operator_codes: [ {
// CHECK-NEXT:    version: 1
// CHECK-NEXT:  } ],
// CHECK-NEXT:  subgraphs: [ {
// CHECK-NEXT:    tensors: [ {

  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"]} {

    return %2 : tensor<1x3x2x2xf32>
  }

  func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {
    %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @conv2d_with_inliner}
        : (tensor<1x3x4x3xf32>) -> tensor<1x3x2x2xf32>
    return %0 : tensor<1x3x2x2xf32>
  }

// CHECK-LABEL: func @serving_default
// CHECK: "tf.PartitionedCall"

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

# 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(

// CHECK-NEXT:      tensor_index: 5
// CHECK-NEXT:    }, {
// CHECK-NEXT:      name: "start_logits",
// CHECK-NEXT:      tensor_index: 6
// CHECK-NEXT:    } ],
// CHECK-NEXT:    signature_key: "serving_default"
// CHECK-NEXT:  } ]
// CHECK-NEXT:}
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 554 : i32}, tf_saved_model.semantics} {

        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"].

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

namespace mlir::quant {

// Returns a public `func::FuncOp` in `module_op` whose name matches either
// `main` or `serving_default`. If `func::FuncOps` with both names exist, the
// function with name "main" takes precedence. Returns null if no such a
// function exists.
func::FuncOp FindMainFuncOp(ModuleOp module_op);

}  // namespace mlir::quant

    tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
          path_map={'serving_default': '/tmp/representative_dataset_path'}
      ).save({'serving_default': representative_dataset})
  )

  # Using in QuantizationOptions.
  quantization_options = tf.quantization.experimental.QuantizationOptions(
      signature_keys=['serving_default'],
      representative_datasets=dataset_file_map,
  )