}
}

// -----

// Tests that no call_once op is added.

module attributes {tf_saved_model.semantics} {
  func.func @no_call_once(%arg0: tensor<i64> {tf_saved_model.index_path = ["x"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]})
  attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    func.return %arg0 : tensor<i64>

// -----

// Tests the case when there's a tf_saved_model.session_initializer and an empty init function.
module attributes {tf_saved_model.semantics} {

  "tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()
// CHECK: "tf_saved_model.session_initializer"()
// CHECK-SAME: initializers = [@init]

  constexpr absl::string_view kModuleCode = R"mlir(
    module attributes {tf_saved_model.semantics} {
      "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()

      func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["restore"], tf_saved_model.initializer_type = "restore_op"} {

// CHECK-LABEL: module @asset
module @asset attributes {tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()

  // CHECK-NOT: "tf_saved_model.asset"
  "tf_saved_model.asset"() {filename = "assets/test0.txt", sym_name = "asset0"} : () -> ()
  "tf_saved_model.asset"() {filename = "assets/test1.txt", sym_name = "asset1"} : () -> ()

  // CHECK: func @init()

  }
}

// -----

module attributes {tf_saved_model.semantics} {

  "tf_saved_model.asset"() {filename = "assets/table.txt", sym_name = "v"} : () -> ()
  "tf_saved_model.asset"() {filename = "assets/table2.txt", sym_name = "w"} : () -> ()

  // CHECK: func @f()
  func.func @f(%arg0: tensor<!tf_type.string> {tf_saved_model.bound_input = @v}, %arg1: tensor<!tf_type.string> {tf_saved_model.bound_input = @w})

// RestoreV2 op created for a single VarHandleOp.

module attributes {tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()

  func.func @init_func_restore_op() -> () attributes {
      tf_saved_model.initializer_type = "restore_op",
      tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"]} {

  func.return %2 : tensor<1x3xf32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]})
  attributes {tf_saved_model.exported_names = ["main"]} {
  // CHECK: tf.VarHandleOp

    module attributes {tf_saved_model.semantics} {
      "tf_saved_model.session_initializer"() <{initializers = []}> : () -> ()
      "tf_saved_model.asset"() <{filename = "assets/vocab_file.txt", sym_name = "__tf_saved_model_asset0_vocab_file.txt"}> : () -> ()

// Test that no ops are added to @main when @tf_quant__save function does
// not exist.

module attributes {tf_saved_model.semantics} {
  func.func @main(%arg: tensor<!tf_type.string> {tf_saved_model.index_path = ["__tf_file_prefix"]}) -> ()
      attributes {tf.entry_function = {inputs = "tf_file_prefix:0", outputs = ""}, tf_saved_model.exported_names = ["main"]} {
    tf_executor.graph {
      tf_executor.fetch
    }
    return
  }
}

# CHECK-SAME:             -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]})
# CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"]

# CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) <{{{.*}}}> {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>