// CHECK:      %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME:   dot_dimension_numbers = #mhlo.dot<
// CHECK-SAME:     lhs_batching_dimensions = [0],
// CHECK-SAME:     rhs_batching_dimensions = [0],
// CHECK-SAME:     lhs_contracting_dimensions = [3],
// CHECK-SAME:     rhs_contracting_dimensions = [1]
// CHECK-SAME: >}> : (tensor<3x72x1x2048xf32>, tensor<3x2048x512xf32>) -> tensor<3x72x1x512xf32>

  // CHECK:  "tf.BatchFunction"
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64
  // CHECK-SAME: num_batch_threads = 2 : i64
  // CHECK-SAME: shared_name = "batch/"

  } : (tensor<1x4xf32>) -> (tensor<1x8xf32>)
  func.return %0 : tensor<1x8xf32>

  //       CHECK: %0 = "tfl.custom"(%arg0) <{
  //  CHECK-SAME:   custom_code = "custom_call.my_custom_op",
  //  CHECK-SAME:   custom_option = #tfl<const_bytes : "0x746869735F69735F615F746573745F737472696E67">
  //  CHECK-SAME: }> : (tensor<1x4xf32>) -> tensor<1x8xf32>
}

// CHECK-LABEL: mhlo_custom_call_test__dont_legalize_dict_backend_config

// CHECK:           "mhlo.dot_general"({{.*}}, {{.*}}) <{
// CHECK-SAME:        lhs_batching_dimensions = [0]
// CHECK-SAME:        rhs_batching_dimensions = [0]
// CHECK-SAME:        lhs_contracting_dimensions = [2]
// CHECK-SAME:        rhs_contracting_dimensions = [1]

// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1}>
// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu_fn_1}
// CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_conv2d_with_bias_fn_1}

  // CHECK-SAME: %arg0,
  // CHECK-SAME: %[[DQ_1]], %[[DQ_2]], %[[DQ_3]], %[[DQ_4]],
  // CHECK-SAME: %[[DQ_5]], %[[DQ_6]], %[[DQ_7]], %[[DQ_8]],
  // CHECK-SAME: %[[NONE]], %[[NONE]], %[[NONE]],
  // CHECK-SAME: %[[DQ_9]], %[[DQ_10]], %[[DQ_9]], %[[DQ_9]],
  // CHECK-SAME: %[[NONE]], %[[NONE]],
  // CHECK-SAME: %[[DQ_11]], %[[DQ_12]],
  // CHECK-SAME: %[[NONE]], %[[NONE]], %[[NONE]], %[[NONE]]) <{

  // CHECK-SAME: Tin = [f32]
  // CHECK-SAME: Tout = [f32]
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false
  // CHECK-SAME: max_batch_size = 6 : i64
  // CHECK-SAME: max_enqueued_batches = 10 : i64

  // CHECK:      "mhlo.gather"({{.*}}) <{
  // CHECK-SAME:   dimension_numbers =
  // CHECK-SAME:     offset_dims = [1, 2]
  // CHECK-SAME:     collapsed_slice_dims = [0]
  // CHECK-SAME:     start_index_map = [0]
  // CHECK-SAME:     index_vector_dim = 1
  // CHECK-SAME:   indices_are_sorted = false
  // CHECK-SAME:   slice_sizes = dense<[1, 4, 128]>
  // CHECK-SAME: (tensor<2x4x128xf32>, tensor<2x1xi32>) -> tensor<2x4x128xf32>

  // CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
  // CHECK-SAME: AddV2
  // CHECK-SAME: op_key = 1
  // CHECK-NOT: __op_key
  %a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
  // CHECK-SAME: Abs
  // CHECK-SAME: op_key = 2
  // CHECK-NOT: __op_key

  }
}

// -----

module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: composite_and_select
  // CHECK-SAME: %[[ARG0:.*]]: tensor<1x2xf32>
  // CHECK-SAME: %[[ARG1:.*]]: tensor<2x3xf32>
  // CHECK-SAME: %[[ARG2:.*]]: tensor<1x3xi1>
  // CHECK-SAME: %[[ARG3:.*]]: tensor<1x3xf32>