// -----

func.func @testXlaHostComputeMlir(%arg0: tensor<2xf32>) -> () {
  "tf._XlaHostComputeMlir"(%arg0) {send_key="", recv_key="", host_mlir_module=""} : (tensor<2xf32>) -> ()
  func.return
}

// -----

func.func @testXlaHostComputeMlir(%arg0: tensor<2xf32>) -> () {

    // CHECK: "tf._XlaHostComputeMlir"
    // CHECK-SAME: -> tensor<*xf32>
    %0 = "tf._XlaHostComputeMlir"(%arg0) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", tpu_core = 0, host_mlir_module = ""} : (tensor<2xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  // CHECK-LABEL: func @xla_host_compute_mlir_shape_inferred

    manual_inputs.push_back(in_manual);
  }

  // Create the _XlaHostComputeMlirOp
  auto host_compute = builder.create<mlir::TF::_XlaHostComputeMlirOp>(
      loc, shard_output_types, manual_inputs,
      /*send_key=*/builder.getStringAttr(args_communication_key),
      /*recv_key=*/builder.getStringAttr(retvals_communication_key),
      /*host_mlir_module=*/builder.getStringAttr(serialized_func_module),

    StrArrayAttr:$ancestors,
    TF_ShapeAttrArray:$shapes,
    OptionalAttr<SymbolRefAttr>:$shape_inference_graph,
    StrAttr:$key,
    DefaultValuedStrAttr<StrAttr, "">:$send_key,
    DefaultValuedStrAttr<StrAttr, "">:$recv_key,
    DefaultValuedOptionalAttr<I64Attr, "1000000">:$cost_estimate_ns,
    DefaultValuedOptionalAttr<I64Attr, "0">:$tpu_core
  );

  let results = (outs

  },  {
    %1 = "tf_device.cluster"() ( {
      %2 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
      %3 = "tf._XlaHostComputeMlir"(%2) {recv_key = "host_compute_channel_0_0_retvals", send_key = "host_compute_channel_0_0_args", tpu_core = 0 : i64} : (tensor<f32>) -> tensor<f32>
      %4 = "tf.AddV2"(%2, %3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
      tf_device.return %4 : tensor<f32>

    %0 = "tf.A"() : () -> tensor<?xi32>
    "tf._XlaHostComputeMlir"(%0) {host_mlir_module = "", recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"} : (tensor<?xi32>) -> ()
    %1 = "tf.C"() : () -> tensor<?xi32>
    return %1 : tensor<?xi32>
  }
}

// -----

// The following xla.OpSharding is used:

A pseudo-op to represent host-side computation in an XLA program.
  }];

  let arguments = (ins
    Arg<Variadic<TF_Tensor>, [{A list of tensors that will be sent to the host.}]>:$inputs,

    StrAttr:$send_key,
    StrAttr:$recv_key,
    DefaultValuedOptionalAttr<StrAttr, "\"\"">:$host_mlir_module,
    DefaultValuedOptionalAttr<BoolAttr, "false">:$manual_sharding
  );

  let results = (outs