1use std::fmt;
2use std::ops::Range;
3
4use tenferro_tensor::{BackendSession, Tensor};
5
6use crate::backend::CompactQrResult;
7use crate::QrOptions;
8
9#[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 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 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 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 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 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 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 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 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 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 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 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 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}