tidu/rules/
ad_rule_error.rs

1use std::error::Error;
2use std::fmt;
3
4/// Identifies which AD rule failed or is unavailable.
5///
6/// # Examples
7///
8/// ```
9/// use tidu::ADRuleKind;
10///
11/// assert_eq!(ADRuleKind::Jvp.as_str(), "jvp");
12/// assert_eq!(ADRuleKind::Transpose.as_str(), "transpose");
13/// ```
14#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
15pub enum ADRuleKind {
16    /// JVP rule for forward linearization.
17    Jvp,
18    /// Transpose / VJP rule for a linear primitive.
19    Transpose,
20}
21
22impl ADRuleKind {
23    /// Returns a stable human-readable rule name.
24    ///
25    /// # Examples
26    ///
27    /// ```
28    /// use tidu::ADRuleKind;
29    ///
30    /// assert_eq!(ADRuleKind::Jvp.as_str(), "jvp");
31    /// ```
32    pub const fn as_str(self) -> &'static str {
33        match self {
34            Self::Jvp => "jvp",
35            Self::Transpose => "transpose",
36        }
37    }
38}
39
40/// Error returned when an AD rule cannot be emitted.
41///
42/// # Examples
43///
44/// ```
45/// use tidu::{ADRuleError, ADRuleKind};
46///
47/// let err = ADRuleError::unsupported("my_crate::fft", ADRuleKind::Jvp);
48/// assert_eq!(err.rule(), ADRuleKind::Jvp);
49/// assert!(err.to_string().contains("my_crate::fft"));
50/// ```
51#[derive(Clone, Debug, PartialEq, Eq)]
52pub enum ADRuleError {
53    /// The requested primitive does not provide the requested AD rule.
54    Unsupported {
55        /// Stable primitive name or extension family identifier.
56        op: String,
57        /// Missing rule kind.
58        rule: ADRuleKind,
59    },
60    /// The requested primitive has an AD rule, but the supplied primal inputs
61    /// violate that rule's documented shape, dtype, or activity contract.
62    InvalidInput {
63        /// Stable primitive name or extension family identifier.
64        op: String,
65        /// Rule kind that rejected the input.
66        rule: ADRuleKind,
67        /// Human-readable validation failure.
68        message: String,
69    },
70}
71
72impl ADRuleError {
73    /// Constructs an unsupported-rule error.
74    ///
75    /// # Examples
76    ///
77    /// ```
78    /// use tidu::{ADRuleError, ADRuleKind};
79    ///
80    /// let err = ADRuleError::unsupported("custom::op", ADRuleKind::Transpose);
81    /// assert_eq!(err.rule(), ADRuleKind::Transpose);
82    /// ```
83    pub fn unsupported(op: impl Into<String>, rule: ADRuleKind) -> Self {
84        Self::Unsupported {
85            op: op.into(),
86            rule,
87        }
88    }
89
90    /// Constructs an invalid-input rule error.
91    ///
92    /// # Examples
93    ///
94    /// ```
95    /// use tidu::{ADRuleError, ADRuleKind};
96    ///
97    /// let err = ADRuleError::invalid_input(
98    ///     "custom::solve",
99    ///     ADRuleKind::Transpose,
100    ///     "expected rank >= 2",
101    /// );
102    /// assert_eq!(err.rule(), ADRuleKind::Transpose);
103    /// ```
104    pub fn invalid_input(
105        op: impl Into<String>,
106        rule: ADRuleKind,
107        message: impl Into<String>,
108    ) -> Self {
109        Self::InvalidInput {
110            op: op.into(),
111            rule,
112            message: message.into(),
113        }
114    }
115
116    /// Returns the AD rule kind associated with this error.
117    ///
118    /// # Examples
119    ///
120    /// ```
121    /// use tidu::{ADRuleError, ADRuleKind};
122    ///
123    /// let err = ADRuleError::unsupported("custom::op", ADRuleKind::Jvp);
124    /// assert_eq!(err.rule(), ADRuleKind::Jvp);
125    /// ```
126    #[cfg_attr(coverage, inline(never))]
127    pub const fn rule(&self) -> ADRuleKind {
128        match self {
129            Self::Unsupported { rule, .. } => *rule,
130            Self::InvalidInput { rule, .. } => *rule,
131        }
132    }
133}
134
135impl fmt::Display for ADRuleError {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        match self {
138            Self::Unsupported { op, rule } => {
139                write!(f, "unsupported {} AD rule for {}", rule.as_str(), op)
140            }
141            Self::InvalidInput { op, rule, message } => {
142                write!(f, "invalid {} AD input for {op}: {message}", rule.as_str())
143            }
144        }
145    }
146}
147
148impl Error for ADRuleError {}
149
150/// Result type used by fallible AD rule emission.
151///
152/// # Examples
153///
154/// ```
155/// use tidu::{ADRuleError, ADRuleKind, ADRuleResult};
156///
157/// fn missing_rule() -> ADRuleResult<()> {
158///     Err(ADRuleError::unsupported("custom::op", ADRuleKind::Transpose))
159/// }
160///
161/// assert!(missing_rule().is_err());
162/// ```
163pub type ADRuleResult<T> = Result<T, ADRuleError>;
tidu/rules/ad_rule_error.rs

tidu/rules/
ad_rule_error.rs
