Skip to content

Error Handling: make XLATensor::Create() return status type.#9544

Merged
ysiraichi merged 1 commit intomasterfrom
ysiraichi/status-for-xlatensor-create
Aug 8, 2025
Merged

Error Handling: make XLATensor::Create() return status type.#9544
ysiraichi merged 1 commit intomasterfrom
ysiraichi/status-for-xlatensor-create

Conversation

@ysiraichi
Copy link
Copy Markdown
Collaborator

This PR updates the signature of XLATensor::Create() to return a status type, enabling better error handling. Specifically, it introduces a new error case: attempting to create an XLATensor with a non-CPU tensor will now return an InvalidArgumentError.

Key Changes:

  • Change the signature of XLATensor::Create() (when tensor data is given), returning a status type for error handling
    • There are other overloads that were left as is, since they don't (directly) fail at this moment
  • Return an InvalidArgumentError when trying to create an XLATensor with a non-CPU tensor
  • Update the calls to XLATensor::Create(), wrapping them with GetValueOrThrow()

Example:

a = torch.rand(5, device="xla")
b = torch.rand(5, device="cuda")
torch.add(a, b)

Before:

Traceback (most recent call last):
  File "test.py", line 9, in test
    torch.add(a, b)
RuntimeError: Check failed: tensor.device().type() == at::kCPU (cuda vs. cpu) (at torch_xla/csrc/tensor.cpp:66)
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()
        torch_xla::runtime::internal::ErrorGenerator::operator&(std::ostream const&) const
        torch_xla::XLATensor::Create(at::Tensor const&, torch::lazy::BackendDevice const&)
        torch_xla::bridge::GetOrCreateXlaTensor(at::Tensor const&, torch::lazy::BackendDevice const&)

        torch_xla::XLANativeFunctions::add(at::Tensor const&, at::Tensor const&, c10::Scalar const&)

        c10::BoxedKernel::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const
        c10::KernelFunction::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const
        c10::Dispatcher::callBoxed(c10::OperatorHandle const&, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const



        c10::BoxedKernel::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const


        at::_ops::add_Tensor::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&)





        at::_ops::add_Tensor::call(at::Tensor const&, at::Tensor const&, c10::Scalar const&)
        at::Tensor::add(at::Tensor const&, c10::Scalar const&) const



        _PyObject_MakeTpCall
        _PyEval_EvalFrameDefault
        PyEval_EvalCode



        _PyRun_SimpleFileObject
        _PyRun_AnyFileObject
        Py_RunMain
        Py_BytesMain
        __libc_start_main
        _start
*** End stack trace ***

After:

Traceback (most recent call last):
  File "test.py", line 9, in test
    torch.add(a, b)
RuntimeError: Could not create an XLATensor out of the provided tensor. Expected tensor data to be on CPU. Got: CUDA. Consider moving the tensor to CPU.

Status Propagation Trace:
    From: Create at torch_xla/csrc/tensor.cpp:67 (error: Could not create an XLATensor out of the provided tensor. Expected tensor data to be on CPU. Got: CUDA. Consider moving the tensor to CPU.)

Exception raised from MaybeThrow at torch_xla/csrc/status.cpp:123 (most recent call first):
C++ CapturedTraceback:
#4 std::enable_if<is_invocable_r_v<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, THPModule_initExtension(_object*, _object*)::{lambda()#1}&>, std::enable_if>::type std::__invoke_r<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, THPModule_initExtension(_object*, _object*)::{lambda()#1}&>(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocato
r<char> >&&, (THPModule_initExtension(_object*, _object*)::{lambda()#1}&)...) from /usr/include/c++/10/bits/invoke.h:115
#5 std::_Function_handler<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > (), THPModule_initExtension(_object*, _object*)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from /usr/include/c++/10/bits/std_function.h:292
#6 std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>::operator()() const from /usr/include/c++/10/bits/std_function.h:622
#7 c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}::operator()() const from /home/ysiraichi/pytorch/build/../c10/util/Logging.cpp:42
#8 std::shared_ptr<c10::PrecomputedLazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > > std::__invoke_impl<std::shared_ptr<c10::PrecomputedLazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > >, c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}&>(std::__inv
oke_other, c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}&) from /usr/include/c++/10/bits/invoke.h:60
#9 std::enable_if<is_invocable_r_v<std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > const>, c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}&>, std::enable_if>::type std::__invoke_r<std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, st
d::allocator<char> > > const>, c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}&>(std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > const>&&, (c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lam
bda()#1}&)...) from /usr/include/c++/10/bits/invoke.h:113
#10 std::_Function_handler<std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > const> (), c10::SetStackTraceFetcher(std::function<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > ()>)::{lambda()#1}>::_M_invoke(std::_Any_data const&) from /usr/include/c++/10/bits/std_function.h:292
#11 std::function<std::shared_ptr<c10::LazyValue<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > > const> ()>::operator()() const from /usr/include/c++/10/bits/std_function.h:622
#12 c10::(anonymous namespace)::PyTorchStyleBacktrace::PyTorchStyleBacktrace(c10::SourceLocation) from /home/ysiraichi/pytorch/build/../c10/util/Logging.cpp:92
#13 void __gnu_cxx::new_allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>::construct<c10::(anonymous namespace)::PyTorchStyleBacktrace, c10::SourceLocation&>(c10::(anonymous namespace)::PyTorchStyleBacktrace*, c10::SourceLocation&) from /usr/include/c++/10/ext/new_allocator.h:150
#14 void std::allocator_traits<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace> >::construct<c10::(anonymous namespace)::PyTorchStyleBacktrace, c10::SourceLocation&>(std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>&, c10::(anonymous namespace)::PyTorchStyleBacktrace*, c10::SourceLocation&) from /usr/include/c++/10/bits/alloc_traits.h:512
#15 std::_Sp_counted_ptr_inplace<c10::(anonymous namespace)::PyTorchStyleBacktrace, std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, (__gnu_cxx::_Lock_policy)2>::_Sp_counted_ptr_inplace<c10::SourceLocation&>(std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr_base.h:551
#16 std::__shared_count<(__gnu_cxx::_Lock_policy)2>::__shared_count<c10::(anonymous namespace)::PyTorchStyleBacktrace, std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, c10::SourceLocation&>(c10::(anonymous namespace)::PyTorchStyleBacktrace*&, std::_Sp_alloc_shared_tag<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace> >, c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr_base.h:
682
#17 std::__shared_ptr<c10::(anonymous namespace)::PyTorchStyleBacktrace, (__gnu_cxx::_Lock_policy)2>::__shared_ptr<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, c10::SourceLocation&>(std::_Sp_alloc_shared_tag<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace> >, c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr_base.h:1371
#18 std::shared_ptr<c10::(anonymous namespace)::PyTorchStyleBacktrace>::shared_ptr<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, c10::SourceLocation&>(std::_Sp_alloc_shared_tag<std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace> >, c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr.h:408
#19 std::shared_ptr<c10::(anonymous namespace)::PyTorchStyleBacktrace> std::allocate_shared<c10::(anonymous namespace)::PyTorchStyleBacktrace, std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace>, c10::SourceLocation&>(std::allocator<c10::(anonymous namespace)::PyTorchStyleBacktrace> const&, c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr.h:860
#20 std::shared_ptr<c10::(anonymous namespace)::PyTorchStyleBacktrace> std::make_shared<c10::(anonymous namespace)::PyTorchStyleBacktrace, c10::SourceLocation&>(c10::SourceLocation&) from /usr/include/c++/10/bits/shared_ptr.h:876
#21 c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) from /home/ysiraichi/pytorch/build/../c10/util/Logging.cpp:114
#22 c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) from /home/ysiraichi/pytorch/build/../c10/util/Exception.cpp:106
#23 torch_xla::MaybeThrow(absl::lts_20230802::Status const&) from ??:0
#24 torch_xla::bridge::GetOrCreateXlaTensor(at::Tensor const&, torch::lazy::BackendDevice const&) from ??:0
#25 torch_xla::(anonymous namespace)::GetBinaryOperands(at::Tensor const&, at::Tensor const&) from aten_xla_type.cpp:0
#26 torch_xla::XLANativeFunctions::add(at::Tensor const&, at::Tensor const&, c10::Scalar const&) from ??:0
#27 c10::impl::make_boxed_from_unboxed_functor<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (at::Tensor const&, at::Tensor const&, c10::Scalar const&), &at::(anonymous namespace)::(anonymous namespace)::wrapper_XLA_Tensor_add>, at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, c10::Scalar const&> >, false>::call(c10::OperatorKernel*, c10::OperatorHandl
e const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) from RegisterXLA.cpp:0
#28 c10::BoxedKernel::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/BoxedKernel_impl.h:48
#29 c10::KernelFunction::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/KernelFunction_impl.h:64
#30 c10::Dispatcher::callBoxed(c10::OperatorHandle const&, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/dispatch/Dispatcher.h:877
#31 c10::OperatorHandle::callBoxed(std::vector<c10::IValue, std::allocator<c10::IValue> >*) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/dispatch/Dispatcher.h:527
#32 (anonymous namespace)::functionalizeFallback(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) from /home/ysiraichi/pytorch/build/../aten/src/ATen/FunctionalizeFallbackKernel.cpp:98
#33 void c10::BoxedKernel::make_boxed_function<&(anonymous namespace)::functionalizeFallback>(c10::OperatorKernel*, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) from /home/ysiraichi/pytorch/build/../aten/src/ATen/FunctionalizeFallbackKernel.cpp:350
#34 c10::BoxedKernel::callBoxed(c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/BoxedKernel_impl.h:48
#35 c10::impl::BoxedKernelWrapper<at::Tensor (at::Tensor const&, at::Tensor const&, c10::Scalar const&), void>::call(c10::BoxedKernel const&, c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/impl/boxing.h:247
#36 at::Tensor c10::KernelFunction::call<at::Tensor, at::Tensor const&, at::Tensor const&, c10::Scalar const&>(c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/KernelFunction_impl.h:157
#37 c10::TypedOperatorHandle<at::Tensor (at::Tensor const&, at::Tensor const&, c10::Scalar const&)>::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/dispatch/Dispatcher.h:613
#38 at::redispatch::add(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) from /home/ysiraichi/pytorch/build/aten/src/ATen/RedispatchFunctions.h:612
#39 torch::autograd::VariableType::(anonymous namespace)::add_Tensor(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&)::{lambda()#1}::operator()() const from /home/ysiraichi/pytorch/build/../torch/csrc/autograd/generated/VariableType_2.cpp:6543
#40 torch::autograd::VariableType::(anonymous namespace)::add_Tensor(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) from /home/ysiraichi/pytorch/build/../torch/csrc/autograd/generated/VariableType_2.cpp:6544
#41 c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&), &torch::autograd::VariableType::(anonymous namespace)::add_Tensor>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&> >, at::Tensor (c10::Dispa
tchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:17
#42 at::Tensor c10::callUnboxedKernelFunction<at::Tensor, at::Tensor const&, at::Tensor const&, c10::Scalar const&>(void*, c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/KernelFunction_impl.h:76
#43 at::Tensor c10::KernelFunction::call<at::Tensor, at::Tensor const&, at::Tensor const&, c10::Scalar const&>(c10::OperatorHandle const&, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) const from /home/ysiraichi/pytorch/build/../aten/src/ATen/core/boxing/KernelFunction_impl.h:149
#44 at::Tensor::add(at::Tensor const&, c10::Scalar const&) const from /home/ysiraichi/pytorch/build/aten/src/ATen/core/TensorBody.h:1669
#45 torch::autograd::THPVariable_add(_object*, _object*, _object*)::{lambda(at::Tensor const&, at::Tensor const&, c10::Scalar const&)#3}::operator()(at::Tensor const&, at::Tensor const&, c10::Scalar const&) const from /home/ysiraichi/pytorch/build/../torch/csrc/autograd/generated/python_torch_functions_2.cpp:1461
#46 torch::autograd::THPVariable_add(_object*, _object*, _object*) from /home/ysiraichi/pytorch/build/../torch/csrc/autograd/generated/python_torch_functions_2.cpp:1463
#47 cfunction_call from /usr/src/python/Objects/methodobject.c:537
#48 _PyObject_MakeTpCall from /usr/src/python/Objects/call.c:240
#49 _PyEval_EvalFrameDefault from /usr/src/python/Python/bytecodes.c:2715
#50 PyEval_EvalCode from /usr/src/python/Python/ceval.c:580
#51 run_eval_code_obj from /usr/src/python/Python/pythonrun.c:1757
#52 run_mod from /usr/src/python/Python/pythonrun.c:1778
#53 pyrun_file from /usr/src/python/Python/pythonrun.c:1674
#54 _PyRun_SimpleFileObject from /usr/src/python/Python/pythonrun.c:459
#55 _PyRun_AnyFileObject from /usr/src/python/Python/pythonrun.c:78
#56 pymain_run_file_obj from /usr/src/python/Modules/main.c:361
#57 Py_BytesMain from /usr/src/python/Modules/main.c:768
#58 __libc_start_main from ./csu/../csu/libc-start.c:308
#59 _start from ??:0

Note: it wasn't possible to create a test for reproducing this error, since we don't have any GPU CI anymore.

@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-xlatensor-create branch from 8c6bf1f to 5b4d600 Compare August 7, 2025 12:22
Copy link
Copy Markdown
Collaborator

@zhanyong-wan zhanyong-wan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

Comment thread torch_xla/csrc/tensor.cpp
@ysiraichi ysiraichi requested a review from zhanyong-wan August 7, 2025 17:27
Copy link
Copy Markdown
Collaborator

@zhanyong-wan zhanyong-wan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@ysiraichi ysiraichi merged commit 095faec into master Aug 8, 2025
23 of 24 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants