"tf.Yield"(%recv) : (tensor<f32>) -> ()
    }, {
      %add = "tf.Add"(%arg1, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
      "tf.Yield"(%add) : (tensor<f32>) -> ()
    }) { is_stateless = true}: (tensor<i1>) -> tensor<f32>
  // CHECK: _xla_token_input_nodes = ["_xla_token_arg_node"]

  func.return %0, %1 : tensor<f32>, tensor<f32>
}


// Next four functions are used to verify handling of a call chain.

    %cst_1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
    %cst_2 = "tf.Const"() {value = dense<0.00235294132> : tensor<f32>} : () -> tensor<f32>
    %cst_3 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
    %cst_4 = "tf.Const"() {value = dense<0.00117647066> : tensor<f32>} : () -> tensor<f32>
    %cst_5 = "tf.Const"() {value = dense<-43> : tensor<i32>} : () -> tensor<i32>

  %8 = "tf.Identity"(%5) {_xla_compile_device_type = "TPU", _replication_info = "c"} : (tensor<f32>) -> tensor<f32>
  %9:2 = "tf.TPUReplicatedOutput"(%6) : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
  %10:2 = "tf.TPUReplicatedOutput"(%7) : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
  %11:2 = "tf.TPUReplicatedOutput"(%8) : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
  func.return
}

// CHECK: %[[XLA_CALL_MODULE:.+]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[CST]])

// MIN-MAX-CHECK-NEXT:  [[add:%.*]] = "tf.PartitionedCall"([[lhs]], [[rhs]])
// MIN-MAX-CHECK-NEXT:  [[res:%.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"([[add]]) <{calibration_method = 1 : i32, id = "composite_conv2d_with_relu6_fn_calibration_method_1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32}> : (tensor<*xf32>) -> (tensor<*xf32>, tensor<f32>, tensor<f32>, tensor<0xi64>)

func.func @quant_raw_data(%arg0: tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<1x10x!quant.uniform<i8:f32, 0.2:2>> {
  %0 = "tfr.cast"(%arg0) : (tensor<1x10x!quant.uniform<i8:f32, 0.1:1>>) -> !tfr.tensor
  %1 = tfr.quant_raw_data(%0) : (!tfr.tensor) -> !tfr.tensor
  %2 = tfr.call @tf__risc(%1) : (!tfr.tensor) -> !tfr.tensor
  %3 = "tfr.cast"(%2) : (!tfr.tensor) -> tensor<1x10x!quant.uniform<i8:f32, 0.3:3>>

}

// CHECK-LABEL: not_fold_cast_fft_to_rfft
func.func @not_fold_cast_fft_to_rfft(%arg0: tensor<10x20x30xcomplex<f64>>) -> tensor<10x20x30xcomplex<f32>> {
  %0 = "tf.Cast"(%arg0) : (tensor<10x20x30xcomplex<f64>>) -> tensor<10x20x30xcomplex<f32>>
  %1 = "tf.FFT"(%0) : (tensor<10x20x30xcomplex<f32>>) -> tensor<10x20x30xcomplex<f32>>
  func.return %1: tensor<10x20x30xcomplex<f32>>

  %0 = "tf.Conj"(%arg0) : (tensor<complex<f32>>) -> tensor<complex<f32>>
  %1 = "tf.Conj"(%0) : (tensor<complex<f32>>) -> tensor<complex<f32>>
  // CHECK: return [[ARG0]]
  func.return %1: tensor<complex<f32>>
}

// CHECK-LABEL: func @testTripleConj
// CHECK-SAME:  ([[ARG0:%.+]]: tensor<complex<f32>>)
func.func @testTripleConj(%arg0: tensor<complex<f32>>) -> tensor<complex<f32>> {

  func.return %arg0 : tensor<!tf_type.resource>
}
func.func @if_else(%arg0: tensor<!tf_type.resource>, %arg1: tensor<!tf_type.resource>) -> tensor<!tf_type.resource> {
  %elem = "tf._SomeOp"() : () -> tensor<f32>
  %push = "tf.StackPushV2"(%arg1, %elem) {swap_memory = false} : (tensor<!tf_type.resource>, tensor<f32>) -> tensor<f32>
  func.return %arg1 : tensor<!tf_type.resource>
}

// -----

// CHECK:       ENTRY %main.{{[0-9]+}} ([[ARG0:.*]]: f32[2], [[ARG1:.*]]: f32[2]) -> (f32[2]) {
// CHECK-NEXT:    %[[ARG1]] = f32[2]{0} parameter(1)
// CHECK-NEXT:    %[[ARG0]] = f32[2]{0} parameter(0)
// CHECK-NEXT:    ROOT %tuple.{{[0-9]+}} = (f32[2]{0}) tuple(f32[2]{0} %[[ARG0]])
// CHECK-NEXT:  }

// CHECK:       // InputMapping {0, 1}
// CHECK-NEXT:  // XlaInputShape f32[2]
// CHECK-NEXT:  // XlaInputShape f32[2]
// CHECK-NEXT:  // XlaOutputShape (f32[2])