using func::FuncOp;
using mlir::ModuleOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

  %begin0 = arith.constant dense<0> : tensor<4xi64>
  %shape0 = arith.constant dense<[2,3,4,4]> : tensor<4xi64>
  %begin1 = arith.constant dense<1> : tensor<4xi64>
  %shape1 = arith.constant dense<[1,2,3,4]> : tensor<4xi64>
  %0 = "tfl.slice"(%arg0, %begin0, %shape0) : (tensor<2x3x4x5xf32>, tensor<4xi64>, tensor<4xi64>) -> tensor<2x3x4x4xf32>

    func.func @main(%arg0: tensor<f32>) {
      %elem = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
      %shape = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64>
      %zero = "tf.Fill"(%shape, %elem) : (tensor<0xi64>, tensor<f32>) -> tensor<f32>
      tf_device.replicate([%arg0, %zero] as %arg1: tensor<f32>) {n = 2 : i32} {
        %arg1_id = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({

  %7 = "mhlo.broadcast_in_dim"(%3) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<5xf32>) -> tensor<32x28x28x5xf32>
  %8 = mhlo.add %6, %7 : tensor<32x28x28x5xf32>
  %9 = mhlo.maximum %8, %1 : tensor<32x28x28x5xf32>
  %10 = "mhlo.reshape"(%9) : (tensor<32x28x28x5xf32>) -> tensor<32x3920xf32>

using ::xla::XlaComputation;
using ::xla::XlaOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

  %2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
  func.return %1, %2 : tensor<i64>, tensor<*xi64>
}

}

// -----

module attributes {tf_saved_model.semantics} {

// Test hoisting const op.

// CHECK-LABEL: func @_tfrt_resource_init

//        limit_indices = dense<[4, 1]> : tensor<2xi64>,
//        start_indices = dense<0> : tensor<2xi64>,
//        strides = dense<1> : tensor<2xi64>} :
//        (tensor<4x6xf32>) -> tensor<4x1xf32>
// %1 = "mhlo.slice"(%input) {
//        limit_indices = dense<[4, 3]> : tensor<2xi64>,
//        start_indices = dense<[0, 1]> : tensor<2xi64>,
//        strides = dense<1> : tensor<2xi64>} :
//        (tensor<4x6xf32>) -> tensor<4x2xf32>

    %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi64>} : () -> tensor<3xi64>
    // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor<i32>) -> tensor<8x4x14x14x16xf32>

    "tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> ()
    "mhlo.return"(%1) : (tensor<f32>) -> ()
  }) {dimensions = dense<[1]> : tensor<1xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?xf32>
  func.return
}

// -----

// Tests transitive unsupported parent of TF/XLA communication op.

//     -> tensor<5x2xf32>
//
// is lowered to
//
//   %shape = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi64>}
//   %inp0 = "tf.Reshape"(%arg0, %shape)
//     : (tensor<2xf32>, tensor<2xi64>) -> tensor<1x2xf32>
//   %inp1 = "tf.Reshape"(%arg1, %shape)
//     : (tensor<2x2x2xf32>, tensor<2xi64>) -> tensor<4x2xf32>
//   %items0 = "tf.Unpack"(%[[INP0]]) {axis = 0 : i64}
//     : (tensor<1x2xf32>) -> tensor<2xf32>