// CHECK:      [[RECV_SINK_TOKEN:%.*]] = mhlo.after_all [[RECV_RETVAL0_TUPLE]]#1, [[RECV_RETVAL1_TUPLE]]#1
  %0:2 = "tf._XlaHostComputeMlir"(%arg0, %arg1) {recv_key = "host_compute_channel_recv", send_key = "host_compute_channel_send", host_mlir_module = ""} : (tensor<i32>, tensor<i64>) -> (tensor<f32>, tensor<f64>)

  // CHECK:      return [[RECV_RETVAL0_TUPLE]]#0, [[RECV_RETVAL1_TUPLE]]#0 : tensor<f32>, tensor<f64>
  func.return %0#0, %0#1 : tensor<f32>, tensor<f64>
}

  // CHECK-SAME: [[ID_OUTPUT]]
  // CHECK-SAME: key = "host_compute_channel_recv"

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

    %cst_63 = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

      shape_attrs.push_back(TF::ShapeAttr::get(rewriter.getContext(),
                                               mlir::cast<ShapedType>(ty)));
    }

    // Clone the `host_func` in the `host_mlir_module` attribute if it exists
    // and use it for `shape_inference_graph` attribute on XlaHostCompute.
    func::FuncOp cloned_func;
    SymbolTable manager(op->getParentOfType<ModuleOp>());

      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),
      /*manual_sharding=*/builder.getBoolAttr(true));
  host_compute_out_ops.push_back(host_compute);

  // Convert MANUAL sharded outputs to split sharded outputs.

  // Extract the module and function.
  // The '_XlaHostComputeMlir` verifier verifies that `host_mlir_module`
  // attribute is well formed, so we just return in case of an error in
  // extracting the host function since it should never occur.
  StringAttr host_module =
      host_compute_op->getAttrOfType<StringAttr>("host_mlir_module");
  if (host_module.getValue().empty()) return false;

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

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

// -----

    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
    Res<Variadic<TF_Tensor>, [{A list of tensors that will be returned to the device.}]>:$outputs
  );