} else {
    return first_try_buffer.status();
  }
}
}  // namespace

void PjRtDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                              absl::string_view tensor_name,
                                              Device* device,
                                              Tensor* cpu_tensor,
                                              StatusCallback done) {

  TFE_OpAddInput(op, in, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TF_DeleteStatus(status);

  TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(in));
  TFE_OpSetAttrString(op, "tensor_name", "dummy", 5);
  TFE_OpSetAttrString(op, "send_device", send_device.c_str(),
                      send_device.size());
  TFE_OpSetAttrString(op, "recv_device", recv_device.c_str(),

    mlir::BlockArgument arg = bb.getArgument(i);
    std::string tensor_name;
    if (has_input_attr)
      tensor_name = std::string(name_mapper_.GetUniqueName(arg));
    if (tensor_name.empty()) tensor_name = absl::StrCat("arg", i);
    if (!build_tensor_and_buffer(arg, index, tensor_name)) return std::nullopt;
  }

  bool failed_once = false;

    });
    if (!status.ok()) {
      done(status);
    }
  }
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             absl::string_view tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  if (device_tensor->NumElements() == 0) {

  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[recv]])
  // CHECK: tf_mlrt.await [[result_future]]
  %0 = "tf.Recv"() {__op_key = 0: i32, device = "/device:CPU:0", tensor_name = "tensor", send_device = "/device:CPU:0", send_device_incarnation = 0, recv_device = "/device:CPU:0"} : () -> tensor<f32>

  inputs.reserve(assets.size());
  for (const auto& asset : assets) {
    TensorInfo tensor_info;
    tensor_info.set_name(asset.tensor_name);
    tensor_info.set_dtype(DT_STRING);
    tensor_info.mutable_tensor_shape();
    inputs.push_back({asset.tensor_name, tensor_info});
  }

  std::unordered_map<std::string, std::string> tf_name_to_mlir_name;
  TF_ASSIGN_OR_RETURN(auto sub_module,