builder.create<mlir::TF::XlaSendToHostOp>(while_op.getLoc(), condition,
                                                condition_send_recv_key);
      // device_ordinal0 is the ordinal of TPU_REPLICATED_CORE_0 and is only
      // used in the replicated case.
      Value device_ordinal0 = nullptr;
      if (!core_to_device_ordinal.empty())
        device_ordinal0 = core_to_device_ordinal[0];

  return
}
```

will be rewritten as:

```mlir
func @tf_tpu_rewrite(%arg0: tensor<i8>, %arg1: tensor<i8>) {
  %0:2 = tf_device.replicate([%arg0, %arg1] as %arg2: tensor<i8>) {devices = {TPU_REPLICATED_CORE_0 = ["/job:worker/replica:0/task:0/device:TPU:0", "/job:worker/replica:0/task:0/device:TPU:1"], TPU_REPLICATED_HOST_0 = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:CPU:0"]}, n = 2 : i32} {