%shape1 = arith.constant dense<[64]> : tensor<1xi32>
  %0 = "tfl.reshape"(%arg0, %shape0) : (tensor<4x4x4xf32>, tensor<2xi32>) -> tensor<16x4xf32>
  %1 = "tfl.reshape"(%0, %shape1) : (tensor<16x4xf32>, tensor<1xi32>) -> tensor<64xf32>
  %2 = "tfl.reshape"(%0, %shape1) : (tensor<16x4xf32>, tensor<1xi32>) -> tensor<64xf32>
  %3 = arith.addf %1, %2 : tensor<64xf32>
  func.return %3 : tensor<64xf32>

func.func private @simple_callee(%arg0: tensor<2xi32>) -> tensor<2xi32> {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<2xi32>
  }
  func.return %graph#0 : tensor<2xi32>
}

// Test with a TPUPartitionedCallOp that is not inlined.
func.func @dont_inline_tpu_partitioned_call(%arg0: tensor<2xi32> {tf._user_specified_name = "inputs_0"}, %arg1: tensor<2xi32> {tf._user_specified_name = "inputs_1"}) -> tensor<2xi32> {

//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ]


module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> {
  %0 = "tf.Div"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}

// Inline a function that contains only tfl ops.
func.func @func_with_tfl_ops(%arg0 : tensor<2xi32>) -> tensor<2xi32> {
  %0 = "tfl.sub"(%arg0, %arg0) {fused_activation_function = "RELU6"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  %1 = "tfl.add"(%0, %arg0) {fused_activation_function = "RELU6"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %1: tensor<2xi32>
}

// CHECK-LABEL: func @inline_with_arg

        %5:2 = "tf.C"(%4) : (tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>)
        tf_device.return %5#0, %5#1 : tensor<2xi32>, tensor<2xi32>
      }) {num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> (tensor<2xi32>, tensor<2xi32>)
      tf_device.return %2#1 : tensor<2xi32>
    }
    func.return %1 : tensor<2xi32>
  }

// CHECK-DAG: %[[padding_rank_1:.*]] = "tf.Concat"({{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}) : (tensor<i32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<8xi32>

  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %cst_0 = arith.constant dense<4> : tensor<2x2xi32>
  %0 = "tf.SpaceToBatchND"(%arg0, %cst, %cst_0) : (tensor<1x128x128x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x68x68x3xf32>
  %1 = "tf.SpaceToBatchND"(%arg0, %cst, %cst_0) : (tensor<1x128x128x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x68x68x3xf32>

module {
func.func @main(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>, %arg2: tensor<2xi32>) -> tensor<2xi32> {
  %0 = "tfl.custom"(%arg0, %arg1, %arg2) {custom_code = "stablehlo.clamp", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}
}

// CHECK:       module {

func.func @modified(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x401408xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input", outputs = "output"}} {
  %cst = arith.constant dense<[1, 401408]> : tensor<2xi32>
  %0 = "tfl.quantize"(%arg0) {qtype = tensor<1x224x224x3x!quant.uniform<i8:f32, 7.812500e-03>>} : (tensor<1x224x224x3xf32>) -> tensor<1x224x224x3x!quant.uniform<i8:f32, 7.812500e-03>>

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

// CHECK:       module {
// CHECK-NEXT:  func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
// CHECK-NEXT:    %0 = stablehlo.add %arg0, %arg0 : tensor<2xi32>
// CHECK-NEXT:    %1 = stablehlo.subtract %0, %arg0 : tensor<2xi32>
// CHECK-NEXT:    return %1 : tensor<2xi32>