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