// CHECK:      "mhlo.recv"([[INIT_TOKEN]])
  %0 = "tf._XlaHostComputeMlir"() {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : () -> tensor<i32>
  func.return
}

// -----

// Tests `tf._XlaHostComputeMlir` with no results simply forwards its token from
// the generated `mhlo.send`.

// CHECK-LABEL: func @host_compute_one_operand_no_results

          %control_4 = tf_executor.island wraps "tf._XlaHostComputeMlir"(%outputs_2) {host_mlir_module = "", manual_sharding = true, recv_key = "host_compute_channel_0_retvals", send_key = "host_compute_channel_0_args"} : (tensor<1x2xf32>) -> ()

        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 tokens[0];
  } else {
    return builder.create<AfterAllOp>(loc, original_token.getType(), tokens)
        .getResult();
  }
}

// Replaces `tf._XlaHostComputeMlir` with individual `mhlo.send` and `mhlo.recv`
// ops per operand and result. Unique Channel IDs are assigned per transfer.
// Sink tokens are created across all `mhlo.send` ops first and then by
// all `mhlo.recv` ops.