}

// -----

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-LABEL: func @xla_host_compute_mlir_empty_module
  func.func @xla_host_compute_mlir_empty_module(%arg0: tensor<2xf32>) -> tensor<*xf32> {
    // 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>

  // Infers the shape CaseRegion outputs based on the embedded StableHLO module.
  // Returns true if a return type was changed.
  bool InferShapeForXlaCallModule(XlaCallModuleOp op);

  // Infers the shape of _XlaHostComputeMlir based on the host computation
  // module.  Returns true if a return type was changed.
  bool InferShapeForXlaHostComputeMlir(_XlaHostComputeMlirOp op);

  // Infers the shape of function attached to XlaHostCompute.

    tf_device.return
  },  {
    %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>

    return %0#1, %0#0 : tensor<?xi32>, tensor<?x!tf_type.string>
  }
  func.func private @_func() -> (tensor<?xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
    %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>
  }
}

  }];

  let arguments = (ins);

  let results = (outs
    TF_StrTensor:$program
  );
}

def TF__XlaHostComputeMlirOp : TF_Op<"_XlaHostComputeMlir", [TF_RecvSideEffect, TF_SendSideEffect, TF_XlaHostComputeSideEffect]> {
  let summary = [{
A pseudo-op to represent host-side computation in an XLA program.
  }];

  let arguments = (ins