Test + Use *WITH_LOCATION macro for propagating the error.

Author: ysiraichi

test/test_operations.py

@@ -2445,6 +2445,12 @@ def test_isneginf_no_fallback(self):
     t = t.to(torch.float16)
     self._test_no_fallback(torch.isneginf, (t,))
 
+  def test_construct_large_tensor_raises_error(self):
+    a = torch.rand(1024, 1024, 1024, 1024, 1024, device=torch_xla.device())
+
+    with self.assertRaisesRegex(RuntimeError, r"Out of memory allocating \d* bytes"):
+      a.cpu()
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -515,17 +515,14 @@ PjRtComputationClient::TransferFromDevice(absl::Span<const DataPtr> handles) {
     ABSL_CHECK(pjrt_data->buffer != nullptr)
         << "PjRt buffer is null in " << __FUNCTION__;
 
-    xla::Literal& literal =
-        literals.emplace_back(host_output_shape(pjrt_data->buffer.get()));
+    xla::Literal& literal = literals.emplace_back(
+        xla::Literal(host_output_shape(pjrt_data->buffer.get()),
+                     /* allocate_arrays= */ false));
     futures.push_back(pjrt_data->buffer->ToLiteral(&literal));
 
     total_size += literal.size_bytes();
   }
-  auto joined = xla::JoinFutures(futures);
-  XLA_RETURN_IF_ERROR(
-      joined.Await(),
-      absl::StrCat(__FUNCTION__,
-                   ": failed to await future from buffer to literal."));
+  XLA_RETURN_IF_ERROR_WITH_LOCATION(xla::JoinFutures(futures).Await());
   InboundDataMetric()->AddSample(total_size);
 
   return literals;