community/rfcs/RFC: Handling XLA Exceptions in Python at c30f509d7f1ae1a333348b29b9a1ac313dd41fda · tensorflow/community · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
RFC: Handling XLA Exceptions in Python

Status	Propose
Author	Vinay Kumar Burugu
Collaborator	Loki Ravi
PR	https://github.com/tensorflow/community/pull/435

Objective

The aim of this RFC is to:

1. Any exception/status code thrown by XLA to be uniquely named (independent of TensorFlow).
2. XLA exceptions to be readily available to front-end frameworks so that they can be exposed to users.

Motivation

1. Though XLA was developed as a part of Tensorflow, it is also used by PyTorch, JAX, etc. There is a need for exceptions thrown by XLA to be named independent of Tensorflow.
2. Today, XLA defines XlaRuntimeError in the xla client. However, this exception is not readily available to front-ends. Refer to this (https://github.com/google/jax/pull/10676) effort in JAX to expose this Exception to users.

User Benefit

1. *Simplified contract with front-ends*: Add a single line of code to front-ends to import XLA exceptions. For eg: from import xla.exceptions as xla_exceptions.
2. *Better exception names in XLA*: Derive readable exceptions like XlaInvalidArgumentError from XlaError to replace the existing ValueError. Additionally, this would make it easier to catch/trap Exceptions from XLA and attribute errors to XLA.

Design Proposal

This proposal would replace existing std:: and py:: exceptions thrown by the CPP layers with new custom exceptions which will all be extended from XlaError base class. All status objects originating from XLA will be replaced by XLA specific status codes. Equivalent python exceptions will be defined under a common module xla.exceptions which will contain translations for status objects (previously TensorFlow exceptions (https://github.com/tensorflow/tensorflow/blob/87462bfac761435a46641ff2f10ad0b6e5414a4b/tensorflow/python/framework/errors_impl.py)) as well as direct throws from XLA (previously std library exceptions).

This proposal introduces an XlaError base class and all exceptions are extended from this class.

class XlaError : public std::Exception {
public:
    char * what () {
        return "Custom XLA Exception";
    }
}

We will create different exception classes based on the need by extending this base class. For example:

class XlaInvalidValueError : public XlaError {
public:
    char * what () {
        return "Invalid value provided.";
    }
}

We will expose the exceptions under a common module xla.exceptions for use by front-ends. We will additionally modify TensorFlow and PyTorch front-ends to import xla exceptions as import xla.exceptions as xla_exceptions

The effort can be categorized into:

1. Replacing existing status objects

 XlaInternalError (Currently a TensorFlow status object)

Current implementation[Link to code (https://github.com/tensorflow/tensorflow/blob/d1a41c236defc6dd05e6fb31dddea8e0a8b53b96/tensorflow/compiler/xla/runtime/execution_engine.cc#:~:text=return%20InternalError(%22exported%20function%20not%20found%3A%20%25s%22%2C%20function_name)%3B)]:

return InternalError("exported function not found: %s", function_name);

Proposed implementation:

return XlaInternalError("exported function not found: %s", function_name);

Python UX:

import xla.exceptions as xla_exceptions
try:
    #Some line of code
except xla_exceptions.XlaInternalError:
    pass

2. Replacing existing use of std:: and py:: exceptions in the CPP layer

XlaAttributeError (Currently a py:: lib exception)

Current implementation[Link to code (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/xla.cc#:~:text=throw%20py%3A%3Aattribute_error(absl%3A%3AStrCat(%22Unknown%20attribute%20%22%2C%20name))%3B)]:

throw py::attribute_error(absl::StrCat("Unknown attribute ", name));

Proposed implementation:

throw XlaAttributeError(absl::StrCat("Unknown attribute ", name))


Python UX:

import xla.exceptions as xla_exceptions
try:
    #Some line of code
except xla_exceptions.XlaAttributeError:
    pass

3. Exposing them under a common module at the Python layer

We will register these exceptions in Python using pybind.

PYBIND11_MODULE(xla.exceptions, m) {
auto py_xlaerror = py::register_exception<XlaError>(m, "XlaError");
py::register_exception<XlaInvalidArgumentError>(m, "XlaInvalidArgumentError", py_xlaerror.ptr());
py::register_exception<XlaInternalError>(m, "XlaInternalError", py_xlaerror.ptr());
}