// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"
// CKECK-SAME: <{calibration_method = 5 : i32, id = "keeping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)

  %cst = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
  %cst_0 = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64>
  %fill = "tf.Fill"(%cst_0, %cst) : (tensor<1xi64>, tensor<f32>) -> tensor<1xf32>
  tf_device.replicate([%0, %fill] as %arg_r0: tensor<1xf32>) {n = 2 : i32} {
    %1 = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({

  // CHECK: "mhlo.all_reduce"
  // CHECK-SAME: channel_handle = #mhlo.channel_handle<handle = 10001, type = 1>
  // CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // CHECK: mhlo.add
  // CHECK: mhlo.return
  %0 = "tf.CollectiveReduceV2"(%input, %group_size, %group_key, %instance_key) {merge_op = "Add", final_op = "Id"} : (tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<f32>

  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>
    func.return %0 : tensor<?xi64>
  }

    %dimension = "tf.Const"() { value = dense<3> : tensor<1xi64> } : () -> tensor<1xi64>
    %6 = "tf.Sum"(%3, %dimension) { keep_dims = true }: (tensor<1x3x2x2xf32>, tensor<1xi64>) -> tensor<1x3x2x1xf32>
    return %5, %6 : tensor<1x3x2x2xf32>, tensor<1x3x2x1xf32>
  }

// CHECK-LABEL: func @conv_with_multiple_uses

func.func @broadcast_both_operand(%arg0: tensor<7xf32>, %arg1: tensor<5x1xf32>) -> tensor<5x7xf32> {
  %cst = arith.constant dense<[5, 7]> : tensor<2xi64>
  %0 = "tf.BroadcastTo"(%arg0, %cst) : (tensor<7xf32>, tensor<2xi64>) -> tensor<5x7xf32>
  %1 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<5x1xf32>, tensor<2xi64>) -> tensor<5x7xf32>
  %2 = "tf.Add"(%0, %1) : (tensor<5x7xf32>, tensor<5x7xf32>) -> tensor<5x7xf32>
  func.return %2 : tensor<5x7xf32>

  // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, 5, 21]> : tensor<3xi64>
  // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[21, 4]> : tensor<2xi64>
  // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg0, %[[cst]]) : (tensor<2x5x7x3xf32>, tensor<3xi64>) -> tensor<2x5x21xf32>
  // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg1, %[[cst_1]]) : (tensor<7x3x4xf32>, tensor<2xi64>) -> tensor<21x4xf32>

      %out, %ctl = tf_executor.island wraps "tf.Const"() {device = "", value = dense<[1]> : tensor<1xi64>} : () -> tensor<1xi64>
      // expected-error @+1 {{Validation failed for the initializer function: NoOp. All initializer function's fetches should be tf_executor::ControlType. Got: tensor<1xi64>.}}
      tf_executor.fetch %out : tensor<1xi64>
    }
    return
  }

  func.func @main(%arg0: tensor<10x19xf32>, %arg1: tensor<19x10xf32> {mhlo.is_same_data_across_replicas = true}) -> tensor<10x19xf32> {
    %0 = "tf.Shape"(%arg0) : (tensor<10x19xf32>) -> tensor<2xi64>
    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<10x19xf32>
    func.return %1 : tensor<10x19xf32>
  }
}

// Tests that foldable ops are constant-folded to enable legalization of ops

  %1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
  %2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
  // expected-error @+1 {{Node `mhlo.dynamic_iota` must have compile-time constant}}
  %3 = "mhlo.dynamic_iota"(%2) <{iota_dimension = 0 : i64}> : (tensor<1xi64>) -> tensor<?xi32>
  %4 = mhlo.multiply %3, %3 : tensor<?xi32>
  return %4 : tensor<?xi32>