// expected-remark@+1 {{Cost: 1}}
    %14 = "tf.Reshape"(%7, %13) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>, tensor<4xi32>) -> tensor<?x8x?x?xf32>
    // expected-remark@+1 {{Cost: 8}}
    func.return %11, %14 : tensor<4xi32>, tensor<?x8x?x?xf32>
}

// CHECK-LABEL: test_expensive_ops

module {
  func.func @wrap_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    %0 = "tf.PartitionedCall"(%arg0, %arg1) <{f = @composite_conv2d_with_relu6_fn}> {_tfl_quant_trait = "fully_quantizable"}
          : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  func.func @no_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {

    %r_max_without_zero = "tf.Max"(%input_1d, %dim) { keep_dims = true }: (tensor<?xf32>, tensor<1xi64>) -> tensor<1xf32>
    %r_max = "tf.Maximum"(%zero, %r_max_without_zero) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

    %r_min_without_zero = "tf.Min"(%input_1d, %dim) { keep_dims = true }: (tensor<?xf32>, tensor<1xi64>) -> tensor<1xf32>

      return %2 : tensor<?x2xf32>
    }
    func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {

    %1 = "tf.BiasAdd"(%0, %cst_0) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<3xf32>) -> tensor<*xf32>
    %2 = "tf.PartitionedCall"(%arg0, %cst_2) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
    func.return %1, %2: tensor<*xf32>, tensor<*xf32>

      %outputs_59, %control_60 = tf_executor.island wraps "tf.Reshape"(%outputs_57, %outputs_23) {_tpu_replicate = "cluster_sample_sequence", device = ""} : (tensor<*xf32>, tensor<4xi32>) -> tensor<*xf32>

func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
  %keys = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>

    func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> {
      %0 = "tf.Const"() {value = dense<1.42> : tensor<2xf32>} : () -> tensor<2xf32>
      %1 = "tf.Atan2"(%arg0, %0) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
      func.return %0 : tensor<2xf32>
    }
  })";

  TF_ASSERT_OK(LegalizeModule(kModuleWithConstParam));
}

      return %2 : tensor<?x2xf32>
    }
    func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
      return %arg0 : tensor<?x2xf32>
    }
  }
)mlir";

    %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %2 = "tf.Relu6"(%1) : (tensor<*xf32>) -> tensor<*xf32>
    func.return %2 : tensor<*xf32>
  }
  func.func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x2x2x3xf32>, %arg1: tensor<2x2x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {