func.return %0#0 : tensor<*xf32>
}

// CHECK-LABEL: @RewriteSendRecvOps
func.func @RewriteSendRecvOps() -> () {
  // CHECK: key = "recv_key_htod_0"
  %0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>

  // CHECK: key = "send_key_dtoh_0"
  "tf.XlaSendToHost"(%0) {key = "send_key"} : (tensor<i32>) -> ()

  func.return

  // CHECK-SAME: handle =
  // CHECK-SAME: type = 3
  // CHECK-SAME: is_host_transfer = true
  // CHECK-SAME: mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "recv_key_htod_0"}
  // CHECK-SAME: mhlo.sharding = "\08\01\1A\01\01\22\01\00"
  // CHECK-SAME: (!mhlo.token) -> (tensor<i32>, !mhlo.token)
  %0 = "tf.XlaRecvFromHost"() {key = "recv_key", shape = #tf_type.shape<>} : () -> tensor<i32>