func.func @main(%arg0: tensor<3x3x10xbf16>, %arg1: tensor<3xi32>) -> tensor<1x?x4xbf16> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "_arg0,_arg1,_arg2", outputs = "_retval0"}} {
      %cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>

  "tf_device.cluster"() ({

    // CHECK: "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f32>>}> : () -> tensor<complex<f32>>
    // CHECK-NOT: tf.ResourceApplyAdam

  %0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<1xi32>, tensor<1xf32>) -> (tensor<1xi32>, tensor<1xf32>)
  func.return %0#0, %0#1: tensor<1xi32>, tensor<1xf32>
}

// -----

// CHECK-LABEL: func @stopgradient
func.func @stopgradient(%arg0: tensor<1xi32>) -> tensor<1xi32> {
  // CHECK-NEXT:  return %arg0 : tensor<1xi32>
  %0 = "tf.StopGradient"(%arg0) : (tensor<1xi32>) -> tensor<1xi32>

  }
  func.func @partitioned_called_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @stateful_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }

    {
      ^bb0(%carg0: tensor<*xf32>, %carg1: tensor<i32>):
      %ne = "tf.NotEqual"(%carg1, %zero) : (tensor<i32>, tensor<i32>) -> tensor<i1>
      "tf.Yield"(%ne) : (tensor<i1>) -> ()
    },
    {
      // loop body
      ^bb0(%barg0: tensor<*xf32>, %barg1: tensor<i32>):
      %add = "tf.Add"(%barg0, %barg0) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>

  %0 = "tf.Bitcast"(%arg0) : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xf32>
  // CHECK: return %arg0
  func.return %0: tensor<8x16x32x64xf32>
}

// CHECK-LABEL: testDifferentBitcastType
func.func @testDifferentBitcastType(%arg0: tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> {

  // Build the XLA computation graph.
  // "arg0" (T=DT_RESOURCE), "arg1" (T=DT_INT32)
  // "arg0", "arg1" -> "if" (If) -> "ret0", "ret1"
  // "arg0", "arg1" -> "while" (While) -> "ret2", "ret3"
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output arg0 = ops::_Arg(s.WithOpName("arg0"), DT_RESOURCE, 0);
  Output arg1 = ops::_Arg(s.WithOpName("arg1"), DT_BOOL, 1);
  NameAttrList f;
  f.set_name("f1");

    %2:2 = "tf.WhileRegion"(%1, %arg0) ({
      ^bb0(%carg0: tensor<f32>, %carg1: tensor<i32>):
         %limit = arith.constant dense<5> : tensor<i32>
         %cond = "tf.NotEqual"(%carg1, %limit) : (tensor<i32>, tensor<i32>) -> tensor<i1>
         "tf.Yield"(%cond) : (tensor<i1>) -> ()
    },  {
      ^bb0(%barg0: tensor<f32>, %barg1: tensor<i32>):

#ifdef GOOS_windows
#define RARG0 CX
#define RARG1 DX
#define RARG2 R8
#define RARG3 R9
#else
#define RARG0 DI
#define RARG1 SI
#define RARG2 DX
#define RARG3 CX
#endif

// func runtime·domsanread(addr unsafe.Pointer, sz uintptr)
// Called from msanread.
TEXT	runtime·domsanread(SB), NOSPLIT, $0-16
	MOVQ	addr+0(FP), RARG0
	MOVQ	sz+8(FP), RARG1

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
  // CHECK: return %arg0
  func.return %0 : tensor<?x?xf32>
}

// -----

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<*xf32>
  // CHECK: "tfl.squeeze"