// CHECK:         %[[ARG_TUPLE]] = (f32[10,19]{1,0}, f32[19,10]{1,0}) parameter(0), parameter_replication={false,true}
// CHECK:         [[ARG0:%.*]] = f32[10,19]{1,0} get-tuple-element((f32[10,19]{1,0}, f32[19,10]{1,0}) %[[ARG_TUPLE]]), index=0
// CHECK:         [[ARG1:%.*]] = f32[19,10]{1,0} get-tuple-element((f32[10,19]{1,0}, f32[19,10]{1,0}) %[[ARG_TUPLE]]), index=1

  %1:4 = "tf.CustomAggregator"(%arg0) <{id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> {min = 7.547870e-07 : f32, max = 0.999992311 : f32} : (tensor<1x1024xf32>) -> (tensor<1x1024xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)

  %0 = "tf.Const"() {value = dense<2.000000e+00> : tensor<f32>} : () -> tensor<f32>
  %1 = "tf.AddV2"(%0, %0) {device = "/job:worker/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %2 = "tf.AddV2"(%0, %0) {device = "/job:worker/replica:0/task:0/device:CPU:1"} : (tensor<f32>, tensor<f32>) -> tensor<f32>
  func.return
}

// CHECK: func @no_change_test
func.func @no_change_test() -> ()  {

      "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.

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

}

// 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_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.GuaranteeConst"(%arg1) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>
    %outputs_0, %control_1 = tf_executor.island wraps "tf.GuaranteeConst"(%arg2) {T = f32, device = ""} : (tensor<f32>) -> tensor<f32>

    %0 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {device = "", max = 2.000000e-01 : f32, min = -1.000000e-01 : f32, narrow_range = false, num_bits = 8 : i64} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>

    %2 = "quantfork.dcast"(%1) : (tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03, 5.000000e-03, 5.000000e-03}>>) -> tensor<4x3xf32>
    %3 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
    %4 = "quantfork.dcast"(%3) : (tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>) -> tensor<1x4xf32>

      %min_range = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
      %max_range = "tf.Const"() { value = dense<5.0> : tensor<f32> } : () -> tensor<f32>
      %0 = "tf.Dequantize"(%arg0, %min_range, %max_range) : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<f32>) -> tensor<1xf32>
      func.return %0 : tensor<1xf32>
    }
  })mlir";