tenferro_linalg/
householder.rs

1use std::fmt;
2use std::ops::Range;
3
4use tenferro_tensor::{Engine, 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::CpuEngine;
18/// use tenferro_linalg::{QrOptions, TensorLinalgExt};
19/// use tenferro_tensor::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 = tenferro_cpu::CpuEngine::serial();
26/// let qr = a.householder_qr(&mut host)?;
27/// let r = qr.r(QrOptions::default(), &mut host)?;
28/// assert_eq!(r.shape(), &[2, 2]);
29/// # Ok::<(), tenferro_tensor::Error>(())
30/// ```
31#[derive(Clone)]
32pub struct HouseholderQr<T> {
33    pub(crate) packed: T,
34    pub(crate) coeff: T,
35}
36
37impl<T> fmt::Debug for HouseholderQr<T> {
38    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
39        formatter
40            .debug_struct("HouseholderQr")
41            .finish_non_exhaustive()
42    }
43}
44
45impl HouseholderQr<Tensor> {
46    pub(crate) fn from_backend(state: CompactQrResult) -> Self {
47        Self {
48            packed: state.packed,
49            coeff: state.coeff,
50        }
51    }
52
53    /// Construct compact state for the product of compatible factors `Q * R`.
54    ///
55    /// # Errors
56    ///
57    /// Returns `tenferro_tensor::Error::Validation` for incompatible rank,
58    /// shape, dtype, placement, or a non-trapezoidal R factor;
59    /// `tenferro_tensor::Error::Unsupported` when the provider lacks compact
60    /// QR; or `tenferro_tensor::Error::BackendSource` for provider failures.
61    pub fn from_factors(
62        q: &Tensor,
63        r: &Tensor,
64        session: &mut dyn Engine,
65    ) -> tenferro_tensor::Result<Self> {
66        crate::tensor_ext::with_linalg_backend(session, "householder_qr_from_factors", |backend| {
67            backend
68                .householder_qr_from_factors(q, r)
69                .map(Self::from_backend)
70        })
71    }
72
73    /// Append a column block without refactorizing existing columns.
74    ///
75    /// # Errors
76    ///
77    /// Returns `tenferro_tensor::Error::Validation` for incompatible shape,
78    /// dtype, placement, or malformed state; `tenferro_tensor::Error::Unsupported`
79    /// when append is unavailable; or `tenferro_tensor::Error::BackendSource`
80    /// for reflector or factorization failures.
81    pub fn append_columns(
82        &self,
83        block: &Tensor,
84        session: &mut dyn Engine,
85    ) -> tenferro_tensor::Result<Self> {
86        crate::tensor_ext::with_linalg_backend(session, "householder_qr_append", |backend| {
87            backend
88                .householder_qr_append(&self.packed, &self.coeff, block)
89                .map(Self::from_backend)
90        })
91    }
92
93    /// Extract the thin upper-trapezoidal factor.
94    ///
95    /// # Errors
96    ///
97    /// Returns `tenferro_tensor::Error::Validation` for malformed state,
98    /// `tenferro_tensor::Error::Unsupported` for an unavailable provider path,
99    /// or `tenferro_tensor::Error::BackendSource` for extraction failures.
100    pub fn r(
101        &self,
102        options: QrOptions,
103        session: &mut dyn Engine,
104    ) -> tenferro_tensor::Result<Tensor> {
105        crate::tensor_ext::with_linalg_backend(session, "householder_qr_r", |backend| {
106            backend.householder_qr_r(&self.packed, &self.coeff, options)
107        })
108    }
109
110    /// Materialize a contiguous range of Q columns, up to the full-Q width.
111    ///
112    /// For an `m x n` input with `k = min(m, n)`, columns `0..k` are the thin-Q
113    /// factor and columns `k..m` span the orthogonal complement of the input's
114    /// column space — the cheaper route to a nullspace basis than a full SVD.
115    /// `QrOptions::gauge` is defined by R's diagonal, so `PositiveDiagonal`
116    /// fixes the first `k` columns only; a complement column has no diagonal to
117    /// fix and is returned as the reflector product produced it.
118    ///
119    /// # Errors
120    ///
121    /// Returns `tenferro_tensor::Error::Validation` when the range is outside
122    /// the full-Q width `0..m` or state metadata is malformed,
123    /// `tenferro_tensor::Error::Unsupported` for an unavailable provider path,
124    /// or `tenferro_tensor::Error::BackendSource` for execution failures.
125    pub fn q_columns(
126        &self,
127        columns: Range<usize>,
128        options: QrOptions,
129        session: &mut dyn Engine,
130    ) -> tenferro_tensor::Result<Tensor> {
131        crate::tensor_ext::with_linalg_backend(session, "householder_qr_q_columns", |backend| {
132            backend.householder_qr_q_columns(&self.packed, &self.coeff, columns, options)
133        })
134    }
135}
136
137#[cfg(feature = "autodiff")]
138impl HouseholderQr<tenferro_ad::EagerTensor> {
139    pub(crate) fn from_eager_outputs(
140        packed: tenferro_ad::EagerTensor,
141        coeff: tenferro_ad::EagerTensor,
142    ) -> Self {
143        Self { packed, coeff }
144    }
145
146    /// Construct compact state from compatible eager factors.
147    ///
148    /// # Examples
149    /// ```rust
150    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
151    /// use tenferro_linalg::HouseholderQr;
152    /// let ctx = EagerRuntime::new()?;
153    /// let state = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
154    ///     let q = s.constant_from(Tensor::from_vec_col_major([1, 1], vec![1.0_f64])?)?;
155    ///     let r = s.constant_from(Tensor::from_vec_col_major([1, 1], vec![2.0_f64])?)?;
156    ///     HouseholderQr::<EagerTensor>::from_factors(&q, &r, s)
157    /// })?;
158    /// assert!(format!("{state:?}").starts_with("HouseholderQr"));
159    /// # Ok::<(), tenferro_ad::Error>(())
160    /// ```
161    /// # Errors
162    ///
163    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
164    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
165    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
166    pub fn from_factors(
167        q: &tenferro_ad::EagerTensor,
168        r: &tenferro_ad::EagerTensor,
169        session: &mut tenferro_ad::EagerSession<'_>,
170    ) -> tenferro_ad::Result<Self> {
171        eager_state(crate::eager_ext::apply_linalg_eager_in_session(
172            session,
173            crate::extension::LinalgOp::HouseholderQrFromFactors,
174            &[q, r],
175        )?)
176    }
177
178    /// Append an eager column block functionally.
179    ///
180    /// # Examples
181    /// ```rust
182    /// use tenferro_ad::{EagerRuntime, Tensor};
183    /// use tenferro_linalg::EagerSessionLinalgExt;
184    /// let ctx = EagerRuntime::new()?;
185    /// let state = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
186    ///     let a = s.constant_from(Tensor::from_vec_col_major([2, 1], vec![1.0_f64, 0.0])?)?;
187    ///     let b = s.constant_from(Tensor::from_vec_col_major([2, 1], vec![0.0_f64, 1.0])?)?;
188    ///     s.householder_qr(&a)?.append_columns(&b, s)
189    /// })?;
190    /// assert!(format!("{state:?}").starts_with("HouseholderQr"));
191    /// # Ok::<(), tenferro_ad::Error>(())
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 append_columns(
199        &self,
200        block: &tenferro_ad::EagerTensor,
201        session: &mut tenferro_ad::EagerSession<'_>,
202    ) -> tenferro_ad::Result<Self> {
203        eager_state(crate::eager_ext::apply_linalg_eager_in_session(
204            session,
205            crate::extension::LinalgOp::HouseholderQrAppend,
206            &[&self.packed, &self.coeff, block],
207        )?)
208    }
209
210    /// Extract eager R.
211    ///
212    /// # Examples
213    /// ```rust
214    /// use tenferro_ad::{EagerRuntime, Tensor};
215    /// use tenferro_linalg::{EagerSessionLinalgExt, QrOptions};
216    /// let ctx = EagerRuntime::new()?;
217    /// let r = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
218    ///     let a = s.constant_from(Tensor::from_vec_col_major([2, 1], vec![1.0_f64, 2.0])?)?;
219    ///     s.householder_qr(&a)?.r(QrOptions::default(), s)
220    /// })?;
221    /// assert_eq!(r.shape(), &[1, 1]);
222    /// # Ok::<(), tenferro_ad::Error>(())
223    /// ```
224    /// # Errors
225    ///
226    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
227    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
228    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
229    pub fn r(
230        &self,
231        options: QrOptions,
232        session: &mut tenferro_ad::EagerSession<'_>,
233    ) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
234        eager_one(
235            crate::eager_ext::apply_linalg_eager_in_session(
236                session,
237                crate::extension::LinalgOp::HouseholderQrR {
238                    gauge: options.gauge,
239                },
240                &[&self.packed, &self.coeff],
241            )?,
242            "householder_qr_r",
243        )
244    }
245
246    /// Materialize eager Q columns, up to the full-Q width.
247    ///
248    /// # Examples
249    /// ```rust
250    /// use tenferro_ad::{EagerRuntime, Tensor};
251    /// use tenferro_linalg::{EagerSessionLinalgExt, QrOptions};
252    /// let ctx = EagerRuntime::new()?;
253    /// let q = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
254    ///     let a = s.constant_from(Tensor::from_vec_col_major([2, 1], vec![1.0_f64, 2.0])?)?;
255    ///     s.householder_qr(&a)?.q_columns(0..1, QrOptions::default(), s)
256    /// })?;
257    /// assert_eq!(q.shape(), &[2, 1]);
258    /// # Ok::<(), tenferro_ad::Error>(())
259    /// ```
260    /// Columns `0..k` are the thin-Q factor; columns `k..m` span the orthogonal
261    /// complement of the input's column space. `PositiveDiagonal` fixes only
262    /// the first `k` columns, because the gauge comes from R's diagonal.
263    /// Differentiating through a range that reaches past `k` is unsupported and
264    /// surfaces a typed AD error rather than a silently wrong derivative.
265    ///
266    /// # Errors
267    ///
268    /// Returns `tenferro_ad::Error::Validation` for known invalid metadata,
269    /// `tenferro_ad::Error::Extension` for unsupported or provider failures,
270    /// or `tenferro_ad::Error::RuntimeState` when eager execution is unavailable.
271    pub fn q_columns(
272        &self,
273        columns: Range<usize>,
274        options: QrOptions,
275        session: &mut tenferro_ad::EagerSession<'_>,
276    ) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
277        eager_one(
278            crate::eager_ext::apply_linalg_eager_in_session(
279                session,
280                crate::extension::LinalgOp::HouseholderQrQColumns {
281                    start: columns.start,
282                    end: columns.end,
283                    gauge: options.gauge,
284                },
285                &[&self.packed, &self.coeff],
286            )?,
287            "householder_qr_q_columns",
288        )
289    }
290}
291
292#[cfg(feature = "autodiff")]
293pub(crate) fn eager_state(
294    outputs: Vec<tenferro_ad::EagerTensor>,
295) -> tenferro_ad::Result<HouseholderQr<tenferro_ad::EagerTensor>> {
296    let mut outputs = outputs.into_iter();
297    match (outputs.next(), outputs.next(), outputs.next()) {
298        (Some(packed), Some(coeff), None) => Ok(HouseholderQr::from_eager_outputs(packed, coeff)),
299        _ => Err(tenferro_ad::Error::Internal(
300            "compact Householder QR returned an unexpected output count".into(),
301        )),
302    }
303}
304
305#[cfg(feature = "autodiff")]
306fn eager_one(
307    outputs: Vec<tenferro_ad::EagerTensor>,
308    op: &'static str,
309) -> tenferro_ad::Result<tenferro_ad::EagerTensor> {
310    let mut outputs = outputs.into_iter();
311    match (outputs.next(), outputs.next()) {
312        (Some(output), None) => Ok(output),
313        _ => Err(tenferro_ad::Error::Internal(format!(
314            "{op} returned an unexpected output count"
315        ))),
316    }
317}
318
319impl HouseholderQr<tenferro_runtime::TracedTensor> {
320    pub(crate) fn from_traced_outputs(
321        packed: tenferro_runtime::TracedTensor,
322        coeff: tenferro_runtime::TracedTensor,
323    ) -> Self {
324        Self { packed, coeff }
325    }
326
327    /// Construct compact state from compatible traced factors.
328    ///
329    /// # Errors
330    ///
331    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
332    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
333    ///
334    /// # Deferred errors
335    ///
336    /// Symbolic shape constraints and backend provider failures may be reported
337    /// during compile or execution.
338    pub fn from_factors(
339        q: &tenferro_runtime::TracedTensor,
340        r: &tenferro_runtime::TracedTensor,
341    ) -> tenferro_runtime::Result<Self> {
342        crate::validation::ensure_float_or_complex("householder_qr_from_factors", q.dtype())?;
343        crate::validation::ensure_float_or_complex("householder_qr_from_factors", r.dtype())?;
344        traced_state(
345            crate::extension::LinalgOp::HouseholderQrFromFactors,
346            &[q, r],
347        )
348    }
349
350    /// Append a traced column block functionally.
351    ///
352    /// # Errors
353    ///
354    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
355    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
356    ///
357    /// # Deferred errors
358    ///
359    /// Symbolic shape constraints and backend provider failures may be reported
360    /// during compile or execution.
361    pub fn append_columns(
362        &self,
363        block: &tenferro_runtime::TracedTensor,
364    ) -> tenferro_runtime::Result<Self> {
365        crate::validation::ensure_float_or_complex("householder_qr_append", block.dtype())?;
366        traced_state(
367            crate::extension::LinalgOp::HouseholderQrAppend,
368            &[&self.packed, &self.coeff, block],
369        )
370    }
371
372    /// Extract traced R.
373    ///
374    /// # Errors
375    ///
376    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
377    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
378    ///
379    /// # Deferred errors
380    ///
381    /// Symbolic shape constraints and backend provider failures may be reported
382    /// during compile or execution.
383    pub fn r(
384        &self,
385        options: QrOptions,
386    ) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
387        traced_one(
388            crate::extension::LinalgOp::HouseholderQrR {
389                gauge: options.gauge,
390            },
391            &[&self.packed, &self.coeff],
392            "householder_qr_r",
393        )
394    }
395
396    /// Materialize traced Q columns, up to the full-Q width.
397    ///
398    /// Columns `0..k` are the thin-Q factor; columns `k..m` span the orthogonal
399    /// complement of the input's column space. `PositiveDiagonal` fixes only
400    /// the first `k` columns, because the gauge comes from R's diagonal.
401    /// Differentiating through a range that reaches past `k` is unsupported and
402    /// surfaces a typed AD error rather than a silently wrong derivative.
403    ///
404    /// # Errors
405    ///
406    /// Returns `tenferro_runtime::Error::Validation` for known invalid metadata
407    /// or `tenferro_runtime::Error::Extension` for unsupported operation state.
408    ///
409    /// # Deferred errors
410    ///
411    /// Symbolic shape constraints and backend provider failures may be reported
412    /// during compile or execution.
413    pub fn q_columns(
414        &self,
415        columns: Range<usize>,
416        options: QrOptions,
417    ) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
418        traced_one(
419            crate::extension::LinalgOp::HouseholderQrQColumns {
420                start: columns.start,
421                end: columns.end,
422                gauge: options.gauge,
423            },
424            &[&self.packed, &self.coeff],
425            "householder_qr_q_columns",
426        )
427    }
428}
429
430fn traced_outputs(
431    op: crate::extension::LinalgOp,
432    inputs: &[&tenferro_runtime::TracedTensor],
433) -> tenferro_runtime::Result<Vec<tenferro_runtime::TracedTensor>> {
434    tenferro_runtime::extension::apply(
435        std::sync::Arc::new(crate::extension::LinalgExtensionOp::new(op)),
436        inputs,
437    )
438}
439
440fn traced_state(
441    op: crate::extension::LinalgOp,
442    inputs: &[&tenferro_runtime::TracedTensor],
443) -> tenferro_runtime::Result<HouseholderQr<tenferro_runtime::TracedTensor>> {
444    let mut outputs = traced_outputs(op, inputs)?.into_iter();
445    match (outputs.next(), outputs.next(), outputs.next()) {
446        (Some(packed), Some(coeff), None) => Ok(HouseholderQr::from_traced_outputs(packed, coeff)),
447        _ => Err(tenferro_runtime::Error::Internal(
448            "compact Householder QR returned an unexpected output count".into(),
449        )),
450    }
451}
452
453fn traced_one(
454    op: crate::extension::LinalgOp,
455    inputs: &[&tenferro_runtime::TracedTensor],
456    name: &'static str,
457) -> tenferro_runtime::Result<tenferro_runtime::TracedTensor> {
458    let mut outputs = traced_outputs(op, inputs)?.into_iter();
459    match (outputs.next(), outputs.next()) {
460        (Some(output), None) => Ok(output),
461        _ => Err(tenferro_runtime::Error::Internal(format!(
462            "{name} returned an unexpected output count"
463        ))),
464    }
465}
tenferro_linalg/householder.rs

tenferro_linalg/
householder.rs
