Skip to main content

tenferro_linalg/
householder.rs

1use std::fmt;
2use std::ops::Range;
3
4use tenferro_tensor::{BackendSession, Tensor};
5
6use crate::backend::CompactQrResult;
7use crate::QrOptions;
8
9/// Opaque compact Householder QR state.
10///
11/// The packed reflector tensors remain private so callers cannot accidentally
12/// treat provider state as ordinary tensor values.
13///
14/// # Examples
15///
16/// ```rust
17/// use tenferro_cpu::CpuBackend;
18/// use tenferro_linalg::{QrOptions, TensorLinalgExt};
19/// use tenferro_tensor::{BackendSessionHost, Tensor};
20///
21/// let a = Tensor::from_vec_col_major(
22///     vec![3, 2],
23///     vec![1.0_f64, 0.0, 1.0, 0.0, 1.0, 1.0],
24/// )?;
25/// let mut host = CpuBackend::new();
26/// let r = host.with_backend_session(|session| {
27///     let qr = a.householder_qr(session)?;
28///     qr.r(QrOptions::default(), session)
29/// })?;
30/// assert_eq!(r.shape(), &[2, 2]);
31/// # Ok::<(), tenferro_tensor::Error>(())
32/// ```
33#[derive(Clone)]
34pub struct HouseholderQr<T> {
35    pub(crate) packed: T,
36    pub(crate) coeff: T,
37}
38
39impl<T> fmt::Debug for HouseholderQr<T> {
40    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
41        formatter
42            .debug_struct("HouseholderQr")
43            .finish_non_exhaustive()
44    }
45}
46
47impl HouseholderQr<Tensor> {
48    pub(crate) fn from_backend(state: CompactQrResult) -> Self {
49        Self {
50            packed: state.packed,
51            coeff: state.coeff,
52        }
53    }
54
55    /// Construct compact state for the product of compatible factors `Q * R`.
56    ///
57    /// # Errors
58    ///
59    /// Returns `tenferro_tensor::Error::Validation` for incompatible rank,
60    /// shape, dtype, placement, or a non-trapezoidal R factor;
61    /// `tenferro_tensor::Error::Unsupported` when the provider lacks compact
62    /// QR; or `tenferro_tensor::Error::BackendSource` for provider failures.
63    pub fn from_factors(
64        q: &Tensor,
65        r: &Tensor,
66        session: &mut dyn BackendSession,
67    ) -> tenferro_tensor::Result<Self> {
68        crate::tensor_ext::with_linalg_backend(session, "householder_qr_from_factors", |backend| {
69            backend
70                .householder_qr_from_factors(q, r)
71                .map(Self::from_backend)
72        })
73    }
74
75    /// Append a column block without refactorizing existing columns.
76    ///
77    /// # Errors
78    ///
79    /// Returns `tenferro_tensor::Error::Validation` for incompatible shape,
80    /// dtype, placement, or malformed state; `tenferro_tensor::Error::Unsupported`
81    /// when append is unavailable; or `tenferro_tensor::Error::BackendSource`
82    /// for reflector or factorization failures.
83    pub fn append_columns(
84        &self,
85        block: &Tensor,
86        session: &mut dyn BackendSession,
87    ) -> tenferro_tensor::Result<Self> {
88        crate::tensor_ext::with_linalg_backend(session, "householder_qr_append", |backend| {
89            backend
90                .householder_qr_append(&self.packed, &self.coeff, block)
91                .map(Self::from_backend)
92        })
93    }
94
95    /// Extract the thin upper-trapezoidal factor.
96    ///
97    /// # Errors
98    ///
99    /// Returns `tenferro_tensor::Error::Validation` for malformed state,
100    /// `tenferro_tensor::Error::Unsupported` for an unavailable provider path,
101    /// or `tenferro_tensor::Error::BackendSource` for extraction failures.
102    pub fn r(
103        &self,
104        options: QrOptions,
105        session: &mut dyn BackendSession,
106    ) -> tenferro_tensor::Result<Tensor> {
107        crate::tensor_ext::with_linalg_backend(session, "householder_qr_r", |backend| {
108            backend.householder_qr_r(&self.packed, &self.coeff, options)
109        })
110    }
111
112    /// Materialize a contiguous range of thin-Q columns.
113    ///
114    /// # Errors
115    ///
116    /// Returns `tenferro_tensor::Error::Validation` when the range is outside
117    /// the thin-Q width or state metadata is malformed,
118    /// `tenferro_tensor::Error::Unsupported` for an unavailable provider path,
119    /// or `tenferro_tensor::Error::BackendSource` for execution failures.
120    pub fn q_columns(
121        &self,
122        columns: Range<usize>,
123        options: QrOptions,
124        session: &mut dyn BackendSession,
125    ) -> tenferro_tensor::Result<Tensor> {
126        crate::tensor_ext::with_linalg_backend(session, "householder_qr_q_columns", |backend| {
127            backend.householder_qr_q_columns(&self.packed, &self.coeff, columns, options)
128        })
129    }
130}
131
132#[cfg(feature = "autodiff")]
133impl HouseholderQr<tenferro_ad::EagerTensor> {
134    pub(crate) fn from_eager_outputs(
135        packed: tenferro_ad::EagerTensor,
136        coeff: tenferro_ad::EagerTensor,
137    ) -> Self {
138        Self { packed, coeff }
139    }
140
141    /// Construct compact state from compatible eager factors.
142    ///
143    /// # Errors
144    ///
145    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
146    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
147    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
148    pub fn from_factors(
149        q: &tenferro_ad::EagerTensor,
150        r: &tenferro_ad::EagerTensor,
151    ) -> tenferro_ad::Result<Self> {
152        eager_state(crate::eager_ext::apply_linalg_eager(
153            crate::extension::LinalgOp::HouseholderQrFromFactors,
154            &[q, r],
155        )?)
156    }
157
158    /// Append an eager column block functionally.
159    ///
160    /// # Errors
161    ///
162    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
163    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
164    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
165    pub fn append_columns(&self, block: &tenferro_ad::EagerTensor) -> tenferro_ad::Result<Self> {
166        eager_state(crate::eager_ext::apply_linalg_eager(
167            crate::extension::LinalgOp::HouseholderQrAppend,
168            &[&self.packed, &self.coeff, block],
169        )?)
170    }
171
172    /// Extract eager R.
173    ///
174    /// # Errors
175    ///
176    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
177    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
178    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
179    pub fn r(&self, options: QrOptions) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
180        eager_one(
181            crate::eager_ext::apply_linalg_eager(
182                crate::extension::LinalgOp::HouseholderQrR {
183                    gauge: options.gauge,
184                },
185                &[&self.packed, &self.coeff],
186            )?,
187            "householder_qr_r",
188        )
189    }
190
191    /// Materialize eager thin-Q columns.
192    ///
193    /// # Errors
194    ///
195    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
196    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
197    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
198    pub fn q_columns(
199        &self,
200        columns: Range<usize>,
201        options: QrOptions,
202    ) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
203        eager_one(
204            crate::eager_ext::apply_linalg_eager(
205                crate::extension::LinalgOp::HouseholderQrQColumns {
206                    start: columns.start,
207                    end: columns.end,
208                    gauge: options.gauge,
209                },
210                &[&self.packed, &self.coeff],
211            )?,
212            "householder_qr_q_columns",
213        )
214    }
215}
216
217#[cfg(feature = "autodiff")]
218fn eager_state(
219    outputs: Vec<tenferro_ad::EagerTensor>,
220) -> tenferro_ad::Result<HouseholderQr<tenferro_ad::EagerTensor>> {
221    let mut outputs = outputs.into_iter();
222    match (outputs.next(), outputs.next(), outputs.next()) {
223        (Some(packed), Some(coeff), None) => Ok(HouseholderQr::from_eager_outputs(packed, coeff)),
224        _ => Err(tenferro_ad::Error::Internal(
225            "compact Householder QR returned an unexpected output count".into(),
226        )),
227    }
228}
229
230#[cfg(feature = "autodiff")]
231fn eager_one(
232    outputs: Vec<tenferro_ad::EagerTensor>,
233    op: &'static str,
234) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
235    let mut outputs = outputs.into_iter();
236    match (outputs.next(), outputs.next()) {
237        (Some(output), None) => Ok(output),
238        _ => Err(tenferro_ad::Error::Internal(format!(
239            "{op} returned an unexpected output count"
240        ))),
241    }
242}
243
244impl HouseholderQr<tenferro_runtime::TracedTensor> {
245    pub(crate) fn from_traced_outputs(
246        packed: tenferro_runtime::TracedTensor,
247        coeff: tenferro_runtime::TracedTensor,
248    ) -> Self {
249        Self { packed, coeff }
250    }
251
252    /// Construct compact state from compatible traced factors.
253    ///
254    /// # Errors
255    ///
256    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
257    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
258    ///
259    /// # Deferred errors
260    ///
261    /// Symbolic shape constraints and backend provider failures may be reported
262    /// during compile or execution.
263    pub fn from_factors(
264        q: &tenferro_runtime::TracedTensor,
265        r: &tenferro_runtime::TracedTensor,
266    ) -> tenferro_runtime::Result<Self> {
267        crate::validation::ensure_float_or_complex("householder_qr_from_factors", q.dtype())?;
268        crate::validation::ensure_float_or_complex("householder_qr_from_factors", r.dtype())?;
269        traced_state(
270            crate::extension::LinalgOp::HouseholderQrFromFactors,
271            &[q, r],
272        )
273    }
274
275    /// Append a traced column block functionally.
276    ///
277    /// # Errors
278    ///
279    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
280    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
281    ///
282    /// # Deferred errors
283    ///
284    /// Symbolic shape constraints and backend provider failures may be reported
285    /// during compile or execution.
286    pub fn append_columns(
287        &self,
288        block: &tenferro_runtime::TracedTensor,
289    ) -> tenferro_runtime::Result<Self> {
290        crate::validation::ensure_float_or_complex("householder_qr_append", block.dtype())?;
291        traced_state(
292            crate::extension::LinalgOp::HouseholderQrAppend,
293            &[&self.packed, &self.coeff, block],
294        )
295    }
296
297    /// Extract traced R.
298    ///
299    /// # Errors
300    ///
301    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
302    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
303    ///
304    /// # Deferred errors
305    ///
306    /// Symbolic shape constraints and backend provider failures may be reported
307    /// during compile or execution.
308    pub fn r(
309        &self,
310        options: QrOptions,
311    ) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
312        traced_one(
313            crate::extension::LinalgOp::HouseholderQrR {
314                gauge: options.gauge,
315            },
316            &[&self.packed, &self.coeff],
317            "householder_qr_r",
318        )
319    }
320
321    /// Materialize traced thin-Q columns.
322    ///
323    /// # Errors
324    ///
325    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
326    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
327    ///
328    /// # Deferred errors
329    ///
330    /// Symbolic shape constraints and backend provider failures may be reported
331    /// during compile or execution.
332    pub fn q_columns(
333        &self,
334        columns: Range<usize>,
335        options: QrOptions,
336    ) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
337        traced_one(
338            crate::extension::LinalgOp::HouseholderQrQColumns {
339                start: columns.start,
340                end: columns.end,
341                gauge: options.gauge,
342            },
343            &[&self.packed, &self.coeff],
344            "householder_qr_q_columns",
345        )
346    }
347}
348
349fn traced_outputs(
350    op: crate::extension::LinalgOp,
351    inputs: &[&tenferro_runtime::TracedTensor],
352) -> tenferro_runtime::Result<Vec<tenferro_runtime::TracedTensor>> {
353    tenferro_runtime::extension::apply(
354        std::sync::Arc::new(crate::extension::LinalgExtensionOp::new(op)),
355        inputs,
356    )
357}
358
359fn traced_state(
360    op: crate::extension::LinalgOp,
361    inputs: &[&tenferro_runtime::TracedTensor],
362) -> tenferro_runtime::Result<HouseholderQr<tenferro_runtime::TracedTensor>> {
363    let mut outputs = traced_outputs(op, inputs)?.into_iter();
364    match (outputs.next(), outputs.next(), outputs.next()) {
365        (Some(packed), Some(coeff), None) => Ok(HouseholderQr::from_traced_outputs(packed, coeff)),
366        _ => Err(tenferro_runtime::Error::Internal(
367            "compact Householder QR returned an unexpected output count".into(),
368        )),
369    }
370}
371
372fn traced_one(
373    op: crate::extension::LinalgOp,
374    inputs: &[&tenferro_runtime::TracedTensor],
375    name: &'static str,
376) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
377    let mut outputs = traced_outputs(op, inputs)?.into_iter();
378    match (outputs.next(), outputs.next()) {
379        (Some(output), None) => Ok(output),
380        _ => Err(tenferro_runtime::Error::Internal(format!(
381            "{name} returned an unexpected output count"
382        ))),
383    }
384}