tenferro_ops/
shape_constraint.rs

1use tenferro_tensor::{DType, ErrorKind, ValidationKind};
2
3use crate::SymDim;
4
5/// A relation between two symbolic shape expressions.
6///
7/// # Examples
8///
9/// ```rust
10/// use tenferro_ops::ShapeRelation;
11///
12/// assert_eq!(ShapeRelation::Equal, ShapeRelation::Equal);
13/// ```
14#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
15pub enum ShapeRelation {
16    /// The two symbolic expressions must evaluate to the same dimension.
17    Equal,
18}
19
20/// A shape relation recorded by an extension during metadata inference.
21///
22/// The constraint stores expressions as provided. It does not attempt to
23/// normalize or solve them.
24#[doc(hidden)]
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct ExtensionShapeConstraint {
27    relation: ShapeRelation,
28    lhs: SymDim,
29    rhs: SymDim,
30}
31
32impl ExtensionShapeConstraint {
33    fn equal(lhs: SymDim, rhs: SymDim) -> Self {
34        Self {
35            relation: ShapeRelation::Equal,
36            lhs,
37            rhs,
38        }
39    }
40
41    /// Return the relation imposed by this constraint.
42    #[doc(hidden)]
43    pub fn relation(&self) -> ShapeRelation {
44        self.relation
45    }
46
47    /// Return the left-hand symbolic expression.
48    #[doc(hidden)]
49    pub fn lhs(&self) -> &SymDim {
50        &self.lhs
51    }
52
53    /// Return the right-hand symbolic expression.
54    #[doc(hidden)]
55    pub fn rhs(&self) -> &SymDim {
56        &self.rhs
57    }
58}
59
60/// Errors produced while extension shape requirements inspect input metadata.
61///
62/// # Examples
63///
64/// ```rust
65/// use tenferro_ops::ExtensionShapeError;
66///
67/// let error = ExtensionShapeError::InputOutOfBounds {
68///     family_id: "example.v1",
69///     input: 2,
70///     input_count: 1,
71/// };
72/// assert!(error.to_string().contains("example.v1"));
73/// ```
74#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
75pub enum ExtensionShapeError {
76    /// An input metadata index was outside the supplied input list.
77    #[error(
78        "extension family {family_id:?} input index {input} out of bounds for {input_count} inputs"
79    )]
80    InputOutOfBounds {
81        /// Stable extension family identifier.
82        family_id: &'static str,
83        /// Requested input index.
84        input: usize,
85        /// Number of available inputs.
86        input_count: usize,
87    },
88    /// An axis was outside the selected input's rank.
89    #[error(
90        "extension family {family_id:?} axis {axis} out of bounds for input {input} rank {rank}"
91    )]
92    AxisOutOfBounds {
93        /// Stable extension family identifier.
94        family_id: &'static str,
95        /// Selected input index.
96        input: usize,
97        /// Requested axis index.
98        axis: usize,
99        /// Rank of the selected input.
100        rank: usize,
101    },
102    /// Two inputs required to have the same shape had different ranks.
103    #[error(
104        "extension family {family_id:?} requires inputs {lhs_input} and {rhs_input} to have the same shape, but their ranks are {lhs_rank} and {rhs_rank}"
105    )]
106    RankMismatch {
107        /// Stable extension family identifier.
108        family_id: &'static str,
109        /// Left-hand input index.
110        lhs_input: usize,
111        /// Left-hand input rank.
112        lhs_rank: usize,
113        /// Right-hand input index.
114        rhs_input: usize,
115        /// Right-hand input rank.
116        rhs_rank: usize,
117    },
118}
119
120impl From<ExtensionShapeError> for tenferro_tensor::Error {
121    fn from(error: ExtensionShapeError) -> Self {
122        let (family_id, kind) = match &error {
123            ExtensionShapeError::InputOutOfBounds { family_id, .. } => (
124                *family_id,
125                ErrorKind::Validation(ValidationKind::InvalidArgument),
126            ),
127            ExtensionShapeError::AxisOutOfBounds { family_id, .. } => (
128                *family_id,
129                ErrorKind::Validation(ValidationKind::AxisOutOfBounds),
130            ),
131            ExtensionShapeError::RankMismatch { family_id, .. } => (
132                *family_id,
133                ErrorKind::Validation(ValidationKind::RankMismatch),
134            ),
135        };
136        tenferro_tensor::Error::extension("extension", family_id, kind, error)
137    }
138}
139
140/// Input metadata and equality requirements for one extension inference call.
141///
142/// The context records requirements declaratively. It does not prove, reject,
143/// normalize, or solve them.
144///
145/// # Examples
146///
147/// ```rust
148/// use tenferro_ops::ExtensionShapeContext;
149///
150/// fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
151///     let lhs = ctx.input_axis(0, 0)?;
152///     let rhs = ctx.input_axis(1, 0)?;
153///     ctx.require_equal(lhs, 2 * rhs)?;
154///     Ok(())
155/// }
156/// ```
157#[derive(Debug)]
158pub struct ExtensionShapeContext<'a> {
159    family_id: &'static str,
160    input_dtypes: &'a [DType],
161    input_shapes: &'a [&'a [SymDim]],
162    constraints: Vec<ExtensionShapeConstraint>,
163}
164
165impl<'a> ExtensionShapeContext<'a> {
166    /// Construct a context for the internal extension inference driver.
167    #[doc(hidden)]
168    pub fn new_for_inference(
169        family_id: &'static str,
170        input_dtypes: &'a [DType],
171        input_shapes: &'a [&'a [SymDim]],
172    ) -> Self {
173        Self {
174            family_id,
175            input_dtypes,
176            input_shapes,
177            constraints: Vec::new(),
178        }
179    }
180
181    /// Return the dtype of one extension input.
182    ///
183    /// # Examples
184    ///
185    /// ```rust
186    /// use tenferro_ops::ExtensionShapeContext;
187    /// use tenferro_tensor::DType;
188    ///
189    /// fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<DType> {
190    ///     Ok(ctx.input_dtype(0)?)
191    /// }
192    /// ```
193    ///
194    /// # Errors
195    ///
196    /// Returns [`ExtensionShapeError::InputOutOfBounds`] when `input` is not
197    /// present in the extension metadata.
198    pub fn input_dtype(&self, input: usize) -> Result<DType, ExtensionShapeError> {
199        self.input_dtypes
200            .get(input)
201            .copied()
202            .ok_or(ExtensionShapeError::InputOutOfBounds {
203                family_id: self.family_id,
204                input,
205                input_count: self.input_dtypes.len(),
206            })
207    }
208
209    /// Return the symbolic shape of one extension input.
210    ///
211    /// # Examples
212    ///
213    /// ```rust
214    /// use tenferro_ops::ExtensionShapeContext;
215    ///
216    /// fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
217    ///     let _shape = ctx.input_shape(0)?;
218    ///     Ok(())
219    /// }
220    /// ```
221    ///
222    /// # Errors
223    ///
224    /// Returns [`ExtensionShapeError::InputOutOfBounds`] when `input` is not
225    /// present in the extension metadata.
226    pub fn input_shape(&self, input: usize) -> Result<&[SymDim], ExtensionShapeError> {
227        self.input_shapes
228            .get(input)
229            .copied()
230            .ok_or(ExtensionShapeError::InputOutOfBounds {
231                family_id: self.family_id,
232                input,
233                input_count: self.input_shapes.len(),
234            })
235    }
236
237    /// Return one symbolic axis expression from an extension input.
238    ///
239    /// # Examples
240    ///
241    /// ```rust
242    /// use tenferro_ops::ExtensionShapeContext;
243    /// use tenferro_ops::SymDim;
244    ///
245    /// fn infer(ctx: &ExtensionShapeContext<'_>) -> tenferro_tensor::Result<SymDim> {
246    ///     Ok(ctx.input_axis(0, 0)?)
247    /// }
248    /// ```
249    ///
250    /// # Errors
251    ///
252    /// Returns [`ExtensionShapeError::InputOutOfBounds`] for an unknown input
253    /// or [`ExtensionShapeError::AxisOutOfBounds`] for an axis outside that
254    /// input's rank.
255    pub fn input_axis(&self, input: usize, axis: usize) -> Result<SymDim, ExtensionShapeError> {
256        let shape = self.input_shape(input)?;
257        shape
258            .get(axis)
259            .cloned()
260            .ok_or(ExtensionShapeError::AxisOutOfBounds {
261                family_id: self.family_id,
262                input,
263                axis,
264                rank: shape.len(),
265            })
266    }
267
268    /// Record equality of two symbolic dimension expressions.
269    ///
270    /// This method records the expressions without trying to solve them.
271    ///
272    /// # Examples
273    ///
274    /// ```rust
275    /// use tenferro_ops::ExtensionShapeContext;
276    ///
277    /// fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
278    ///     let lhs = ctx.input_axis(0, 0)?;
279    ///     let rhs = ctx.input_axis(1, 0)?;
280    ///     ctx.require_equal(lhs, 2 * rhs)?;
281    ///     Ok(())
282    /// }
283    /// ```
284    ///
285    /// # Errors
286    ///
287    /// Returns `Ok(())`; this method only records a symbolic equality and does
288    /// not evaluate it or produce an [`ExtensionShapeError`].
289    pub fn require_equal(&mut self, lhs: SymDim, rhs: SymDim) -> Result<(), ExtensionShapeError> {
290        self.constraints
291            .push(ExtensionShapeConstraint::equal(lhs, rhs));
292        Ok(())
293    }
294
295    /// Record equality of two input axes.
296    ///
297    /// # Examples
298    ///
299    /// ```rust
300    /// use tenferro_ops::ExtensionShapeContext;
301    ///
302    /// fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
303    ///     ctx.require_axes_equal((0, 0), (1, 0))?;
304    ///     Ok(())
305    /// }
306    /// ```
307    ///
308    /// # Errors
309    ///
310    /// Returns [`ExtensionShapeError::InputOutOfBounds`] or
311    /// [`ExtensionShapeError::AxisOutOfBounds`] when either referenced axis is
312    /// absent from the extension metadata.
313    pub fn require_axes_equal(
314        &mut self,
315        lhs: (usize, usize),
316        rhs: (usize, usize),
317    ) -> Result<(), ExtensionShapeError> {
318        let lhs = self.input_axis(lhs.0, lhs.1)?;
319        let rhs = self.input_axis(rhs.0, rhs.1)?;
320        self.require_equal(lhs, rhs)
321    }
322
323    /// Require two extension inputs to have the same rank and axis extents.
324    ///
325    /// A rank mismatch is returned before any equality is recorded.
326    ///
327    /// # Examples
328    ///
329    /// ```rust
330    /// use tenferro_ops::ExtensionShapeContext;
331    ///
332    /// fn infer(ctx: &mut ExtensionShapeContext<'_>) -> tenferro_tensor::Result<()> {
333    ///     ctx.require_same_shape(0, 1)?;
334    ///     Ok(())
335    /// }
336    /// ```
337    ///
338    /// # Errors
339    ///
340    /// Returns [`ExtensionShapeError::InputOutOfBounds`] when an input is
341    /// absent, or [`ExtensionShapeError::RankMismatch`] when the two inputs do
342    /// not have the same rank.
343    pub fn require_same_shape(
344        &mut self,
345        lhs_input: usize,
346        rhs_input: usize,
347    ) -> Result<(), ExtensionShapeError> {
348        let lhs_shape = self.input_shape(lhs_input)?;
349        let rhs_shape = self.input_shape(rhs_input)?;
350        if lhs_shape.len() != rhs_shape.len() {
351            return Err(ExtensionShapeError::RankMismatch {
352                family_id: self.family_id,
353                lhs_input,
354                lhs_rank: lhs_shape.len(),
355                rhs_input,
356                rhs_rank: rhs_shape.len(),
357            });
358        }
359
360        let equalities: Vec<_> = lhs_shape
361            .iter()
362            .cloned()
363            .zip(rhs_shape.iter().cloned())
364            .map(|(lhs, rhs)| ExtensionShapeConstraint::equal(lhs, rhs))
365            .collect();
366        self.constraints.extend(equalities);
367        Ok(())
368    }
369
370    /// Borrow the constraints collected by this inference call.
371    #[doc(hidden)]
372    pub fn constraints(&self) -> &[ExtensionShapeConstraint] {
373        &self.constraints
374    }
375
376    /// Consume the context and return its collected constraints.
377    #[doc(hidden)]
378    pub fn into_constraints(self) -> Vec<ExtensionShapeConstraint> {
379        self.constraints
380    }
381}
382
383#[cfg(test)]
384mod tests;
tenferro_ops/shape_constraint.rs

tenferro_ops/
shape_constraint.rs
