Skip to main content

tenferro_xla/
error.rs

1use std::error::Error as StdError;
2use std::path::PathBuf;
3
4use tenferro_tensor::{DType, ErrorKind, ValidationKind};
5
6/// Erased typed source used only at the XLA plugin boundary.
7pub type BoxError = Box<dyn StdError + Send + Sync + 'static>;
8
9/// Error type for StableHLO lowering and runtime PJRT plugin loading.
10///
11/// # Examples
12///
13/// ```
14/// use tenferro_runtime::{DType, GraphCompiler, TracedTensor};
15/// use tenferro_xla::{lower_compiled_to_stablehlo, Error};
16///
17/// let x = TracedTensor::input_symbolic_shape(DType::I64, 1).unwrap();
18/// let mut compiler = GraphCompiler::new();
19/// let y = x.neg().unwrap();
20/// let program = compiler
21///     .compile_with_input_specs(&y, &[(&x, DType::I64, &[2])])
22///     .unwrap();
23/// let err = lower_compiled_to_stablehlo(&program).unwrap_err();
24/// assert!(matches!(err, Error::UnsupportedDType { .. }));
25/// ```
26#[derive(Debug, thiserror::Error)]
27pub enum Error {
28    #[error("XLA lowering does not support dtype {dtype:?} in {context}")]
29    UnsupportedDType { dtype: DType, context: &'static str },
30    #[error("XLA lowering does not support ExecOp::{op}: {reason}")]
31    UnsupportedOp {
32        op: &'static str,
33        reason: &'static str,
34    },
35    #[error(
36        "XLA lowering supports only exact static shapes; ExecOp::{op} output {output_index} axis {axis} is {kind}"
37    )]
38    NonStaticShape {
39        op: &'static str,
40        output_index: usize,
41        axis: usize,
42        kind: &'static str,
43    },
44    #[error("invalid XLA program: {message}")]
45    InvalidProgram { message: String },
46    #[error("XLA tensor input/output error: {0}")]
47    Tensor(#[from] tenferro_tensor::Error),
48    #[error("XLA extension standard-op lowering failed: {source}")]
49    ExtensionLowering {
50        #[source]
51        source: tenferro_ops::ext_op::ExtensionLoweringError,
52    },
53    #[error("PJRT support requires enabling the tenferro-xla `pjrt` feature")]
54    PjrtFeatureDisabled,
55    #[error("PJRT execution requires an executor created from a loaded plugin")]
56    PjrtPluginNotLoaded,
57    #[error("PJRT call {call} failed: {message}")]
58    PjrtCall { call: &'static str, message: String },
59    #[error("environment variable {var} is not set; set it to a PJRT plugin .so path")]
60    MissingEnv { var: &'static str },
61    #[error("failed to load PJRT plugin from {path}: {source}")]
62    PluginLoad {
63        path: PathBuf,
64        #[source]
65        source: BoxError,
66    },
67}
68
69impl Error {
70    /// Return the stable coarse classification for this XLA failure.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tenferro_tensor::ErrorKind;
76    /// use tenferro_xla::Error;
77    ///
78    /// assert_eq!(
79    ///     Error::UnsupportedOp { op: "Maximum", reason: "not lowered" }.kind(),
80    ///     ErrorKind::Unsupported,
81    /// );
82    /// ```
83    #[must_use]
84    pub fn kind(&self) -> ErrorKind {
85        match self {
86            Self::UnsupportedDType { .. } | Self::UnsupportedOp { .. } => ErrorKind::Unsupported,
87            Self::NonStaticShape { .. } => ErrorKind::Validation(ValidationKind::ShapeMismatch),
88            Self::InvalidProgram { .. } => ErrorKind::Validation(ValidationKind::InvalidArgument),
89            Self::Tensor(source) => source.kind(),
90            Self::ExtensionLowering { source } => source.kind(),
91            Self::PjrtFeatureDisabled | Self::PjrtPluginNotLoaded | Self::MissingEnv { .. } => {
92                ErrorKind::RuntimeState
93            }
94            Self::PluginLoad { .. } => ErrorKind::Io,
95            Self::PjrtCall { .. } => ErrorKind::BackendFailure,
96        }
97    }
98}
99
100/// Result alias for `tenferro-xla`.
101///
102/// # Examples
103///
104/// ```
105/// use tenferro_xla::Result;
106///
107/// fn ok() -> Result<()> {
108///     Ok(())
109/// }
110///
111/// ok().unwrap();
112/// ```
113pub type Result<T> = std::result::Result<T, Error>;
114
115#[cfg(test)]
116mod tests;