// CHECK:            %[[HOST_OUTPUT:[0-9]*]] = "tf._XlaHostComputeMlir"(%[[A_OUTPUT]])
    // CHECK:            %[[D_OUTPUT:[0-9]*]] = "tf.D"(%[[HOST_OUTPUT]])
    // CHECK:            %[[HOST_OUTPUT2:[0-9]*]] = "tf._XlaHostComputeMlir"(%[[B_OUTPUT]])
    // CHECK:            %[[HOST_OUTPUT3:[0-9]*]] = "tf._XlaHostComputeMlir"()
    // CHECK:            %[[HOST_OUTPUT4:[0-9]*]] = "tf._XlaHostComputeMlir"(%[[HOST_OUTPUT3]], %[[HOST_OUTPUT2]])

}

// -----

// _XlaHostComputeMlir with manual_sharding should not fall back to
// XlaHostCompute, because XlaHostCompute does not support manual_sharding.
// Instead, it is skipped and the MlirXlaOpKernel is expected to handle it.

func.func @HostComputeManualNoFallback(%arg0: tensor<i32>) -> () {
  // CHECK: "tf._XlaHostComputeMlir"

  // 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

  if (op->getNumOperands() != func.getFunctionType().getNumInputs())
    return op.emitError()
           << "'host_func' has " << func.getFunctionType().getNumInputs()
           << " inputs and '_XlaHostComputeMlir' has " << op->getNumOperands()
           << " operands.  Number of operands/inputs should be the same.";

  if (op->getNumResults() != func.getFunctionType().getNumResults())

          %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>) -> ()

// Reshape should not be executed on TPU as all are marked by outside
// compilation. And there should be no host-device communication.
// CHECK: tf._TPUCompileMlir
// CHECK-NOT: tf.Reshape
// CHECK-NOT: tf._XlaHostComputeMlir

    %cst_34 = "tf.Const"() <{value = dense<[3, 1120, -1]> : tensor<3xi32>}> : () -> tensor<3xi32>
    %cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

// . So the TPU program should contain tf.Rank op and there should be no
// host-device communication.
// CHECK: tf._TPUCompileMlir
// CHECK-SAME: tf.Rank
// CHECK-NOT: tf._XlaHostComputeMlir
// CHECK-NOT: tf._XlaRecvAtHost
// CHECK-NOT: tf._XlaSendFromHost

}

// Create an `_XlaHostComputeMlir` for the map_outside_compilation case. Inputs
// are converted from split sharding to MANUAL sharding and outputs are
// converted from MANUAL sharding to split sharding. Set `common_split_sharding`
// if it has not yet been set. Output `full_outputs`, which is the outputs of
// the `_XlaHostComputeMlir` and add the
// `_XlaHostComputeMlir` to `host_compute_out_ops`.

        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>