func.func private @some_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
    return %arg0 : tensor<1x2x2x3xf32>
  }

  // CHECK-SKIP: func.func private @some_other_func
  func.func private @some_other_func(%arg0: tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32> attributes {tf._noinline = true} {
    return %arg0 : tensor<1x2x2x3xf32>
  }

// Check that we only outline the outermost partitioned call with
// `_xla_compile_device_type` in nested calls to a device cluster. The callee
// of any partitioned call outside of a device cluster is marked with
// `tf._noinline = true` to prevent it getting inlined to perseve runtime
// performance.
// CHECK-LABEL: func.func @nested_calls
func.func @nested_calls(%arg0: tensor<i32>) -> tensor<i32> attributes {tf.entry_function = {}} {

  func.return %1 : tensor<i32>
}

func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
  %1 = "tf.AddV2"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  %2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>

  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<f32>>>) -> tensor<f32>
  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<7xf32>, tensor<f32>) -> tensor<10xf32>
  func.return %2 : tensor<10xf32>

  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

func.func private @simple_callee_with_noinline() -> tensor<2xi32> attributes {tf._noinline = true} {
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

// Test that simple TF operations can be inlined.

// CHECK-LABEL: func @inline_simple(

    return %arg0 : tensor<1x2x2x3xf32>
  }
}

// CHECK: module
// CHECK: tf.StatefulPartitionedCall
// CHECK: func.func private @some_func
// CHECK-NOT: func.call

// -----

// This test makes sure functions without tf._noinline=true is inlined.

module {

  return %x : tensor<i32>
}

// -----

// Test for XlaLaunch

func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
  %1 = "tf.AddV2"(%arg0, %arg1) {__op_key = 0: i32} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %1 : tensor<1x3xf32>
}

// saved using the SaveV2 op.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertSaveOpPass();

// Creates a pass that marks functions with the attribute `tf._noinline = true`
// to avoid being inlined by the `InlinerPass`. `noinline_functions` is the name
// of the functions to mark.
std::unique_ptr<OperationPass<func::FuncOp>> CreateMarkFunctionsNoinlinePass(

  // CHECK-SAME: f = @some_other_func
  // CHECK: return

  func.func private @some_func(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<3x3xf32> attributes {tf._noinline = true} {
    %0 = stablehlo.add %arg0, %arg1 : tensor<3x3xf32>
    return %0 : tensor<3x3xf32>
  }
  // CHECK: func.func private @some_func
  // CHECK: tf.XlaCallModule

  attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x5x5x1024>], tf._noinline = true, tf._original_func_name = "__inference_main_540"} {
    %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
    %cst_1 = "tf.Const"() {value = dense<4> : tensor<i32>} : () -> tensor<i32>