1use crate::defaults::DynIndex;
2use crate::index_like::IndexLike;
3use crate::index_ops::{common_ind_positions, prepare_contraction, prepare_contraction_pairs};
4use crate::tensor_like::LinearizationOrder;
5use crate::AnyScalar;
6use anyhow::{Context, Result};
7use num_complex::{Complex32, Complex64};
8use num_traits::Zero;
9use rand::Rng;
10use rand_distr::{Distribution, StandardNormal};
11use std::cell::RefCell;
12use std::cmp::Reverse;
13use std::collections::{HashMap, HashSet};
14use std::env;
15use std::sync::{Arc, OnceLock};
16use std::time::{Duration, Instant};
17use tenferro::{
18 DType, DotGeneralConfig, Tensor as NativeTensor, TensorRead, TensorValue, TensorView,
19};
20use tenferro_ad::{extension::adopt_untracked_eager_value, EagerRuntime, EagerTensor};
21use tenferro_einsum::{EagerEinsumExt, EinsumSubscripts};
22use tenferro_linalg::{EagerTensorLinalgExt, RankRevealingQrOptions};
23use tensor4all_tensorbackend::{
24 contract_native_tensor, default_eager_ctx, dense_native_tensor_from_col_major,
25 dense_native_tensor_from_col_major_owned, diag_native_tensor_from_col_major,
26 native_tensor_primal_to_diag, storage_payload_native_read_input, storage_to_native_tensor,
27 ExecutionContext, NativeTensorReadInput, TensorElement,
28};
29use tensor4all_tensorbackend::{
30 src_error_estimate as backend_src_error_estimate,
31 src_error_estimate_general as backend_src_error_estimate_general, Matrix,
32};
33use tensor4all_tensorbackend::{IncrementalQr, IncrementalQrScalar};
34use tensor4all_tensorbackend::{Storage, StorageKind};
35
36use super::contract::{ContractionOptions, PairwiseContractionOptions};
37use super::structured_contraction::{
38 normalize_payload_read_for_roots, storage_from_payload_native, storage_payload_native,
39 OperandLayout, StructuredContractionPlan, StructuredContractionSpec,
40};
41
42fn conjugate_eager(
43 inner: &EagerTensor,
44) -> std::result::Result<EagerTensor, Arc<dyn std::error::Error + Send + Sync + 'static>> {
45 inner.conj().map_err(|source| Arc::new(source) as _)
46}
47
48#[derive(Debug, Default, Clone)]
49struct PairwiseContractProfileEntry {
50 calls: usize,
51 total_time: Duration,
52 total_bytes: usize,
53}
54
55#[derive(Debug, Clone)]
77pub struct TensorHermitianEigendecomposition {
78 pub eigenvalues: Vec<f64>,
80 pub eigenvectors: IdxTensor,
82 pub eigenvector_index: DynIndex,
84}
85
86thread_local! {
87 static PAIRWISE_CONTRACT_PROFILE_STATE: RefCell<HashMap<&'static str, PairwiseContractProfileEntry>> =
88 RefCell::new(HashMap::new());
89}
90
91fn pairwise_contract_profile_enabled() -> bool {
92 static ENABLED: OnceLock<bool> = OnceLock::new();
93 *ENABLED.get_or_init(|| env::var("T4A_PROFILE_PAIRWISE_CONTRACT").is_ok())
94}
95
96fn record_pairwise_contract_profile(section: &'static str, elapsed: Duration) {
97 if !pairwise_contract_profile_enabled() {
98 return;
99 }
100 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
101 let mut state = state.borrow_mut();
102 let entry = state.entry(section).or_default();
103 entry.calls += 1;
104 entry.total_time += elapsed;
105 });
106}
107
108fn record_pairwise_contract_profile_bytes(section: &'static str, bytes: usize) {
109 if !pairwise_contract_profile_enabled() {
110 return;
111 }
112 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
113 let mut state = state.borrow_mut();
114 let entry = state.entry(section).or_default();
115 entry.total_bytes += bytes;
116 });
117}
118
119fn profile_pairwise_contract_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
120 if !pairwise_contract_profile_enabled() {
121 return f();
122 }
123 let started = Instant::now();
124 let result = f();
125 record_pairwise_contract_profile(section, started.elapsed());
126 result
127}
128
129pub fn reset_pairwise_contract_profile() {
131 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| state.borrow_mut().clear());
132}
133
134pub fn print_and_reset_pairwise_contract_profile() {
136 if !pairwise_contract_profile_enabled() {
137 return;
138 }
139 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
140 let mut entries: Vec<_> = state
141 .borrow()
142 .iter()
143 .map(|(section, entry)| (*section, entry.clone()))
144 .collect();
145 state.borrow_mut().clear();
146 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
147
148 eprintln!("=== IdxTensor pairwise contract profile ===");
149 for (section, entry) in entries {
150 let per_call_us = if entry.calls == 0 {
151 0.0
152 } else {
153 entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64
154 };
155 eprintln!(
156 "{section}: calls={} total={:.6}ms per_call={:.3}us bytes={}",
157 entry.calls,
158 entry.total_time.as_secs_f64() * 1.0e3,
159 per_call_us,
160 entry.total_bytes,
161 );
162 }
163 });
164}
165
166fn tensor_profile_bytes(dtype: DType, shape: &[usize]) -> usize {
167 let element_size = match dtype {
168 DType::F32 => 4,
169 DType::F64 => 8,
170 DType::C32 => 8,
171 DType::C64 => 16,
172 DType::I32 => 4,
173 DType::I64 => 8,
174 DType::Bool => 1,
175 };
176 shape
177 .iter()
178 .try_fold(1usize, |bytes, &dim| bytes.checked_mul(dim))
179 .and_then(|elements| elements.checked_mul(element_size))
180 .unwrap_or(usize::MAX)
181}
182
183pub trait RandomScalar: TensorElement {
187 fn random_value<R: Rng>(rng: &mut R) -> Self;
189}
190
191impl RandomScalar for f64 {
192 fn random_value<R: Rng>(rng: &mut R) -> Self {
193 StandardNormal.sample(rng)
194 }
195}
196
197impl RandomScalar for Complex64 {
198 fn random_value<R: Rng>(rng: &mut R) -> Self {
199 Complex64::new(StandardNormal.sample(rng), StandardNormal.sample(rng))
200 }
201}
202
203pub fn compute_permutation_from_indices(
229 original_indices: &[DynIndex],
230 new_indices: &[DynIndex],
231) -> std::result::Result<Vec<usize>, IdxTensorError> {
232 if !(new_indices.len() == original_indices.len()) {
233 return Err(
234 anyhow::anyhow!("new_indices length must match original_indices length").into(),
235 );
236 };
237
238 let mut perm = Vec::with_capacity(new_indices.len());
239 let mut used = std::collections::HashSet::new();
240
241 for new_idx in new_indices {
242 let pos = original_indices
245 .iter()
246 .position(|old_idx| old_idx == new_idx)
247 .ok_or_else(|| {
248 anyhow::anyhow!("new_indices must be a permutation of original_indices")
249 })?;
250
251 if !(used.insert(pos)) {
252 return Err(anyhow::anyhow!("duplicate index in new_indices").into());
253 };
254 perm.push(pos);
255 }
256
257 Ok(perm)
258}
259
260#[derive(Clone)]
265pub(crate) struct StructuredPayload {
266 payload: Arc<EagerTensor>,
267 payload_dims: Vec<usize>,
268 axis_classes: Vec<usize>,
269}
270
271#[derive(Debug, Clone, thiserror::Error)]
290pub enum TensorStorageError {
291 #[error("failed to materialize IdxTensor storage: {source}")]
293 Materialization {
294 #[source]
296 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
297 },
298 #[error(
300 "compact IdxTensor storage does not support dtype {dtype}; the eager payload remains authoritative"
301 )]
302 UnsupportedDtype {
303 dtype: &'static str,
305 },
306 #[error("failed to conjugate IdxTensor storage: {source}")]
309 Conjugation {
310 #[source]
312 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
313 },
314}
315
316#[derive(Debug, Clone, thiserror::Error)]
342pub enum IdxTensorError {
343 #[error("IdxTensor storage operation failed: {source}")]
345 Storage {
346 #[source]
348 source: TensorStorageError,
349 },
350 #[error("IdxTensor materialization failed: {source}")]
352 Materialization {
353 #[source]
355 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
356 },
357 #[error("IdxTensor scalar extraction failed: {source}")]
359 ScalarExtraction {
360 #[source]
362 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
363 },
364 #[error("IdxTensor scalar type mismatch: expected {expected}, got {actual}")]
367 ScalarTypeMismatch {
368 expected: &'static str,
370 actual: String,
372 },
373 #[error("IdxTensor shape mismatch during {operation}: expected {expected}, got {actual}")]
377 ShapeMismatch {
378 operation: &'static str,
380 expected: String,
382 actual: String,
384 },
385 #[error("IdxTensor subtraction failed: {source}")]
387 Subtraction {
388 #[source]
390 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
391 },
392 #[error("IdxTensor {operation} received NaN input")]
395 NaNInput {
396 operation: &'static str,
398 },
399 #[error("IdxTensor tolerance {name} is invalid: {value}")]
401 InvalidTolerance {
402 name: &'static str,
404 value: f64,
406 },
407 #[error("IdxTensor {operation} failed: {source}")]
409 Operation {
410 operation: &'static str,
412 #[source]
414 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
415 },
416}
417
418impl From<anyhow::Error> for IdxTensorError {
419 fn from(source: anyhow::Error) -> Self {
420 Self::operation("IdxTensor", source)
421 }
422}
423
424impl From<TensorStorageError> for IdxTensorError {
425 fn from(source: TensorStorageError) -> Self {
426 Self::Storage { source }
427 }
428}
429
430impl From<tenferro_ad::Error> for IdxTensorError {
431 fn from(source: tenferro_ad::Error) -> Self {
432 Self::Materialization {
433 source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
434 }
435 }
436}
437
438impl From<tensor4all_tensorbackend::EagerContextError> for IdxTensorError {
439 fn from(source: tensor4all_tensorbackend::EagerContextError) -> Self {
440 Self::Materialization {
441 source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
442 }
443 }
444}
445
446impl From<tensor4all_tensorbackend::BridgeError> for IdxTensorError {
447 fn from(source: tensor4all_tensorbackend::BridgeError) -> Self {
448 Self::Materialization {
449 source: Arc::new(source),
450 }
451 }
452}
453
454impl IdxTensorError {
455 fn boxed(error: anyhow::Error) -> Arc<dyn std::error::Error + Send + Sync + 'static> {
456 Arc::from(error.into_boxed_dyn_error())
457 }
458
459 fn materialization(error: anyhow::Error) -> Self {
460 Self::Materialization {
461 source: Self::boxed(error),
462 }
463 }
464
465 fn scalar_extraction(error: anyhow::Error) -> Self {
466 Self::ScalarExtraction {
467 source: Self::boxed(error),
468 }
469 }
470
471 fn operation(operation: &'static str, error: anyhow::Error) -> Self {
472 Self::Operation {
473 operation,
474 source: Self::boxed(error),
475 }
476 }
477}
478
479#[derive(Clone)]
480pub(crate) enum IdxTensorStorage {
481 Materialized(Arc<Storage>),
482 Eager {
483 inner: Arc<EagerTensor>,
484 axis_classes: Vec<usize>,
485 },
486 Compact(Arc<StructuredPayload>),
488 Deferred {
491 source: Box<Self>,
492 error: Arc<TensorStorageError>,
493 },
494}
495
496impl IdxTensorStorage {
497 fn from_storage(storage: Arc<Storage>) -> Self {
498 Self::Materialized(storage)
499 }
500
501 fn from_eager_dense(inner: EagerTensor, rank: usize) -> Self {
502 Self::Eager {
503 inner: Arc::new(inner),
504 axis_classes: IdxTensor::dense_axis_classes(rank),
505 }
506 }
507
508 fn eager(&self) -> Option<&EagerTensor> {
509 match self {
510 Self::Materialized(_) => None,
511 Self::Eager { inner, .. } => Some(inner.as_ref()),
512 Self::Compact(payload) => Some(payload.payload.as_ref()),
513 Self::Deferred { source, .. } => source.eager(),
514 }
515 }
516
517 fn deferred_error(&self) -> Option<&TensorStorageError> {
518 match self {
519 Self::Deferred { error, .. } => Some(error.as_ref()),
520 _ => None,
521 }
522 }
523
524 fn with_deferred_error(self, error: TensorStorageError) -> Self {
525 if self.deferred_error().is_some() {
526 self
527 } else {
528 Self::Deferred {
529 source: Box::new(self),
530 error: Arc::new(error),
531 }
532 }
533 }
534
535 fn axis_classes(&self) -> &[usize] {
536 match self {
537 Self::Materialized(storage) => storage.axis_classes(),
538 Self::Eager { axis_classes, .. } => axis_classes,
539 Self::Compact(payload) => &payload.axis_classes,
540 Self::Deferred { source, .. } => source.axis_classes(),
541 }
542 }
543
544 fn payload_dims(&self) -> &[usize] {
545 match self {
546 Self::Materialized(storage) => storage.payload_dims(),
547 Self::Eager { inner, .. } => inner.shape(),
548 Self::Compact(payload) => &payload.payload_dims,
549 Self::Deferred { source, .. } => source.payload_dims(),
550 }
551 }
552
553 fn payload_strides_vec(&self) -> Vec<isize> {
554 match self {
555 Self::Materialized(storage) => storage.payload_strides().to_vec(),
556 Self::Eager { inner, .. } => {
557 IdxTensor::col_major_strides(inner.shape()).unwrap_or_default()
558 }
559 Self::Compact(payload) => {
560 IdxTensor::col_major_strides(&payload.payload_dims).unwrap_or_default()
561 }
562 Self::Deferred { source, .. } => source.payload_strides_vec(),
563 }
564 }
565
566 fn is_f64(&self) -> bool {
567 match self {
568 Self::Materialized(storage) => storage.is_f64(),
569 Self::Eager { inner, .. } => inner.dtype() == DType::F64,
570 Self::Compact(payload) => payload.payload.dtype() == DType::F64,
571 Self::Deferred { source, .. } => source.is_f64(),
572 }
573 }
574
575 fn is_c64(&self) -> bool {
576 match self {
577 Self::Materialized(storage) => storage.is_c64(),
578 Self::Eager { inner, .. } => inner.dtype() == DType::C64,
579 Self::Compact(payload) => payload.payload.dtype() == DType::C64,
580 Self::Deferred { source, .. } => source.is_c64(),
581 }
582 }
583
584 fn dtype(&self) -> Option<DType> {
585 match self {
586 Self::Materialized(storage) => Some(if storage.is_c64() {
587 DType::C64
588 } else {
589 DType::F64
590 }),
591 Self::Eager { inner, .. } => Some(inner.dtype()),
592 Self::Compact(payload) => Some(payload.payload.dtype()),
593 Self::Deferred { source, .. } => source.dtype(),
594 }
595 }
596
597 fn is_complex(&self) -> bool {
598 match self {
599 Self::Materialized(storage) => storage.is_complex(),
600 Self::Eager { inner, .. } => matches!(inner.dtype(), DType::C32 | DType::C64),
601 Self::Compact(payload) => {
602 matches!(payload.payload.dtype(), DType::C32 | DType::C64)
603 }
604 Self::Deferred { source, .. } => source.is_complex(),
605 }
606 }
607
608 fn is_diag(&self) -> bool {
609 match self {
610 Self::Materialized(storage) => storage.is_diag(),
611 Self::Eager { axis_classes, .. } => IdxTensor::is_diag_axis_classes(axis_classes),
612 Self::Compact(payload) => IdxTensor::is_diag_axis_classes(&payload.axis_classes),
613 Self::Deferred { source, .. } => source.is_diag(),
614 }
615 }
616
617 fn storage_kind(&self) -> StorageKind {
618 match self {
619 Self::Materialized(storage) => storage.storage_kind(),
620 Self::Eager { axis_classes, .. } => {
621 if axis_classes.iter().copied().eq(0..axis_classes.len()) {
622 StorageKind::Dense
623 } else if IdxTensor::is_diag_axis_classes(axis_classes) {
624 StorageKind::Diagonal
625 } else {
626 StorageKind::Structured
627 }
628 }
629 Self::Compact(payload) => {
630 if payload
631 .axis_classes
632 .iter()
633 .copied()
634 .eq(0..payload.axis_classes.len())
635 {
636 StorageKind::Dense
637 } else if IdxTensor::is_diag_axis_classes(&payload.axis_classes) {
638 StorageKind::Diagonal
639 } else {
640 StorageKind::Structured
641 }
642 }
643 Self::Deferred { source, .. } => source.storage_kind(),
644 }
645 }
646
647 fn materialize_eager_payload(inner: &EagerTensor) -> Result<NativeTensor> {
648 let native = inner.duplicate_value()?;
649 if native.is_col_major_contiguous()? {
650 return Ok(native);
651 }
652 let read = inner.tensor_read();
653 Ok(inner.runtime().with_execution_session(|session| {
654 session.to_contiguous_read(TensorRead::from_view(read.tensor_view()))
655 })??)
656 }
657
658 fn materialize(
659 &self,
660 logical_rank: usize,
661 ) -> std::result::Result<Arc<Storage>, TensorStorageError> {
662 match self {
663 Self::Materialized(storage) => Ok(Arc::clone(storage)),
664 Self::Eager {
665 inner,
666 axis_classes,
667 } => {
668 let native = Self::materialize_eager_payload(inner).map_err(|source| {
669 TensorStorageError::Materialization {
670 source: Arc::from(source.into_boxed_dyn_error()),
671 }
672 })?;
673 let dtype = native.dtype();
674 if matches!(dtype, DType::F32 | DType::C32) {
675 return Err(TensorStorageError::UnsupportedDtype {
676 dtype: IdxTensor::dtype_name(dtype),
677 });
678 }
679 IdxTensor::storage_from_native_with_axis_classes(
680 &native,
681 axis_classes,
682 logical_rank,
683 )
684 .map(Arc::new)
685 .map_err(|source| TensorStorageError::Materialization {
686 source: Arc::from(source.into_boxed_dyn_error()),
687 })
688 }
689 Self::Compact(payload) => {
690 let native =
691 Self::materialize_eager_payload(&payload.payload).map_err(|source| {
692 TensorStorageError::Materialization {
693 source: Arc::from(source.into_boxed_dyn_error()),
694 }
695 })?;
696 let dtype = native.dtype();
697 if matches!(dtype, DType::F32 | DType::C32) {
698 return Err(TensorStorageError::UnsupportedDtype {
699 dtype: IdxTensor::dtype_name(dtype),
700 });
701 }
702 IdxTensor::storage_from_native_with_axis_classes(
703 &native,
704 &payload.axis_classes,
705 logical_rank,
706 )
707 .map(Arc::new)
708 .map_err(|source| TensorStorageError::Materialization {
709 source: Arc::from(source.into_boxed_dyn_error()),
710 })
711 }
712 Self::Deferred { error, .. } => Err((**error).clone()),
713 }
714 }
715
716 fn scale_eager_payload(&self, scalar: &AnyScalar) -> Result<Self> {
717 let (payload, payload_dims, axis_classes) = match self {
718 Self::Materialized(storage) => {
719 let native = if storage.is_f64() {
720 let values = storage
721 .payload_f64_col_major_vec()
722 .map_err(anyhow::Error::new)?;
723 dense_native_tensor_from_col_major(&values, storage.payload_dims())?
724 } else {
725 let values = storage
726 .payload_c64_col_major_vec()
727 .map_err(anyhow::Error::new)?;
728 dense_native_tensor_from_col_major(&values, storage.payload_dims())?
729 };
730 (
731 EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
732 storage.payload_dims().to_vec(),
733 storage.axis_classes().to_vec(),
734 )
735 }
736 Self::Eager {
737 inner,
738 axis_classes,
739 } => (
740 (**inner).clone(),
741 inner.shape().to_vec(),
742 axis_classes.clone(),
743 ),
744 Self::Compact(compact) => (
745 (*compact.payload).clone(),
746 compact.payload_dims.clone(),
747 compact.axis_classes.clone(),
748 ),
749 Self::Deferred { error, .. } => return Err(anyhow::Error::new((**error).clone())),
750 };
751 let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
752 let target_dtype = IdxTensor::scale_target_dtype(payload.dtype(), scalar_inner.dtype())?;
753 let payload = if payload.dtype() == target_dtype {
754 payload
755 } else {
756 payload.cast(target_dtype)?
757 };
758 let scalar_inner = if scalar_inner.dtype() == target_dtype {
759 scalar_inner.clone()
760 } else {
761 scalar_inner.cast(target_dtype)?
762 };
763 let scaled = if payload.shape().is_empty() {
764 payload.mul(&scalar_inner)?
765 } else {
766 let subscripts = IdxTensor::scale_subscripts(payload.shape().len())?;
767 [&payload, &scalar_inner].einsum_subscripts(&subscripts)?
768 };
769 match self {
770 Self::Eager { .. } => Ok(Self::Eager {
771 inner: Arc::new(scaled),
772 axis_classes,
773 }),
774 Self::Compact(_) | Self::Materialized(_) => {
775 Ok(Self::Compact(Arc::new(StructuredPayload {
776 payload: Arc::new(scaled),
777 payload_dims,
778 axis_classes,
779 })))
780 }
781 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
782 }
783 }
784
785 fn conjugate_with<F>(&self, conjugate: &F) -> std::result::Result<Self, TensorStorageError>
786 where
787 F: Fn(
788 &EagerTensor,
789 ) -> std::result::Result<
790 EagerTensor,
791 Arc<dyn std::error::Error + Send + Sync + 'static>,
792 >,
793 {
794 match self {
795 Self::Materialized(storage) => Ok(Self::Materialized(Arc::new(storage.conj()))),
796 Self::Eager {
797 inner,
798 axis_classes,
799 } => conjugate(inner)
800 .map(|conjugated| Self::Eager {
801 inner: Arc::new(conjugated),
802 axis_classes: axis_classes.clone(),
803 })
804 .map_err(|source| TensorStorageError::Conjugation { source }),
805 Self::Compact(payload) => conjugate(payload.payload.as_ref())
806 .map(|conjugated| {
807 Self::Compact(Arc::new(StructuredPayload {
808 payload: Arc::new(conjugated),
809 payload_dims: payload.payload_dims.clone(),
810 axis_classes: payload.axis_classes.clone(),
811 }))
812 })
813 .map_err(|source| TensorStorageError::Conjugation { source }),
814 Self::Deferred { error, .. } => Err((**error).clone()),
815 }
816 }
817
818 fn sum_scalar(&self) -> Result<AnyScalar> {
819 match self {
820 Self::Materialized(storage) => {
821 if storage.is_f64() {
822 Ok(AnyScalar::new_real(storage.sum::<f64>()))
823 } else {
824 let value = storage.sum::<Complex64>();
825 Ok(AnyScalar::new_complex(value.re, value.im))
826 }
827 }
828 Self::Eager { inner, .. } => IdxTensor::native_sum_scalar(inner),
829 Self::Compact(payload) => IdxTensor::native_sum_scalar(&payload.payload),
830 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
831 }
832 }
833
834 fn nonfinite_flags(&self) -> Result<(bool, bool)> {
835 match self {
836 Self::Materialized(storage) => Ok(storage.payload_nonfinite_flags()),
837 Self::Eager { inner, .. } => IdxTensor::native_nonfinite_flags(inner),
838 Self::Compact(payload) => IdxTensor::native_nonfinite_flags(&payload.payload),
839 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
840 }
841 }
842
843 fn payload_value_at(&self, payload_coords: &[usize]) -> Result<Complex64> {
844 match self {
845 Self::Materialized(storage) => storage
846 .scalar_at(payload_coords)
847 .map(Complex64::from)
848 .map_err(anyhow::Error::new),
849 Self::Eager { inner, .. } => {
850 IdxTensor::native_complex_payload_value_at(inner, payload_coords)
851 }
852 Self::Compact(payload) => {
853 IdxTensor::native_complex_payload_value_at(&payload.payload, payload_coords)
854 }
855 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
856 }
857 }
858
859 fn for_each_payload_value(&self, mut f: impl FnMut(Complex64)) -> Result<()> {
860 let payload_dims = self.payload_dims();
861 let payload_len = checked_product(payload_dims)?;
862 let mut payload_coords = vec![0usize; payload_dims.len()];
863 for _ in 0..payload_len {
864 f(self.payload_value_at(&payload_coords)?);
865 let mut carry = true;
866 for (coordinate, &dim) in payload_coords.iter_mut().zip(payload_dims.iter()) {
867 if !carry {
868 break;
869 }
870 *coordinate += 1;
871 if *coordinate == dim {
872 *coordinate = 0;
873 } else {
874 carry = false;
875 }
876 }
877 }
878 Ok(())
879 }
880
881 fn compact_payload(&self) -> Option<&StructuredPayload> {
882 match self {
883 Self::Compact(payload) => Some(payload.as_ref()),
884 Self::Deferred { source, .. } => source.compact_payload(),
885 _ => None,
886 }
887 }
888}
889
890#[derive(Debug, thiserror::Error)]
905pub enum StructuredSelectorError {
906 #[error("copy-selector bond dimensions differ: left={left}, right={right}")]
908 BondDimensionMismatch {
909 left: usize,
911 right: usize,
913 },
914 #[error("copy-selector {axis} dimension must be positive")]
916 ZeroDimension {
917 axis: &'static str,
919 },
920 #[error("selected site value {value} is outside 0..{site_dim}")]
922 SelectedValueOutOfBounds {
923 value: usize,
925 site_dim: usize,
927 },
928 #[error("copy-selector payload size overflows usize for dimensions {bond_dim} x {site_dim}")]
930 PayloadSizeOverflow {
931 bond_dim: usize,
933 site_dim: usize,
935 },
936 #[error("copy-selector bond stride {bond_dim} exceeds isize::MAX")]
938 StrideOverflow {
939 bond_dim: usize,
941 },
942 #[error("could not allocate copy-selector payload with {elements} elements")]
944 AllocationFailed {
945 elements: usize,
947 },
948 #[error("invalid copy-selector storage: {message}")]
950 InvalidStorage {
951 message: String,
953 },
954}
955
956#[derive(Clone)]
1000pub struct IdxTensor {
1001 pub indices: Vec<DynIndex>,
1003 pub(crate) storage: IdxTensorStorage,
1007 pub(crate) eager_cache: Arc<OnceLock<Arc<EagerTensor>>>,
1009}
1010
1011impl IdxTensor {
1012 fn dense_axis_classes(rank: usize) -> Vec<usize> {
1013 (0..rank).collect()
1014 }
1015
1016 fn dtype_name(dtype: DType) -> &'static str {
1017 match dtype {
1018 DType::F32 => "f32",
1019 DType::F64 => "f64",
1020 DType::C32 => "c32",
1021 DType::C64 => "c64",
1022 DType::I32 => "i32",
1023 DType::I64 => "i64",
1024 DType::Bool => "bool",
1025 }
1026 }
1027
1028 fn scalar_dtype(&self) -> Result<DType> {
1029 if let Some(inner) = self.storage.eager() {
1030 return Ok(inner.dtype());
1031 }
1032 if self.storage.is_f64() {
1033 Ok(DType::F64)
1034 } else if self.storage.is_c64() {
1035 Ok(DType::C64)
1036 } else {
1037 Err(anyhow::anyhow!(
1038 "unable to determine IdxTensor scalar dtype"
1039 ))
1040 }
1041 }
1042
1043 fn diag_axis_classes(rank: usize) -> Vec<usize> {
1044 if rank == 0 {
1045 vec![]
1046 } else {
1047 vec![0; rank]
1048 }
1049 }
1050
1051 fn canonicalize_axis_classes(axis_classes: &[usize]) -> Vec<usize> {
1052 let mut map = std::collections::HashMap::new();
1053 let mut next = 0usize;
1054 axis_classes
1055 .iter()
1056 .map(|&class_id| {
1057 *map.entry(class_id).or_insert_with(|| {
1058 let canonical = next;
1059 next += 1;
1060 canonical
1061 })
1062 })
1063 .collect()
1064 }
1065
1066 fn permute_axis_classes(&self, perm: &[usize]) -> Vec<usize> {
1067 let axis_classes = self.storage.axis_classes();
1068 let permuted: Vec<usize> = perm.iter().map(|&index| axis_classes[index]).collect();
1069 Self::canonicalize_axis_classes(&permuted)
1070 }
1071
1072 fn normalize_insert_axis(op: &str, axis: isize, rank: usize) -> Result<usize> {
1073 let normalized = if axis < 0 {
1074 rank as isize + 1 + axis
1075 } else {
1076 axis
1077 };
1078 if !(normalized >= 0 && normalized <= rank as isize) {
1079 return Err(anyhow::anyhow!(
1080 "{op}: axis {axis} is out of bounds for inserting into rank {rank}"
1081 ));
1082 };
1083 Ok(normalized as usize)
1084 }
1085
1086 fn is_diag_axis_classes(axis_classes: &[usize]) -> bool {
1087 axis_classes.len() >= 2 && axis_classes.iter().all(|&class_id| class_id == 0)
1088 }
1089
1090 fn validate_axis_classes(axis_classes: &[usize], rank: usize) -> Result<()> {
1091 if axis_classes.len() != rank {
1092 return Err(anyhow::anyhow!(
1093 "axis-class rank {} does not match tensor rank {rank}",
1094 axis_classes.len()
1095 ));
1096 }
1097 if Self::canonicalize_axis_classes(axis_classes) != axis_classes {
1098 return Err(anyhow::anyhow!(
1099 "axis classes must be canonical first-occurrence labels: {axis_classes:?}"
1100 ));
1101 }
1102 Ok(())
1103 }
1104
1105 fn einsum_subscripts_from_usize_ids(
1106 inputs: &[Vec<usize>],
1107 output: &[usize],
1108 ) -> Result<EinsumSubscripts> {
1109 let input_labels = inputs
1110 .iter()
1111 .map(|ids| {
1112 ids.iter()
1113 .map(|&id| {
1114 u32::try_from(id)
1115 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1116 })
1117 .collect::<Result<Vec<_>>>()
1118 })
1119 .collect::<Result<Vec<_>>>()?;
1120 let output_labels = output
1121 .iter()
1122 .map(|&id| {
1123 u32::try_from(id)
1124 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1125 })
1126 .collect::<Result<Vec<_>>>()?;
1127 let input_refs = input_labels.iter().map(Vec::as_slice).collect::<Vec<_>>();
1128 Ok(EinsumSubscripts::new(&input_refs, &output_labels))
1129 }
1130
1131 fn build_binary_einsum_subscripts(
1132 lhs_rank: usize,
1133 axes_a: &[usize],
1134 rhs_rank: usize,
1135 axes_b: &[usize],
1136 ) -> Result<EinsumSubscripts> {
1137 if !(axes_a.len() == axes_b.len()) {
1138 return Err(anyhow::anyhow!(
1139 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1140 axes_a,
1141 axes_b
1142 ));
1143 };
1144
1145 let mut lhs_ids = vec![usize::MAX; lhs_rank];
1146 let mut rhs_ids = vec![usize::MAX; rhs_rank];
1147 let mut next_id = 0usize;
1148
1149 let mut seen_lhs = vec![false; lhs_rank];
1150 let mut seen_rhs = vec![false; rhs_rank];
1151
1152 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1153 if !(lhs_axis < lhs_rank) {
1154 return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1155 };
1156 if !(rhs_axis < rhs_rank) {
1157 return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1158 };
1159 if !(!seen_lhs[lhs_axis]) {
1160 return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1161 };
1162 if !(!seen_rhs[rhs_axis]) {
1163 return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1164 };
1165 seen_lhs[lhs_axis] = true;
1166 seen_rhs[rhs_axis] = true;
1167 lhs_ids[lhs_axis] = next_id;
1168 rhs_ids[rhs_axis] = next_id;
1169 next_id += 1;
1170 }
1171
1172 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1173 for id in &mut lhs_ids {
1174 if *id == usize::MAX {
1175 *id = next_id;
1176 output_ids.push(next_id);
1177 next_id += 1;
1178 }
1179 }
1180 for id in &mut rhs_ids {
1181 if *id == usize::MAX {
1182 *id = next_id;
1183 output_ids.push(next_id);
1184 next_id += 1;
1185 }
1186 }
1187
1188 Self::einsum_subscripts_from_usize_ids(&[lhs_ids, rhs_ids], &output_ids)
1189 }
1190
1191 fn binary_dot_general_config(axes_a: &[usize], axes_b: &[usize]) -> Result<DotGeneralConfig> {
1192 if !(axes_a.len() == axes_b.len()) {
1193 return Err(anyhow::anyhow!(
1194 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1195 axes_a,
1196 axes_b
1197 ));
1198 };
1199 Ok(DotGeneralConfig {
1200 lhs_contracting_dims: axes_a.to_vec(),
1201 rhs_contracting_dims: axes_b.to_vec(),
1202 lhs_batch_dims: vec![],
1203 rhs_batch_dims: vec![],
1204 })
1205 }
1206
1207 fn binary_contraction_axis_classes(
1208 lhs_axis_classes: &[usize],
1209 axes_a: &[usize],
1210 rhs_axis_classes: &[usize],
1211 axes_b: &[usize],
1212 ) -> Result<Vec<usize>> {
1213 debug_assert_eq!(axes_a.len(), axes_b.len());
1214
1215 fn find(parent: &mut [usize], value: usize) -> usize {
1216 if parent[value] != value {
1217 parent[value] = find(parent, parent[value]);
1218 }
1219 parent[value]
1220 }
1221
1222 fn union(parent: &mut [usize], lhs: usize, rhs: usize) {
1223 let lhs_root = find(parent, lhs);
1224 let rhs_root = find(parent, rhs);
1225 if lhs_root != rhs_root {
1226 parent[rhs_root] = lhs_root;
1227 }
1228 }
1229
1230 let lhs_payload_rank = match lhs_axis_classes.iter().copied().max() {
1231 Some(value) => value
1232 .checked_add(1)
1233 .ok_or_else(|| anyhow::anyhow!("left payload rank overflows usize"))?,
1234 None => 0,
1235 };
1236 let rhs_payload_rank = match rhs_axis_classes.iter().copied().max() {
1237 Some(value) => value
1238 .checked_add(1)
1239 .ok_or_else(|| anyhow::anyhow!("right payload rank overflows usize"))?,
1240 None => 0,
1241 };
1242 let rhs_offset = lhs_payload_rank;
1243 let parent_len = lhs_payload_rank
1244 .checked_add(rhs_payload_rank)
1245 .ok_or_else(|| anyhow::anyhow!("payload rank sum overflows usize"))?;
1246 let mut parent: Vec<usize> = (0..parent_len).collect();
1247
1248 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1249 union(
1250 &mut parent,
1251 lhs_axis_classes[lhs_axis],
1252 rhs_offset
1253 .checked_add(rhs_axis_classes[rhs_axis])
1254 .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?,
1255 );
1256 }
1257
1258 let mut lhs_contracted = vec![false; lhs_axis_classes.len()];
1259 for &axis in axes_a {
1260 lhs_contracted[axis] = true;
1261 }
1262 let mut rhs_contracted = vec![false; rhs_axis_classes.len()];
1263 for &axis in axes_b {
1264 rhs_contracted[axis] = true;
1265 }
1266
1267 let mut root_to_class = std::collections::HashMap::new();
1268 let mut next_class = 0usize;
1269 let mut axis_classes = Vec::with_capacity(
1270 lhs_axis_classes
1271 .len()
1272 .saturating_add(rhs_axis_classes.len()),
1273 );
1274
1275 for (axis, &class_id) in lhs_axis_classes.iter().enumerate() {
1276 if !lhs_contracted[axis] {
1277 let root = find(&mut parent, class_id);
1278 let class = *root_to_class.entry(root).or_insert_with(|| {
1279 let value = next_class;
1280 next_class += 1;
1281 value
1282 });
1283 axis_classes.push(class);
1284 }
1285 }
1286 for (axis, &class_id) in rhs_axis_classes.iter().enumerate() {
1287 if !rhs_contracted[axis] {
1288 let rhs_class = rhs_offset
1289 .checked_add(class_id)
1290 .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?;
1291 let root = find(&mut parent, rhs_class);
1292 let class = *root_to_class.entry(root).or_insert_with(|| {
1293 let value = next_class;
1294 next_class += 1;
1295 value
1296 });
1297 axis_classes.push(class);
1298 }
1299 }
1300
1301 Ok(axis_classes)
1302 }
1303
1304 fn scale_subscripts(rank: usize) -> Result<EinsumSubscripts> {
1305 let ids: Vec<usize> = (0..rank).collect();
1306 Self::einsum_subscripts_from_usize_ids(&[ids.clone(), Vec::new()], &ids)
1307 }
1308
1309 fn scale_target_dtype(payload: DType, scalar: DType) -> Result<DType> {
1310 let target = match payload {
1311 DType::F32 => match scalar {
1312 DType::C32 | DType::C64 => DType::C32,
1313 DType::F32 | DType::F64 => DType::F32,
1314 dtype => {
1315 return Err(anyhow::anyhow!(
1316 "unsupported scalar dtype {dtype:?} for f32 scaling"
1317 ));
1318 }
1319 },
1320 DType::C32 => match scalar {
1321 DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C32,
1322 dtype => {
1323 return Err(anyhow::anyhow!(
1324 "unsupported scalar dtype {dtype:?} for c32 scaling"
1325 ));
1326 }
1327 },
1328 DType::F64 => match scalar {
1329 DType::C32 | DType::C64 => DType::C64,
1330 DType::F32 | DType::F64 => DType::F64,
1331 dtype => {
1332 return Err(anyhow::anyhow!(
1333 "unsupported scalar dtype {dtype:?} for f64 scaling"
1334 ));
1335 }
1336 },
1337 DType::C64 => match scalar {
1338 DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C64,
1339 dtype => {
1340 return Err(anyhow::anyhow!(
1341 "unsupported scalar dtype {dtype:?} for c64 scaling"
1342 ));
1343 }
1344 },
1345 dtype => {
1346 return Err(anyhow::anyhow!(
1347 "unsupported tensor dtype {dtype:?} for scaling"
1348 ));
1349 }
1350 };
1351 Ok(target)
1352 }
1353
1354 fn validate_indices(indices: &[DynIndex]) -> Result<()> {
1355 let mut seen = HashSet::new();
1356 for idx in indices {
1357 if !(seen.insert(idx.clone())) {
1358 return Err(anyhow::anyhow!("Tensor indices must all be unique"));
1359 };
1360 }
1361 Ok(())
1362 }
1363
1364 fn validate_diag_dims(dims: &[usize]) -> Result<()> {
1365 if !dims.is_empty() {
1366 let first_dim = dims[0];
1367 for (i, &dim) in dims.iter().enumerate() {
1368 if !(dim == first_dim) {
1369 return Err(anyhow::anyhow!("DiagTensor requires all indices to have the same dimension, but dims[{i}] = {dim} != dims[0] = {first_dim}"));
1370 };
1371 }
1372 }
1373 Ok(())
1374 }
1375
1376 fn seed_native_payload(storage: &Storage, dims: &[usize]) -> Result<NativeTensor> {
1377 Ok(storage_to_native_tensor(storage, dims)?)
1378 }
1379
1380 fn empty_eager_cache() -> Arc<OnceLock<Arc<EagerTensor>>> {
1381 Arc::new(OnceLock::new())
1382 }
1383
1384 fn eager_cache_with(inner: EagerTensor) -> Arc<OnceLock<Arc<EagerTensor>>> {
1385 let cache = Arc::new(OnceLock::new());
1386 let _ = cache.set(Arc::new(inner));
1387 cache
1388 }
1389
1390 fn compact_payload_inner(&self) -> Result<EagerTensor> {
1391 self.ensure_storage_ready()?;
1392 if let Some(inner) = self.storage.eager() {
1393 return Ok(inner.clone());
1394 }
1395 Ok(EagerTensor::from_tensor_in(
1396 storage_payload_native(self.storage.materialize(self.indices.len())?.as_ref())?,
1397 default_eager_ctx()?,
1398 )?)
1399 }
1400
1401 fn dense_inner_from_payload(
1402 payload: &EagerTensor,
1403 axis_classes: &[usize],
1404 logical_dims: &[usize],
1405 ) -> Result<EagerTensor> {
1406 let payload_rank = match axis_classes.iter().copied().max() {
1407 Some(class_id) => class_id
1408 .checked_add(1)
1409 .ok_or_else(|| anyhow::anyhow!("structured payload class rank overflows usize"))?,
1410 None => 0,
1411 };
1412 if !(payload.shape().len() == payload_rank) {
1413 return Err(anyhow::anyhow!(
1414 "structured payload rank {} does not match axis classes {:?}",
1415 payload.shape().len(),
1416 axis_classes
1417 ));
1418 };
1419 if !(logical_dims.len() == axis_classes.len()) {
1420 return Err(anyhow::anyhow!(
1421 "logical rank {} does not match axis class rank {}",
1422 logical_dims.len(),
1423 axis_classes.len()
1424 ));
1425 };
1426
1427 if axis_classes == Self::dense_axis_classes(logical_dims.len()) {
1428 if !(payload.shape() == logical_dims) {
1429 return Err(anyhow::anyhow!(
1430 "dense payload dims {:?} do not match logical dims {:?}",
1431 payload.shape(),
1432 logical_dims
1433 ));
1434 };
1435 return Ok(payload.clone());
1436 }
1437
1438 let mut first_axis_by_class = vec![None; payload_rank];
1439 let mut dense = payload.clone();
1440 for (logical_axis, &class_id) in axis_classes.iter().enumerate() {
1441 let first_axis = match first_axis_by_class[class_id] {
1442 Some(first_axis) => first_axis,
1443 None => {
1444 first_axis_by_class[class_id] = Some(logical_axis);
1445 continue;
1446 }
1447 };
1448 dense = dense.embed_diag(first_axis, logical_axis)?;
1449 }
1450 if !(dense.shape() == logical_dims) {
1451 return Err(anyhow::anyhow!(
1452 "expanded structured payload dims {:?} do not match logical dims {:?}",
1453 dense.shape(),
1454 logical_dims
1455 ));
1456 };
1457 Ok(dense)
1458 }
1459
1460 fn tracked_compact_payload_value(&self) -> Option<&StructuredPayload> {
1461 self.storage
1462 .deferred_error()
1463 .is_none()
1464 .then_some(self.storage.compact_payload())
1465 .flatten()
1466 .filter(|value| value.payload.tracks_grad())
1467 }
1468
1469 fn ensure_storage_ready(&self) -> Result<()> {
1470 if let Some(error) = self.storage.deferred_error() {
1471 return Err(anyhow::Error::new(error.clone()));
1472 }
1473 Ok(())
1474 }
1475
1476 fn compact_payload_is_logical_dense(&self, payload_dims: &[usize]) -> bool {
1477 self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len())
1478 && payload_dims == self.dims()
1479 }
1480
1481 fn uses_tracked_compact_storage(&self) -> bool {
1482 self.tracked_compact_payload_value()
1483 .is_some_and(|value| !self.compact_payload_is_logical_dense(&value.payload_dims))
1484 }
1485
1486 fn ensure_shape_packing_preserves_ad(&self, op_name: &str) -> Result<()> {
1487 self.ensure_storage_ready()?;
1488 if !(!self.uses_tracked_compact_storage()) {
1489 return Err(anyhow::anyhow!("{op_name}: structured AD tensors with compact storage are not supported because materializing compact storage would detach gradients"));
1490 };
1491 Ok(())
1492 }
1493
1494 fn operand_indices_for_contraction(&self, conjugate: bool) -> Vec<DynIndex> {
1495 if conjugate {
1496 self.indices.iter().map(|index| index.conj()).collect()
1497 } else {
1498 self.indices.clone()
1499 }
1500 }
1501
1502 fn build_binary_contraction_labels(
1503 lhs_rank: usize,
1504 axes_a: &[usize],
1505 rhs_rank: usize,
1506 axes_b: &[usize],
1507 ) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
1508 if !(axes_a.len() == axes_b.len()) {
1509 return Err(anyhow::anyhow!(
1510 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1511 axes_a,
1512 axes_b
1513 ));
1514 };
1515
1516 let mut lhs_ids = vec![usize::MAX; lhs_rank];
1517 let mut rhs_ids = vec![usize::MAX; rhs_rank];
1518 let mut next_id = 0usize;
1519
1520 let mut seen_lhs = vec![false; lhs_rank];
1521 let mut seen_rhs = vec![false; rhs_rank];
1522
1523 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1524 if !(lhs_axis < lhs_rank) {
1525 return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1526 };
1527 if !(rhs_axis < rhs_rank) {
1528 return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1529 };
1530 if !(!seen_lhs[lhs_axis]) {
1531 return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1532 };
1533 if !(!seen_rhs[rhs_axis]) {
1534 return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1535 };
1536 seen_lhs[lhs_axis] = true;
1537 seen_rhs[rhs_axis] = true;
1538 lhs_ids[lhs_axis] = next_id;
1539 rhs_ids[rhs_axis] = next_id;
1540 next_id += 1;
1541 }
1542
1543 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1544 for id in &mut lhs_ids {
1545 if *id == usize::MAX {
1546 *id = next_id;
1547 output_ids.push(next_id);
1548 next_id += 1;
1549 }
1550 }
1551 for id in &mut rhs_ids {
1552 if *id == usize::MAX {
1553 *id = next_id;
1554 output_ids.push(next_id);
1555 next_id += 1;
1556 }
1557 }
1558
1559 Ok((lhs_ids, rhs_ids, output_ids))
1560 }
1561
1562 fn build_payload_einsum_subscripts(
1563 input_roots: &[Vec<usize>],
1564 output_roots: &[usize],
1565 ) -> Result<EinsumSubscripts> {
1566 Self::einsum_subscripts_from_usize_ids(input_roots, output_roots)
1567 }
1568
1569 fn normalize_eager_payload_for_roots(
1570 payload: &EagerTensor,
1571 roots: &[usize],
1572 ) -> Result<(Option<EagerTensor>, Vec<usize>)> {
1573 if !(payload.shape().len() == roots.len()) {
1574 return Err(anyhow::anyhow!(
1575 "payload rank {} does not match root label count {}",
1576 payload.shape().len(),
1577 roots.len()
1578 ));
1579 };
1580
1581 let mut current_payload = None;
1582 let mut current_roots = roots.to_vec();
1583 while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(¤t_roots) {
1584 let source = current_payload.as_ref().unwrap_or(payload);
1585 current_payload = Some(source.extract_diag(axis_a, axis_b)?);
1586 current_roots.remove(axis_b);
1587 }
1588
1589 Ok((current_payload, current_roots))
1590 }
1591
1592 fn first_duplicate_pair(values: &[usize]) -> Option<(usize, usize)> {
1593 let mut first_axis_by_value = std::collections::HashMap::new();
1594 for (axis, &value) in values.iter().enumerate() {
1595 if let Some(&first_axis) = first_axis_by_value.get(&value) {
1596 return Some((first_axis, axis));
1597 }
1598 first_axis_by_value.insert(value, axis);
1599 }
1600 None
1601 }
1602
1603 fn from_structured_payload_inner(
1604 indices: Vec<DynIndex>,
1605 payload_inner: EagerTensor,
1606 payload_dims: Vec<usize>,
1607 axis_classes: Vec<usize>,
1608 ) -> Result<Self> {
1609 Self::validate_indices(&indices)?;
1610 if payload_inner.shape() != payload_dims {
1611 return Err(anyhow::anyhow!(
1612 "structured payload dims {:?} do not match planned payload dims {:?}",
1613 payload_inner.shape(),
1614 payload_dims
1615 ));
1616 }
1617 if axis_classes == Self::dense_axis_classes(indices.len()) {
1618 return Self::from_inner_with_axis_classes(indices, payload_inner, axis_classes);
1619 }
1620 let structured_payload = Arc::new(StructuredPayload {
1621 payload: Arc::new(payload_inner),
1622 payload_dims,
1623 axis_classes,
1624 });
1625 Ok(Self {
1626 indices,
1627 storage: IdxTensorStorage::Compact(structured_payload),
1628 eager_cache: Self::empty_eager_cache(),
1629 })
1630 }
1631
1632 fn contract_structured_payloads(
1633 &self,
1634 other: &Self,
1635 result_indices: Vec<DynIndex>,
1636 axes_a: &[usize],
1637 axes_b: &[usize],
1638 ) -> Result<Self> {
1639 let (lhs_labels, rhs_labels, output_labels) = Self::build_binary_contraction_labels(
1640 self.indices.len(),
1641 axes_a,
1642 other.indices.len(),
1643 axes_b,
1644 )?;
1645 Self::contract_structured_payloads_nary(
1646 &[self, other],
1647 result_indices,
1648 vec![lhs_labels, rhs_labels],
1649 output_labels,
1650 )
1651 }
1652
1653 pub(crate) fn contract_structured_payloads_nary(
1654 operands: &[&Self],
1655 result_indices: Vec<DynIndex>,
1656 input_labels: Vec<Vec<usize>>,
1657 output_labels: Vec<usize>,
1658 ) -> Result<Self> {
1659 if !(!operands.is_empty()) {
1660 return Err(anyhow::anyhow!("structured contraction needs operands"));
1661 };
1662 for operand in operands {
1663 operand.ensure_storage_ready()?;
1664 }
1665 let layouts = operands
1666 .iter()
1667 .map(|operand| {
1668 OperandLayout::new(operand.dims(), operand.storage.axis_classes().to_vec())
1669 })
1670 .collect::<Result<Vec<_>>>()?;
1671 let spec = StructuredContractionSpec {
1672 input_labels,
1673 output_labels,
1674 retained_labels: Default::default(),
1675 };
1676 let plan = StructuredContractionPlan::new(&layouts, &spec)?;
1677 let any_grad = operands.iter().any(|operand| operand.tracks_grad());
1678
1679 if any_grad {
1680 let dtypes = operands
1681 .iter()
1682 .map(|operand| {
1683 operand.storage.dtype().ok_or_else(|| {
1684 anyhow::anyhow!("structured contraction operand has no scalar dtype")
1685 })
1686 })
1687 .collect::<Result<Vec<_>>>()?;
1688 let target = Self::common_eager_dtype(&dtypes)?;
1689 let mut payloads = Vec::with_capacity(operands.len());
1690 for operand in operands {
1691 let payload = operand.compact_payload_inner()?;
1692 payloads.push(if payload.dtype() == target {
1693 payload
1694 } else {
1695 payload.cast(target)?
1696 });
1697 }
1698
1699 let mut normalized = Vec::with_capacity(payloads.len());
1700 let mut labels = Vec::with_capacity(payloads.len());
1701 for (operand_idx, (payload, operand_plan)) in
1702 payloads.iter().zip(plan.operand_plans.iter()).enumerate()
1703 {
1704 let (payload, roots) =
1705 Self::normalize_eager_payload_for_roots(payload, &operand_plan.class_roots)?;
1706 normalized.push(payload.unwrap_or_else(|| payloads[operand_idx].clone()));
1707 labels.push(roots);
1708 }
1709 let refs = normalized.iter().collect::<Vec<_>>();
1710 let subscripts =
1711 Self::build_payload_einsum_subscripts(&labels, &plan.output_payload_roots)?;
1712 let payload = refs.as_slice().einsum_subscripts(&subscripts)?;
1713 return Self::from_structured_payload_inner(
1714 result_indices,
1715 payload,
1716 plan.output_payload_dims,
1717 plan.output_axis_classes,
1718 );
1719 }
1720
1721 let storage_owners = operands
1725 .iter()
1726 .map(|operand| match &operand.storage {
1727 IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1728 IdxTensorStorage::Deferred { source, .. } => match source.as_ref() {
1729 IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1730 _ => None,
1731 },
1732 _ => None,
1733 })
1734 .collect::<Vec<_>>();
1735 let mut inputs = Vec::with_capacity(operands.len());
1736 for (operand_idx, operand) in operands.iter().enumerate() {
1737 if let Some(storage) = storage_owners[operand_idx].as_ref() {
1738 inputs.push(storage_payload_native_read_input(storage.as_ref())?);
1739 } else {
1740 let inner = operand
1741 .storage
1742 .eager()
1743 .ok_or_else(|| anyhow::anyhow!("structured operand has no compact payload"))?;
1744 inputs.push(NativeTensorReadInput::Borrowed(inner.tensor_read()));
1745 }
1746 }
1747 let mut normalized = Vec::with_capacity(inputs.len());
1748 let mut labels = Vec::with_capacity(inputs.len());
1749 for (input, operand_plan) in inputs.into_iter().zip(plan.operand_plans.iter()) {
1750 let (input, roots) =
1751 normalize_payload_read_for_roots(input, &operand_plan.class_roots)?;
1752 normalized.push(input);
1753 labels.push(roots);
1754 }
1755 let refs = normalized
1756 .iter()
1757 .zip(labels.iter())
1758 .map(|(input, labels)| (input, labels.as_slice()))
1759 .collect::<Vec<_>>();
1760 let payload = tensor4all_tensorbackend::einsum_native_tensor_reads(
1761 &refs,
1762 &plan.output_payload_roots,
1763 )?;
1764 let mut common: Option<Arc<EagerRuntime>> = None;
1770 let mut mixed = false;
1771 for operand in operands {
1772 if let Some(inner) = operand.storage.eager() {
1773 match &common {
1774 None => common = Some(Arc::clone(inner.runtime())),
1775 Some(runtime) => {
1776 if runtime.id() != inner.ctx_id() {
1777 mixed = true;
1778 }
1779 }
1780 }
1781 }
1782 }
1783 let payload_inner = match (common, mixed) {
1784 (Some(runtime), false) => EagerTensor::from_tensor_in(payload, runtime)?,
1785 _ => EagerTensor::from_tensor_in(payload, default_eager_ctx()?)?,
1786 };
1787 Self::from_structured_payload_inner(
1788 result_indices,
1789 payload_inner,
1790 plan.output_payload_dims,
1791 plan.output_axis_classes,
1792 )
1793 }
1794
1795 fn common_eager_dtype(dtypes: &[DType]) -> Result<DType> {
1796 let target = if dtypes.contains(&DType::C64)
1797 || (dtypes.contains(&DType::C32) && dtypes.contains(&DType::F64))
1798 {
1799 DType::C64
1800 } else if dtypes.contains(&DType::F64) || dtypes.contains(&DType::C32) {
1801 if dtypes.contains(&DType::C32) {
1802 DType::C32
1803 } else {
1804 DType::F64
1805 }
1806 } else {
1807 DType::F32
1808 };
1809 if !(dtypes
1810 .iter()
1811 .all(|dtype| matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64)))
1812 {
1813 return Err(anyhow::anyhow!(
1814 "structured contraction supports only f32, f64, c32, and c64 operands"
1815 ));
1816 };
1817 Ok(target)
1818 }
1819
1820 fn should_use_structured_payload_contract(&self, other: &Self) -> bool {
1821 self.tracks_grad()
1822 || other.tracks_grad()
1823 || self.storage.axis_classes() != Self::dense_axis_classes(self.indices.len())
1824 || other.storage.axis_classes() != Self::dense_axis_classes(other.indices.len())
1825 }
1826
1827 fn storage_from_native_with_axis_classes(
1828 native: &NativeTensor,
1829 axis_classes: &[usize],
1830 logical_rank: usize,
1831 ) -> Result<Storage> {
1832 if matches!(native.dtype(), DType::F32 | DType::C32) {
1833 return Err(anyhow::anyhow!(
1834 "compact IdxTensor storage does not support dtype {:?}; retain the eager payload",
1835 native.dtype()
1836 ));
1837 }
1838 if Self::is_diag_axis_classes(axis_classes) {
1839 match native.dtype() {
1840 DType::F64 | DType::I32 | DType::I64 | DType::Bool => Storage::from_diag_col_major(
1841 native_tensor_primal_to_diag::<f64>(native)?,
1842 logical_rank,
1843 ),
1844 DType::C64 => Storage::from_diag_col_major(
1845 native_tensor_primal_to_diag::<Complex64>(native)?,
1846 logical_rank,
1847 ),
1848 DType::F32 | DType::C32 => Err(anyhow::anyhow!(
1849 "compact IdxTensor storage does not support dtype {:?}",
1850 native.dtype()
1851 )),
1852 }
1853 } else {
1854 storage_from_payload_native(native.duplicate()?, native.shape(), axis_classes.to_vec())
1855 }
1856 }
1857
1858 fn dense_selected_diag_payload<T: TensorElement + Copy + Zero>(
1859 payload: Vec<T>,
1860 kept_dims: &[usize],
1861 selected_positions: &[usize],
1862 ) -> Result<Vec<T>> {
1863 let output_len = checked_product(kept_dims)?;
1864 let mut data = vec![T::zero(); output_len];
1865 if output_len == 0 {
1866 return Ok(data);
1867 }
1868
1869 let Some((&first_position, rest)) = selected_positions.split_first() else {
1870 return Ok(data);
1871 };
1872 if rest.iter().any(|&position| position != first_position) {
1873 return Ok(data);
1874 }
1875
1876 let value = payload[first_position];
1877 if kept_dims.is_empty() {
1878 data[0] = value;
1879 return Ok(data);
1880 }
1881
1882 let mut offset = 0usize;
1883 let mut stride = 1usize;
1884 for &dim in kept_dims {
1885 let term = first_position
1886 .checked_mul(stride)
1887 .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1888 offset = offset
1889 .checked_add(term)
1890 .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1891 stride = stride
1892 .checked_mul(dim)
1893 .ok_or_else(|| anyhow::anyhow!("diagonal selection stride overflow"))?;
1894 }
1895 data[offset] = value;
1896 Ok(data)
1897 }
1898
1899 fn select_diag_indices(
1900 &self,
1901 kept_indices: Vec<DynIndex>,
1902 kept_dims: Vec<usize>,
1903 positions: &[usize],
1904 ) -> Result<Self> {
1905 if self.storage.is_f64() {
1906 let storage = self.storage.materialize(self.indices.len())?;
1907 let payload = storage
1908 .payload_f64_col_major_vec()
1909 .map_err(anyhow::Error::new)?;
1910 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1911 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1912 } else if self.storage.is_c64() {
1913 let storage = self.storage.materialize(self.indices.len())?;
1914 let payload = storage
1915 .payload_c64_col_major_vec()
1916 .map_err(anyhow::Error::new)?;
1917 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1918 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1919 } else if self.storage.dtype() == Some(DType::F32) {
1920 let inner = self
1921 .storage
1922 .eager()
1923 .ok_or_else(|| anyhow::anyhow!("failed to read f32 diagonal payload"))?;
1924 let payload = inner.value()?.as_slice::<f32>()?.to_vec();
1925 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1926 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1927 } else if self.storage.dtype() == Some(DType::C32) {
1928 let inner = self
1929 .storage
1930 .eager()
1931 .ok_or_else(|| anyhow::anyhow!("failed to read c32 diagonal payload"))?;
1932 let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
1933 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1934 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1935 } else {
1936 Err(anyhow::anyhow!("unsupported diagonal storage scalar type"))
1937 }
1938 }
1939
1940 fn col_major_strides(dims: &[usize]) -> Result<Vec<isize>> {
1941 let mut strides = Vec::with_capacity(dims.len());
1942 let mut stride = 1isize;
1943 for &dim in dims {
1944 strides.push(stride);
1945 let dim = isize::try_from(dim)
1946 .map_err(|_| anyhow::anyhow!("dimension does not fit in isize"))?;
1947 stride = stride
1948 .checked_mul(dim)
1949 .ok_or_else(|| anyhow::anyhow!("column-major stride overflow"))?;
1950 }
1951 Ok(strides)
1952 }
1953
1954 fn zero_structured_selection<T>(
1955 kept_indices: Vec<DynIndex>,
1956 kept_dims: &[usize],
1957 ) -> Result<Self>
1958 where
1959 T: TensorElement + Zero,
1960 {
1961 let output_len = checked_product(kept_dims)?;
1962 Self::from_dense(kept_indices, vec![T::zero(); output_len]).map_err(anyhow::Error::from)
1963 }
1964
1965 fn selected_structured_class_positions(
1966 axis_classes: &[usize],
1967 payload_rank: usize,
1968 selected_axes: &[usize],
1969 positions: &[usize],
1970 ) -> Option<Vec<Option<usize>>> {
1971 let mut selected_class_positions = vec![None; payload_rank];
1972 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
1973 let class_id = axis_classes[axis];
1974 if let Some(existing) = selected_class_positions[class_id] {
1975 if existing != position {
1976 return None;
1977 }
1978 } else {
1979 selected_class_positions[class_id] = Some(position);
1980 }
1981 }
1982 Some(selected_class_positions)
1983 }
1984
1985 fn select_structured_indices_typed<T, F>(
1986 &self,
1987 payload: Vec<T>,
1988 kept_axes: &[usize],
1989 kept_indices: Vec<DynIndex>,
1990 kept_dims: Vec<usize>,
1991 selected: (&[usize], &[usize]),
1992 make_output: F,
1993 ) -> Result<Self>
1994 where
1995 T: TensorElement + Zero,
1996 F: FnOnce(Vec<T>, Vec<usize>, Vec<isize>, Vec<usize>) -> Result<Self>,
1997 {
1998 let (selected_axes, positions) = selected;
1999 let payload_dims = self.storage.payload_dims();
2000 let axis_classes = self.storage.axis_classes();
2001 let payload_rank = payload_dims.len();
2002 let Some(selected_class_positions) = Self::selected_structured_class_positions(
2003 axis_classes,
2004 payload_rank,
2005 selected_axes,
2006 positions,
2007 ) else {
2008 return Self::zero_structured_selection::<T>(kept_indices, &kept_dims);
2009 };
2010
2011 let selected_class_kept = kept_axes
2012 .iter()
2013 .any(|&axis| selected_class_positions[axis_classes[axis]].is_some());
2014 if selected_class_kept {
2015 return self.select_structured_indices_dense(
2016 payload,
2017 kept_axes,
2018 kept_indices,
2019 kept_dims,
2020 &selected_class_positions,
2021 );
2022 }
2023
2024 let mut old_to_new_class = vec![None; payload_rank];
2025 let mut output_payload_dims = Vec::new();
2026 let mut output_axis_classes = Vec::with_capacity(kept_axes.len());
2027 for &axis in kept_axes {
2028 let class_id = axis_classes[axis];
2029 let new_class = match old_to_new_class[class_id] {
2030 Some(new_class) => new_class,
2031 None => {
2032 let new_class = output_payload_dims.len();
2033 old_to_new_class[class_id] = Some(new_class);
2034 output_payload_dims.push(payload_dims[class_id]);
2035 new_class
2036 }
2037 };
2038 output_axis_classes.push(new_class);
2039 }
2040
2041 let output_len = checked_product(&output_payload_dims)?;
2042 let mut output_payload = Vec::with_capacity(output_len);
2043 for linear in 0..output_len {
2044 let output_payload_index = decode_col_major_linear(linear, &output_payload_dims)?;
2045 let mut input_payload_index = vec![0usize; payload_rank];
2046 for class_id in 0..payload_rank {
2047 input_payload_index[class_id] =
2048 if let Some(position) = selected_class_positions[class_id] {
2049 position
2050 } else if let Some(new_class) = old_to_new_class[class_id] {
2051 output_payload_index[new_class]
2052 } else {
2053 return Err(anyhow::anyhow!(
2054 "structured payload class {class_id} is neither selected nor kept"
2055 ));
2056 };
2057 }
2058 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2059 output_payload.push(payload[input_linear]);
2060 }
2061
2062 let output_strides = Self::col_major_strides(&output_payload_dims)?;
2063 make_output(
2064 output_payload,
2065 output_payload_dims,
2066 output_strides,
2067 output_axis_classes,
2068 )
2069 }
2070
2071 fn select_structured_indices_dense<T>(
2072 &self,
2073 payload: Vec<T>,
2074 kept_axes: &[usize],
2075 kept_indices: Vec<DynIndex>,
2076 kept_dims: Vec<usize>,
2077 selected_class_positions: &[Option<usize>],
2078 ) -> Result<Self>
2079 where
2080 T: TensorElement + Zero,
2081 {
2082 let payload_dims = self.storage.payload_dims();
2083 let axis_classes = self.storage.axis_classes();
2084 let output_len = checked_product(&kept_dims)?;
2085 let mut output = Vec::with_capacity(output_len);
2086
2087 for linear in 0..output_len {
2088 let kept_position = decode_col_major_linear(linear, &kept_dims)?;
2089 let mut input_payload_index = selected_class_positions.to_vec();
2090 let mut is_structural_zero = false;
2091
2092 for (&axis, &position) in kept_axes.iter().zip(kept_position.iter()) {
2093 let class_id = axis_classes[axis];
2094 match input_payload_index[class_id] {
2095 Some(existing) if existing != position => {
2096 is_structural_zero = true;
2097 break;
2098 }
2099 Some(_) => {}
2100 None => input_payload_index[class_id] = Some(position),
2101 }
2102 }
2103
2104 if is_structural_zero {
2105 output.push(T::zero());
2106 continue;
2107 }
2108
2109 let input_payload_index = input_payload_index
2110 .into_iter()
2111 .enumerate()
2112 .map(|(class_id, position)| {
2113 position.ok_or_else(|| {
2114 anyhow::anyhow!(
2115 "structured payload class {class_id} is neither selected nor kept"
2116 )
2117 })
2118 })
2119 .collect::<Result<Vec<_>>>()?;
2120 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2121 output.push(payload[input_linear]);
2122 }
2123
2124 Self::from_dense(kept_indices, output).map_err(anyhow::Error::from)
2125 }
2126
2127 fn select_structured_indices(
2128 &self,
2129 kept_axes: &[usize],
2130 kept_indices: Vec<DynIndex>,
2131 kept_dims: Vec<usize>,
2132 selected_axes: &[usize],
2133 positions: &[usize],
2134 ) -> Result<Self> {
2135 if self.storage.is_f64() {
2136 let storage = self.storage.materialize(self.indices.len())?;
2137 let payload = storage
2138 .payload_f64_col_major_vec()
2139 .map_err(anyhow::Error::new)?;
2140 let output_indices = kept_indices.clone();
2141 self.select_structured_indices_typed(
2142 payload,
2143 kept_axes,
2144 kept_indices,
2145 kept_dims,
2146 (selected_axes, positions),
2147 move |payload, dims, strides, classes| {
2148 let storage = Storage::new_structured(payload, dims, strides, classes)?;
2149 Self::from_storage(output_indices, Arc::new(storage))
2150 .map_err(anyhow::Error::from)
2151 },
2152 )
2153 } else if self.storage.is_c64() {
2154 let storage = self.storage.materialize(self.indices.len())?;
2155 let payload = storage
2156 .payload_c64_col_major_vec()
2157 .map_err(anyhow::Error::new)?;
2158 let output_indices = kept_indices.clone();
2159 self.select_structured_indices_typed(
2160 payload,
2161 kept_axes,
2162 kept_indices,
2163 kept_dims,
2164 (selected_axes, positions),
2165 move |payload, dims, strides, classes| {
2166 let storage = Storage::new_structured(payload, dims, strides, classes)?;
2167 Self::from_storage(output_indices, Arc::new(storage))
2168 .map_err(anyhow::Error::from)
2169 },
2170 )
2171 } else if self.storage.dtype() == Some(DType::F32) {
2172 let inner = self
2173 .storage
2174 .eager()
2175 .ok_or_else(|| anyhow::anyhow!("failed to read f32 structured payload"))?;
2176 let payload = inner.value()?.as_slice::<f32>()?.to_vec();
2177 let output_indices = kept_indices.clone();
2178 self.select_structured_indices_typed(
2179 payload,
2180 kept_axes,
2181 kept_indices,
2182 kept_dims,
2183 (selected_axes, positions),
2184 move |payload, dims, _strides, classes| {
2185 let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2186 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2187 Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2188 },
2189 )
2190 } else if self.storage.dtype() == Some(DType::C32) {
2191 let inner = self
2192 .storage
2193 .eager()
2194 .ok_or_else(|| anyhow::anyhow!("failed to read c32 structured payload"))?;
2195 let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
2196 let output_indices = kept_indices.clone();
2197 self.select_structured_indices_typed(
2198 payload,
2199 kept_axes,
2200 kept_indices,
2201 kept_dims,
2202 (selected_axes, positions),
2203 move |payload, dims, _strides, classes| {
2204 let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2205 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2206 Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2207 },
2208 )
2209 } else {
2210 Err(anyhow::anyhow!(
2211 "unsupported structured storage scalar type"
2212 ))
2213 }
2214 }
2215
2216 fn validate_storage_matches_indices(indices: &[DynIndex], storage: &Storage) -> Result<()> {
2217 let dims = Self::expected_dims_from_indices(indices);
2218 let storage_dims = storage.logical_dims();
2219 if storage_dims != dims {
2220 return Err(anyhow::anyhow!(
2221 "storage logical dims {:?} do not match indices dims {:?}",
2222 storage_dims,
2223 dims
2224 ));
2225 }
2226 if storage.is_diag() {
2227 Self::validate_diag_dims(&dims)?;
2228 }
2229 Ok(())
2230 }
2231
2232 fn try_materialized_inner(&self) -> Result<&EagerTensor> {
2233 self.ensure_storage_ready()?;
2234 let logical_dims = self.dims();
2235 if let Some(value) = self.tracked_compact_payload_value() {
2236 if self.compact_payload_is_logical_dense(&value.payload_dims) {
2237 return Ok(value.payload.as_ref());
2238 }
2239 if self.eager_cache.get().is_none() {
2240 let dense = Self::dense_inner_from_payload(
2241 value.payload.as_ref(),
2242 &value.axis_classes,
2243 &logical_dims,
2244 )?;
2245 let _ = self.eager_cache.set(Arc::new(dense));
2246 }
2247 return self
2248 .eager_cache
2249 .get()
2250 .map(|inner| inner.as_ref())
2251 .ok_or_else(|| {
2252 anyhow::anyhow!("IdxTensor structured AD cache was not initialized")
2253 });
2254 }
2255 if let Some(inner) = self.storage.eager() {
2256 if self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len()) {
2257 return Ok(inner);
2258 }
2259 if self.eager_cache.get().is_none() {
2260 let dense = Self::dense_inner_from_payload(
2261 inner,
2262 self.storage.axis_classes(),
2263 &logical_dims,
2264 )?;
2265 let _ = self.eager_cache.set(Arc::new(dense));
2266 }
2267 return self
2268 .eager_cache
2269 .get()
2270 .map(|inner| inner.as_ref())
2271 .ok_or_else(|| {
2272 anyhow::anyhow!("IdxTensor structured eager cache was not initialized")
2273 });
2274 }
2275 if self.eager_cache.get().is_none() {
2276 let native = profile_pairwise_contract_section("materialize_storage_to_native", || {
2277 let storage = self.storage.materialize(self.indices.len())?;
2278 Self::seed_native_payload(storage.as_ref(), &logical_dims)
2279 })
2280 .context("IdxTensor materialization failed")?;
2281 record_pairwise_contract_profile_bytes(
2282 "materialize_storage_to_native",
2283 tensor_profile_bytes(native.dtype(), native.shape()),
2284 );
2285 let _ = self.eager_cache.set(Arc::new(EagerTensor::from_tensor_in(
2286 native,
2287 default_eager_ctx()?,
2288 )?));
2289 }
2290 self.eager_cache
2291 .get()
2292 .map(|inner| inner.as_ref())
2293 .ok_or_else(|| anyhow::anyhow!("IdxTensor materialization cache was not initialized"))
2294 }
2295
2296 pub(crate) fn as_inner(&self) -> Result<&EagerTensor> {
2297 self.try_materialized_inner()
2298 }
2299
2300 #[inline]
2302 fn expected_dims_from_indices(indices: &[DynIndex]) -> Vec<usize> {
2303 indices.iter().map(|idx| idx.dim()).collect()
2304 }
2305
2306 pub fn dims(&self) -> Vec<usize> {
2325 Self::expected_dims_from_indices(&self.indices)
2326 }
2327
2328 pub fn select_indices(
2371 &self,
2372 selected_indices: &[DynIndex],
2373 positions: &[usize],
2374 ) -> std::result::Result<Self, IdxTensorError> {
2375 if selected_indices.len() != positions.len() {
2376 return Err(anyhow::anyhow!(
2377 "selected_indices length {} does not match positions length {}",
2378 selected_indices.len(),
2379 positions.len()
2380 )
2381 .into());
2382 }
2383 if selected_indices.is_empty() {
2384 return Ok(self.clone());
2385 }
2386
2387 let mut selected_axes = Vec::with_capacity(selected_indices.len());
2388 let mut seen_axes = HashSet::with_capacity(selected_indices.len());
2389 for (selected, &position) in selected_indices.iter().zip(positions.iter()) {
2390 let axis = self
2391 .indices
2392 .iter()
2393 .position(|index| index == selected)
2394 .ok_or_else(|| anyhow::anyhow!("selected index is not present in tensor"))?;
2395 if !seen_axes.insert(axis) {
2396 return Err(anyhow::anyhow!("selected index appears more than once").into());
2397 }
2398 let dim = self.indices[axis].dim();
2399 if position >= dim {
2400 return Err(anyhow::anyhow!(
2401 "selected coordinate {position} is out of range for axis {axis} with dim {dim}"
2402 )
2403 .into());
2404 }
2405 selected_axes.push(axis);
2406 }
2407
2408 let kept_axes = self
2409 .indices
2410 .iter()
2411 .enumerate()
2412 .filter(|(axis, _)| !seen_axes.contains(axis))
2413 .map(|(axis, _)| axis)
2414 .collect::<Vec<_>>();
2415 let kept_indices = kept_axes
2416 .iter()
2417 .map(|&axis| self.indices[axis].clone())
2418 .collect::<Vec<_>>();
2419 let kept_dims = kept_axes
2420 .iter()
2421 .map(|&axis| self.indices[axis].dim())
2422 .collect::<Vec<_>>();
2423
2424 if matches!(
2425 self.storage.storage_kind(),
2426 StorageKind::Diagonal | StorageKind::Structured
2427 ) {
2428 self.ensure_shape_packing_preserves_ad("select_indices")?;
2429 }
2430 if self.storage.storage_kind() == StorageKind::Diagonal {
2431 return self
2432 .select_diag_indices(kept_indices, kept_dims, positions)
2433 .map_err(IdxTensorError::from);
2434 }
2435 if self.storage.storage_kind() == StorageKind::Structured {
2436 return self
2437 .select_structured_indices(
2438 &kept_axes,
2439 kept_indices,
2440 kept_dims,
2441 &selected_axes,
2442 positions,
2443 )
2444 .map_err(IdxTensorError::from);
2445 }
2446 if self.storage.storage_kind() != StorageKind::Dense {
2447 return Err(anyhow::anyhow!(
2448 "select_indices got unsupported storage kind {:?}",
2449 self.storage.storage_kind()
2450 )
2451 .into());
2452 }
2453
2454 let mut sliced = self.try_materialized_inner()?.clone();
2459 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
2460 sliced = sliced
2461 .slice_axis(axis, position..position + 1)
2462 .map_err(|error| anyhow::anyhow!("select_indices slicing failed: {error}"))?;
2463 }
2464 Self::from_inner(kept_indices, sliced.reshape(&kept_dims)?).map_err(IdxTensorError::from)
2465 }
2466
2467 pub fn stack_along_new_index(
2500 tensors: &[&Self],
2501 new_index: DynIndex,
2502 axis: isize,
2503 ) -> std::result::Result<Self, IdxTensorError> {
2504 let first = tensors
2505 .first()
2506 .copied()
2507 .ok_or_else(|| anyhow::anyhow!("stack_along_new_index requires at least one tensor"))?;
2508 if !(new_index.dim() == tensors.len()) {
2509 return Err(anyhow::anyhow!(
2510 "stack_along_new_index: new index dim {} does not match tensor count {}",
2511 new_index.dim(),
2512 tensors.len()
2513 )
2514 .into());
2515 };
2516
2517 let base_indices = first.indices.clone();
2518 for tensor in tensors.iter().copied().skip(1) {
2519 if !(tensor.indices == base_indices) {
2520 return Err(anyhow::anyhow!(
2521 "stack_along_new_index: input tensors must have identical index order"
2522 )
2523 .into());
2524 };
2525 }
2526 for &tensor in tensors {
2527 tensor.ensure_shape_packing_preserves_ad("stack_along_new_index")?;
2528 }
2529
2530 let insert_axis =
2531 Self::normalize_insert_axis("stack_along_new_index", axis, base_indices.len())?;
2532 let mut result_indices = base_indices;
2533 result_indices.insert(insert_axis, new_index);
2534
2535 let inner_refs = tensors
2536 .iter()
2537 .map(|tensor| tensor.try_materialized_inner())
2538 .collect::<Result<Vec<_>>>()?;
2539 let stacked = EagerTensor::stack(&inner_refs, axis)?;
2540 Self::from_inner(result_indices, stacked).map_err(IdxTensorError::from)
2541 }
2542
2543 pub fn index_select(
2573 &self,
2574 source_index: &DynIndex,
2575 target_index: DynIndex,
2576 positions: &[usize],
2577 ) -> std::result::Result<Self, IdxTensorError> {
2578 if !(target_index.dim() == positions.len()) {
2579 return Err(anyhow::anyhow!(
2580 "index_select: target index dim {} does not match position count {}",
2581 target_index.dim(),
2582 positions.len()
2583 )
2584 .into());
2585 };
2586 let axis = self
2587 .indices
2588 .iter()
2589 .position(|index| index == source_index)
2590 .ok_or_else(|| anyhow::anyhow!("index_select: source index is not present"))?;
2591 let source_dim = self.indices[axis].dim();
2592 for &position in positions {
2593 if !(position < source_dim) {
2594 return Err(anyhow::anyhow!(
2595 "index_select: position {position} is out of range for source dim {source_dim}"
2596 )
2597 .into());
2598 };
2599 }
2600 self.ensure_shape_packing_preserves_ad("index_select")?;
2601
2602 let axis = isize::try_from(axis)
2603 .map_err(|_| anyhow::anyhow!("index_select: axis does not fit in isize"))?;
2604 let selected = self
2605 .try_materialized_inner()?
2606 .index_select(axis, positions)?;
2607 let mut result_indices = self.indices.clone();
2608 result_indices[axis as usize] = target_index;
2609 Self::from_inner(result_indices, selected).map_err(IdxTensorError::from)
2610 }
2611
2612 pub fn new(
2631 indices: Vec<DynIndex>,
2632 storage: Arc<Storage>,
2633 ) -> std::result::Result<Self, IdxTensorError> {
2634 Self::from_storage(indices, storage)
2635 }
2636
2637 pub fn from_indices(
2658 indices: Vec<DynIndex>,
2659 storage: Arc<Storage>,
2660 ) -> std::result::Result<Self, IdxTensorError> {
2661 Self::new(indices, storage)
2662 }
2663
2664 pub fn from_storage(
2684 indices: Vec<DynIndex>,
2685 storage: Arc<Storage>,
2686 ) -> std::result::Result<Self, IdxTensorError> {
2687 Self::validate_indices(&indices)?;
2688 Self::validate_storage_matches_indices(&indices, storage.as_ref())?;
2689 Ok(Self {
2690 indices,
2691 storage: IdxTensorStorage::from_storage(storage),
2692 eager_cache: Self::empty_eager_cache(),
2693 })
2694 }
2695
2696 pub fn from_structured_storage(
2718 indices: Vec<DynIndex>,
2719 storage: Arc<Storage>,
2720 ) -> std::result::Result<Self, IdxTensorError> {
2721 Self::from_storage(indices, storage)
2722 }
2723
2724 pub fn from_copy_selector<T>(
2780 left: DynIndex,
2781 site: DynIndex,
2782 right: DynIndex,
2783 selected_value: usize,
2784 scale: T,
2785 ) -> std::result::Result<Self, StructuredSelectorError>
2786 where
2787 T: TensorElement + Copy + Zero,
2788 {
2789 if left.dim == 0 {
2790 return Err(StructuredSelectorError::ZeroDimension { axis: "left" });
2791 }
2792 if site.dim == 0 {
2793 return Err(StructuredSelectorError::ZeroDimension { axis: "site" });
2794 }
2795 if right.dim == 0 {
2796 return Err(StructuredSelectorError::ZeroDimension { axis: "right" });
2797 }
2798 if left.dim != right.dim {
2799 return Err(StructuredSelectorError::BondDimensionMismatch {
2800 left: left.dim,
2801 right: right.dim,
2802 });
2803 }
2804 if selected_value >= site.dim {
2805 return Err(StructuredSelectorError::SelectedValueOutOfBounds {
2806 value: selected_value,
2807 site_dim: site.dim,
2808 });
2809 }
2810
2811 let payload_len =
2812 left.dim
2813 .checked_mul(site.dim)
2814 .ok_or(StructuredSelectorError::PayloadSizeOverflow {
2815 bond_dim: left.dim,
2816 site_dim: site.dim,
2817 })?;
2818 let _site_stride = isize::try_from(left.dim)
2819 .map_err(|_| StructuredSelectorError::StrideOverflow { bond_dim: left.dim })?;
2820 let selected_offset = left.dim.checked_mul(selected_value).ok_or(
2821 StructuredSelectorError::PayloadSizeOverflow {
2822 bond_dim: left.dim,
2823 site_dim: site.dim,
2824 },
2825 )?;
2826
2827 let mut payload = Vec::new();
2828 payload.try_reserve_exact(payload_len).map_err(|_| {
2829 StructuredSelectorError::AllocationFailed {
2830 elements: payload_len,
2831 }
2832 })?;
2833 payload.resize(payload_len, T::zero());
2834 for bond in 0..left.dim {
2835 payload[selected_offset + bond] = scale;
2836 }
2837 let payload_native = dense_native_tensor_from_col_major(&payload, &[left.dim, site.dim])
2838 .map_err(|error| StructuredSelectorError::InvalidStorage {
2839 message: error.to_string(),
2840 })?;
2841 let payload_dtype = payload_native.dtype();
2842 let payload_inner = EagerTensor::from_tensor_in(
2843 payload_native,
2844 default_eager_ctx().map_err(|error| StructuredSelectorError::InvalidStorage {
2845 message: error.to_string(),
2846 })?,
2847 )
2848 .map_err(|error| StructuredSelectorError::InvalidStorage {
2849 message: error.to_string(),
2850 })?;
2851 let payload_dims = vec![left.dim, site.dim];
2852 let indices = vec![left, site, right];
2853 if !matches!(
2854 payload_dtype,
2855 DType::F32 | DType::F64 | DType::C32 | DType::C64
2856 ) {
2857 return Err(StructuredSelectorError::InvalidStorage {
2858 message: format!("unsupported selector dtype {:?}", payload_dtype),
2859 });
2860 }
2861 Self::from_structured_payload_inner(indices, payload_inner, payload_dims, vec![0, 1, 0])
2862 .map_err(|error| StructuredSelectorError::InvalidStorage {
2863 message: error.to_string(),
2864 })
2865 }
2866
2867 pub(crate) fn from_native(indices: Vec<DynIndex>, native: NativeTensor) -> Result<Self> {
2869 let axis_classes = Self::dense_axis_classes(indices.len());
2870 Self::from_native_with_axis_classes(indices, native, axis_classes)
2871 }
2872
2873 pub(crate) fn from_native_with_axis_classes(
2874 indices: Vec<DynIndex>,
2875 native: NativeTensor,
2876 axis_classes: Vec<usize>,
2877 ) -> Result<Self> {
2878 Self::from_inner_with_axis_classes(
2879 indices,
2880 EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
2881 axis_classes,
2882 )
2883 }
2884
2885 pub(crate) fn from_untracked_native_with_axis_classes(
2888 indices: Vec<DynIndex>,
2889 native: NativeTensor,
2890 axis_classes: Vec<usize>,
2891 ) -> Result<Self> {
2892 Self::from_inner_with_validated_metadata(
2893 indices,
2894 adopt_untracked_eager_value(default_eager_ctx()?, TensorValue::from_tensor(native))?,
2895 axis_classes,
2896 )
2897 }
2898
2899 pub(crate) fn from_inner(indices: Vec<DynIndex>, inner: EagerTensor) -> Result<Self> {
2900 let axis_classes = Self::dense_axis_classes(indices.len());
2901 Self::from_inner_with_axis_classes(indices, inner, axis_classes)
2902 }
2903
2904 pub fn hermitian_eigendecomposition(
2944 &self,
2945 hermitian_tol: f64,
2946 ) -> std::result::Result<TensorHermitianEigendecomposition, IdxTensorError> {
2947 if !(self.indices.len() == 2) {
2948 return Err(anyhow::anyhow!(
2949 "IdxTensor::hermitian_eigendecomposition requires a rank-2 tensor, got rank {}",
2950 self.indices.len()
2951 )
2952 .into());
2953 };
2954 let dims = self.dims();
2955 if !(dims[0] == dims[1]) {
2956 return Err(anyhow::anyhow!(
2957 "IdxTensor::hermitian_eigendecomposition requires a square matrix, got {}x{}",
2958 dims[0],
2959 dims[1]
2960 )
2961 .into());
2962 };
2963 if !(dims[0] > 0) {
2964 return Err(anyhow::anyhow!(
2965 "IdxTensor::hermitian_eigendecomposition requires a non-empty matrix"
2966 )
2967 .into());
2968 };
2969 if !(hermitian_tol.is_finite() && hermitian_tol >= 0.0) {
2970 return Err(anyhow::anyhow!(
2971 "IdxTensor::hermitian_eigendecomposition requires a finite non-negative tolerance"
2972 )
2973 .into());
2974 };
2975
2976 let input = self.try_materialized_inner()?;
2977 let (values, vectors) = input
2978 .eigh()
2979 .map_err(|source| anyhow::anyhow!("Hermitian eigendecomposition failed: {source}"))?;
2980
2981 let eigenvalue_index = DynIndex::new_dyn(dims[0]);
2982 let eigenvector_index = DynIndex::new_dyn(dims[0]);
2983 let eigenvalue_tensor = Self::from_inner(vec![eigenvalue_index], values)?;
2984 let eigenvalues = Self::read_real_eigenvalues(&eigenvalue_tensor, hermitian_tol)
2985 .with_context(|| {
2986 "IdxTensor::hermitian_eigendecomposition failed to read eigenvalues"
2987 })?;
2988 let eigenvectors = Self::from_inner(
2989 vec![self.indices[0].clone(), eigenvector_index.clone()],
2990 vectors,
2991 )?;
2992
2993 Ok(TensorHermitianEigendecomposition {
2994 eigenvalues,
2995 eigenvectors,
2996 eigenvector_index,
2997 })
2998 }
2999
3000 fn read_real_eigenvalues(values: &Self, hermitian_tol: f64) -> Result<Vec<f64>> {
3001 if values.is_complex() {
3002 values
3003 .to_vec::<Complex64>()?
3004 .into_iter()
3005 .enumerate()
3006 .map(|(index, value)| {
3007 let imaginary = value.im.abs();
3008 let allowed = hermitian_tol * value.norm().max(1.0);
3009 if !matches!(
3010 imaginary.partial_cmp(&allowed),
3011 Some(std::cmp::Ordering::Less) | Some(std::cmp::Ordering::Equal)
3012 ) {
3013 return Err(anyhow::anyhow!("Hermitian eigenvalue {index} has imaginary part {imaginary}, exceeding tolerance {allowed}"));
3014 };
3015 Ok(value.re)
3016 })
3017 .collect()
3018 } else {
3019 values.to_vec::<f64>().map_err(anyhow::Error::from)
3020 }
3021 }
3022
3023 pub(crate) fn from_diag_inner(
3024 indices: Vec<DynIndex>,
3025 payload_inner: EagerTensor,
3026 ) -> Result<Self> {
3027 let dims = Self::expected_dims_from_indices(&indices);
3028 Self::validate_indices(&indices)?;
3029 Self::validate_diag_dims(&dims)?;
3030 let payload_len = checked_product(payload_inner.shape())?;
3031 Self::validate_diag_payload_len(payload_len, &dims)?;
3032 let axis_classes = Self::diag_axis_classes(dims.len());
3033 let diag_inner = payload_inner.embed_diag(0, 1)?;
3034 Self::from_inner_with_axis_classes(indices, diag_inner, axis_classes)
3035 }
3036
3037 fn compact_inner_from_logical(
3038 inner: &EagerTensor,
3039 axis_classes: &[usize],
3040 ) -> Result<EagerTensor> {
3041 let mut payload = inner.clone();
3042 let mut classes = axis_classes.to_vec();
3043 while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(&classes) {
3044 payload = payload.extract_diag(axis_a, axis_b)?;
3045 classes.remove(axis_b);
3046 }
3047 Ok(payload)
3048 }
3049
3050 pub(crate) fn from_inner_with_axis_classes(
3051 indices: Vec<DynIndex>,
3052 inner: EagerTensor,
3053 axis_classes: Vec<usize>,
3054 ) -> Result<Self> {
3055 Self::from_inner_with_axis_classes_impl(indices, inner, axis_classes, false)
3056 }
3057
3058 fn from_inner_with_validated_metadata(
3059 indices: Vec<DynIndex>,
3060 inner: EagerTensor,
3061 axis_classes: Vec<usize>,
3062 ) -> Result<Self> {
3063 Self::from_inner_with_axis_classes_impl(indices, inner, axis_classes, true)
3064 }
3065
3066 fn from_inner_with_axis_classes_impl(
3067 indices: Vec<DynIndex>,
3068 inner: EagerTensor,
3069 axis_classes: Vec<usize>,
3070 validated: bool,
3071 ) -> Result<Self> {
3072 let dims = profile_pairwise_contract_section("from_inner_expected_dims", || {
3073 Self::expected_dims_from_indices(&indices)
3074 });
3075 let dense_axis_classes = axis_classes.iter().copied().eq(0..indices.len());
3076 if validated {
3077 debug_assert!(Self::validate_indices(&indices).is_ok());
3078 if !dense_axis_classes {
3079 Self::validate_axis_classes(&axis_classes, indices.len())?;
3080 }
3081 } else {
3082 profile_pairwise_contract_section("from_inner_validate_indices", || {
3083 Self::validate_indices(&indices)
3084 })?;
3085 Self::validate_axis_classes(&axis_classes, indices.len())?;
3086 }
3087 if dims != inner.shape() {
3088 return Err(anyhow::anyhow!(
3089 "native payload dims {:?} do not match indices dims {:?}",
3090 inner.shape(),
3091 dims
3092 ));
3093 }
3094 if Self::is_diag_axis_classes(&axis_classes) {
3095 profile_pairwise_contract_section("from_inner_validate_diag_dims", || {
3096 Self::validate_diag_dims(&dims)
3097 })?;
3098 }
3099 let storage = if dense_axis_classes {
3100 IdxTensorStorage::from_eager_dense(inner, indices.len())
3101 } else {
3102 let payload = Self::compact_inner_from_logical(&inner, &axis_classes)?;
3103 let payload_dims = payload.shape().to_vec();
3104 IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3105 payload: Arc::new(payload),
3106 payload_dims,
3107 axis_classes,
3108 }))
3109 };
3110 Ok(Self {
3111 indices,
3112 storage,
3113 eager_cache: Self::empty_eager_cache(),
3114 })
3115 }
3116
3117 pub fn indices(&self) -> &[DynIndex] {
3119 &self.indices
3120 }
3121
3122 pub(crate) fn axis_classes(&self) -> &[usize] {
3123 self.storage.axis_classes()
3124 }
3125
3126 #[cfg(feature = "tenferro-cuda")]
3127 pub(crate) fn deferred_storage_error(&self) -> Option<&TensorStorageError> {
3128 self.storage.deferred_error()
3129 }
3130
3131 #[cfg(feature = "tenferro-cuda")]
3132 pub(crate) fn cuda_eager_inner(&self) -> Option<&EagerTensor> {
3133 self.storage
3134 .eager()
3135 .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3136 }
3137
3138 pub(crate) fn eager_runtime(&self) -> Option<Arc<EagerRuntime>> {
3143 self.storage
3144 .eager()
3145 .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3146 .map(|inner| Arc::clone(inner.runtime()))
3147 }
3148
3149 #[cfg(feature = "tenferro-cuda")]
3156 pub(crate) fn is_cuda_resident(&self) -> bool {
3157 self.cuda_eager_inner().is_some_and(|inner| {
3158 inner.tensor_read().placement().memory_kind == tenferro_tensor::MemoryKind::Device
3159 })
3160 }
3161
3162 #[cfg(not(feature = "tenferro-cuda"))]
3166 pub(crate) fn is_cuda_resident(&self) -> bool {
3167 false
3168 }
3169
3170 #[cfg(feature = "tenferro-cuda")]
3171 pub(crate) fn cuda_duplicate_native(&self) -> Result<NativeTensor> {
3172 Ok(self.try_materialized_inner()?.duplicate_value()?)
3173 }
3174
3175 pub fn enable_grad(self) -> std::result::Result<Self, IdxTensorError> {
3181 self.ensure_storage_ready()?;
3182 let eager_payload = self
3185 .storage
3186 .eager()
3187 .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3188 .filter(|inner| inner.shape() == self.storage.payload_dims());
3189 let payload = match eager_payload {
3190 Some(inner) => inner.duplicate_value()?,
3191 None => {
3192 let materialized = self.storage.materialize(self.indices.len())?;
3193 storage_payload_native(materialized.as_ref())
3194 .context("IdxTensor::enable_grad failed")?
3195 }
3196 };
3197 let payload_dims = self.storage.payload_dims().to_vec();
3198 let axis_classes = self.storage.axis_classes().to_vec();
3199 let tracked = Arc::new(EagerTensor::requires_grad_in(
3200 payload,
3201 default_eager_ctx()?,
3202 )?);
3203 let storage = if axis_classes == Self::dense_axis_classes(self.indices.len()) {
3204 IdxTensorStorage::Eager {
3205 inner: tracked,
3206 axis_classes,
3207 }
3208 } else {
3209 IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3210 payload: tracked,
3211 payload_dims,
3212 axis_classes,
3213 }))
3214 };
3215 Ok(Self {
3216 indices: self.indices,
3217 storage,
3218 eager_cache: Self::empty_eager_cache(),
3219 })
3220 }
3221
3222 pub fn tracks_grad(&self) -> bool {
3224 self.storage.eager().is_some_and(EagerTensor::tracks_grad)
3225 || self
3226 .eager_cache
3227 .get()
3228 .is_some_and(|inner| inner.tracks_grad())
3229 }
3230
3231 pub fn grad(&self) -> std::result::Result<Option<Self>, IdxTensorError> {
3237 if let Some(value) = self.tracked_compact_payload_value() {
3238 let Some(gradient) = value.payload.grad()? else {
3239 return Ok(None);
3240 };
3241 let gradient_shape = gradient.shape().to_vec();
3242 let gradient_tensor = gradient.to_tensor()?;
3243 if self.compact_payload_is_logical_dense(&value.payload_dims) {
3244 return Ok(Some(Self::from_native_with_axis_classes(
3245 self.indices.clone(),
3246 gradient_tensor,
3247 value.axis_classes.clone(),
3248 )?));
3249 }
3250 if gradient_shape != value.payload_dims {
3251 return Err(anyhow::anyhow!(
3252 "gradient payload dims {:?} do not match {:?}",
3253 gradient_shape,
3254 value.payload_dims
3255 )
3256 .into());
3257 }
3258 let gradient = EagerTensor::from_tensor_in(gradient_tensor, default_eager_ctx()?)?;
3259 return Ok(Some(Self::from_structured_payload_inner(
3260 self.indices.clone(),
3261 gradient,
3262 value.payload_dims.clone(),
3263 value.axis_classes.clone(),
3264 )?));
3265 }
3266
3267 let Some(gradient) = self.try_materialized_inner()?.grad()? else {
3268 return Ok(None);
3269 };
3270 Ok(Some(Self::from_native_with_axis_classes(
3271 self.indices.clone(),
3272 gradient.to_tensor()?,
3273 self.storage.axis_classes().to_vec(),
3274 )?))
3275 }
3276
3277 pub fn clear_grad(&self) -> std::result::Result<(), IdxTensorError> {
3283 self.ensure_storage_ready()?;
3284 if let Some(value) = self.tracked_compact_payload_value() {
3285 value.payload.clear_grad()?;
3286 }
3287 if let Some(inner) = self.storage.eager() {
3288 inner.clear_grad()?;
3289 }
3290 if let Some(inner) = self.eager_cache.get() {
3291 inner.clear_grad()?;
3292 }
3293 Ok(())
3294 }
3295
3296 pub fn backward(&self) -> std::result::Result<(), IdxTensorError> {
3302 if let Some(value) = self.tracked_compact_payload_value() {
3303 return value.payload.backward().map(|_| ()).map_err(|e| {
3304 IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}"))
3305 });
3306 }
3307 self.try_materialized_inner()?
3308 .backward()
3309 .map(|_| ())
3310 .map_err(|e| IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}")))
3311 }
3312
3313 pub fn detach(&self) -> std::result::Result<Self, IdxTensorError> {
3319 Self::from_inner_with_axis_classes(
3320 self.indices.clone(),
3321 self.try_materialized_inner()?.detach(),
3322 self.storage.axis_classes().to_vec(),
3323 )
3324 .map_err(IdxTensorError::from)
3325 }
3326
3327 pub fn is_simple(&self) -> bool {
3329 true
3330 }
3331
3332 pub fn to_storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3343 self.storage.materialize(self.indices.len())
3344 }
3345
3346 pub fn storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3365 self.storage.materialize(self.indices.len())
3366 }
3367
3368 pub fn storage_kind(&self) -> StorageKind {
3387 self.storage.storage_kind()
3388 }
3389
3390 #[cfg(feature = "tenferro-cuda")]
3419 pub fn cuda_dtype(&self) -> std::result::Result<DType, IdxTensorError> {
3420 self.scalar_dtype().map_err(IdxTensorError::from)
3421 }
3422
3423 pub fn sum(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3439 self.ensure_storage_ready()?;
3440 if self.indices.is_empty() {
3441 return AnyScalar::from_tensor(self.clone()).map_err(IdxTensorError::from);
3442 }
3443 if let Some(payload) = self.storage.eager().filter(|payload| payload.tracks_grad()) {
3444 let axes: Vec<usize> = (0..payload.shape().len()).collect();
3445 let reduced = payload.reduce_sum(Some(&axes))?;
3446 return AnyScalar::from_tensor(Self::from_inner(Vec::new(), reduced)?)
3447 .map_err(IdxTensorError::from);
3448 }
3449 self.storage.sum_scalar().map_err(IdxTensorError::from)
3450 }
3451
3452 pub fn only(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3476 let dims = self.dims();
3477 let total_size = checked_product(&dims)?;
3478 if !(total_size == 1 || dims.is_empty()) {
3479 return Err(anyhow::anyhow!(
3480 "only() requires a scalar tensor (1 element), got {} elements with dims {:?}",
3481 if dims.is_empty() { 1 } else { total_size },
3482 dims
3483 )
3484 .into());
3485 };
3486 self.sum()
3487 }
3488
3489 pub fn permute_indices(
3524 &self,
3525 new_indices: &[DynIndex],
3526 ) -> std::result::Result<Self, IdxTensorError> {
3527 let perm = compute_permutation_from_indices(&self.indices, new_indices)?;
3529 if perm.iter().copied().eq(0..perm.len()) {
3530 return Ok(Self {
3531 indices: new_indices.to_vec(),
3532 storage: self.storage.clone(),
3533 eager_cache: Arc::clone(&self.eager_cache),
3534 });
3535 }
3536
3537 let permuted = self.try_materialized_inner()?.transpose(&perm)?;
3538 let axis_classes = self.permute_axis_classes(&perm);
3539 Self::from_inner_with_axis_classes(new_indices.to_vec(), permuted, axis_classes)
3540 .map_err(IdxTensorError::from)
3541 }
3542
3543 pub fn permute(&self, perm: &[usize]) -> std::result::Result<Self, IdxTensorError> {
3575 if !(perm.len() == self.indices.len()) {
3576 return Err(anyhow::anyhow!("permutation length must match tensor rank").into());
3577 };
3578 let mut seen = HashSet::new();
3579 for &axis in perm {
3580 if !(axis < self.indices.len()) {
3581 return Err(anyhow::anyhow!("permutation axis {axis} out of range").into());
3582 };
3583 if !(seen.insert(axis)) {
3584 return Err(anyhow::anyhow!("duplicate axis {axis} in permutation").into());
3585 };
3586 }
3587 if perm.iter().copied().eq(0..perm.len()) {
3588 return Ok(self.clone());
3589 }
3590
3591 let new_indices: Vec<DynIndex> = perm.iter().map(|&i| self.indices[i].clone()).collect();
3593 let permuted = self.try_materialized_inner()?.transpose(perm)?;
3594 let axis_classes = self.permute_axis_classes(perm);
3595 Self::from_inner_with_axis_classes(new_indices, permuted, axis_classes)
3596 .map_err(IdxTensorError::from)
3597 }
3598
3599 pub(crate) fn try_contract_pairwise_default(&self, other: &Self) -> Result<Self> {
3600 self.try_contract_pairwise_default_with_options(other, PairwiseContractionOptions::new())
3601 }
3602
3603 pub(crate) fn try_contract_pairwise_default_with_options(
3604 &self,
3605 other: &Self,
3606 options: PairwiseContractionOptions,
3607 ) -> Result<Self> {
3608 #[cfg(feature = "tenferro-cuda")]
3614 if self.is_cuda_resident()
3615 && other.is_cuda_resident()
3616 && self.should_use_structured_payload_contract(other)
3617 {
3618 return super::contract::contract_pair_via_plan(self, other, options);
3619 }
3620 let self_indices = profile_pairwise_contract_section("operand_indices", || {
3621 self.operand_indices_for_contraction(options.lhs_conj)
3622 });
3623 let other_indices = profile_pairwise_contract_section("operand_indices", || {
3624 other.operand_indices_for_contraction(options.rhs_conj)
3625 });
3626 let self_dims = profile_pairwise_contract_section("expected_dims", || {
3627 Self::expected_dims_from_indices(&self_indices)
3628 });
3629 let other_dims = profile_pairwise_contract_section("expected_dims", || {
3630 Self::expected_dims_from_indices(&other_indices)
3631 });
3632 let spec = profile_pairwise_contract_section("prepare_contraction", || {
3633 prepare_contraction(&self_indices, &self_dims, &other_indices, &other_dims)
3634 })
3635 .context("contraction preparation failed")?;
3636 let result_axis_classes = profile_pairwise_contract_section("result_axis_classes", || {
3637 Self::binary_contraction_axis_classes(
3638 self.storage.axis_classes(),
3639 &spec.axes_a,
3640 other.storage.axis_classes(),
3641 &spec.axes_b,
3642 )
3643 })?;
3644
3645 if profile_pairwise_contract_section("structured_check", || {
3646 self.should_use_structured_payload_contract(other)
3647 }) {
3648 if options.has_conj() {
3649 let lhs = if options.lhs_conj {
3650 self.conj()
3651 } else {
3652 self.clone()
3653 };
3654 let rhs = if options.rhs_conj {
3655 other.conj()
3656 } else {
3657 other.clone()
3658 };
3659 return profile_pairwise_contract_section("structured_conj_fallback", || {
3660 lhs.try_contract_pairwise_default(&rhs)
3661 });
3662 }
3663 return profile_pairwise_contract_section("structured_payload_contract", || {
3664 self.contract_structured_payloads(
3665 other,
3666 spec.result_indices.into_vec(),
3667 &spec.axes_a,
3668 &spec.axes_b,
3669 )
3670 });
3671 }
3672
3673 if self.indices.is_empty() && other.indices.is_empty() {
3674 if options.has_conj() {
3675 let lhs = if options.lhs_conj {
3676 self.conj()
3677 } else {
3678 self.clone()
3679 };
3680 let rhs = if options.rhs_conj {
3681 other.conj()
3682 } else {
3683 other.clone()
3684 };
3685 return lhs.try_contract_pairwise_default(&rhs);
3686 }
3687 let result = profile_pairwise_contract_section("scalar_mul", || {
3688 Ok::<_, anyhow::Error>(
3689 self.try_materialized_inner()?
3690 .mul(other.try_materialized_inner()?)?,
3691 )
3692 })?;
3693 return profile_pairwise_contract_section("from_inner", || {
3694 Self::from_inner(spec.result_indices.into_vec(), result)
3695 });
3696 }
3697
3698 let self_dtype = self.try_materialized_inner()?.dtype();
3699 let other_dtype = other.try_materialized_inner()?.dtype();
3700 if self_dtype != other_dtype {
3701 if options.has_conj() {
3702 let lhs = if options.lhs_conj {
3703 self.conj()
3704 } else {
3705 self.clone()
3706 };
3707 let rhs = if options.rhs_conj {
3708 other.conj()
3709 } else {
3710 other.clone()
3711 };
3712 return lhs.try_contract_pairwise_default(&rhs);
3713 }
3714 let self_native = self.try_materialized_inner()?.duplicate_value()?;
3715 let other_native = other.try_materialized_inner()?.duplicate_value()?;
3716 let result_native = profile_pairwise_contract_section("native_contract", || {
3717 contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)
3718 })?;
3719 return profile_pairwise_contract_section("from_native", || {
3720 Self::from_untracked_native_with_axis_classes(
3721 spec.result_indices.into_vec(),
3722 result_native,
3723 result_axis_classes,
3724 )
3725 });
3726 }
3727
3728 let config = profile_pairwise_contract_section("build_dot_general_config", || {
3729 Self::binary_dot_general_config(&spec.axes_a, &spec.axes_b)
3730 })?;
3731 let result = profile_pairwise_contract_section("dot_general_with_conj", || {
3732 let lhs = profile_pairwise_contract_section("lhs_try_materialized_inner", || {
3733 self.try_materialized_inner()
3734 })?;
3735 let rhs = profile_pairwise_contract_section("rhs_try_materialized_inner", || {
3736 other.try_materialized_inner()
3737 })?;
3738 profile_pairwise_contract_section("dot_general_execute", || {
3739 lhs.dot_general_with_conj(rhs, config, options.lhs_conj, options.rhs_conj)
3740 })
3741 .map_err(anyhow::Error::from)
3742 })?;
3743 record_pairwise_contract_profile_bytes(
3744 "dot_general_output",
3745 tensor_profile_bytes(result.dtype(), result.shape()),
3746 );
3747 profile_pairwise_contract_section("from_inner_axis_classes", || {
3748 Self::from_inner_with_axis_classes(
3749 spec.result_indices.into_vec(),
3750 result,
3751 result_axis_classes,
3752 )
3753 })
3754 }
3755
3756 pub(crate) fn try_contract_pairwise_retaining(
3761 &self,
3762 other: &Self,
3763 retained_indices: &[DynIndex],
3764 ) -> Result<Self> {
3765 if retained_indices.is_empty()
3766 || self.should_use_structured_payload_contract(other)
3767 || self.try_materialized_inner()?.dtype() != other.try_materialized_inner()?.dtype()
3768 {
3769 let options = ContractionOptions::new().with_retain_indices(retained_indices);
3770 return super::contract::contract_with_options(&[self, other], options)
3771 .map_err(anyhow::Error::from);
3772 }
3773
3774 let common = common_ind_positions(&self.indices, &other.indices);
3775 if common.is_empty() {
3776 let options = ContractionOptions::new().with_retain_indices(retained_indices);
3777 return super::contract::contract_with_options(&[self, other], options)
3778 .map_err(anyhow::Error::from);
3779 }
3780
3781 let mut retained_pairs = Vec::with_capacity(retained_indices.len());
3782 for retained in retained_indices {
3783 let Some(pos_a) = self.indices.iter().position(|index| index == retained) else {
3784 let options = ContractionOptions::new().with_retain_indices(retained_indices);
3785 return super::contract::contract_with_options(&[self, other], options)
3786 .map_err(anyhow::Error::from);
3787 };
3788 let Some(pos_b) = other.indices.iter().position(|index| index == retained) else {
3789 let options = ContractionOptions::new().with_retain_indices(retained_indices);
3790 return super::contract::contract_with_options(&[self, other], options)
3791 .map_err(anyhow::Error::from);
3792 };
3793 if !self.indices[pos_a].is_contractable(&other.indices[pos_b]) {
3794 let options = ContractionOptions::new().with_retain_indices(retained_indices);
3795 return super::contract::contract_with_options(&[self, other], options)
3796 .map_err(anyhow::Error::from);
3797 }
3798 retained_pairs.push((pos_a, pos_b));
3799 }
3800
3801 let retained_pairs_set: HashSet<(usize, usize)> = retained_pairs.iter().copied().collect();
3802 let mut contracting_a = Vec::with_capacity(common.len());
3803 let mut contracting_b = Vec::with_capacity(common.len());
3804 for &(pos_a, pos_b) in &common {
3805 if !retained_pairs_set.contains(&(pos_a, pos_b)) {
3806 contracting_a.push(pos_a);
3807 contracting_b.push(pos_b);
3808 }
3809 }
3810
3811 let config = DotGeneralConfig {
3812 lhs_contracting_dims: contracting_a.clone(),
3813 rhs_contracting_dims: contracting_b.clone(),
3814 lhs_batch_dims: retained_pairs.iter().map(|&(pos_a, _)| pos_a).collect(),
3815 rhs_batch_dims: retained_pairs.iter().map(|&(_, pos_b)| pos_b).collect(),
3816 };
3817 let result = self
3818 .try_materialized_inner()?
3819 .dot_general_with_conj(other.try_materialized_inner()?, config, false, false)
3820 .map_err(|error| anyhow::anyhow!("retained pairwise contraction failed: {error}"))?;
3821
3822 let contracting_a_set: HashSet<usize> = contracting_a.into_iter().collect();
3828 let contracting_b_set: HashSet<usize> = contracting_b.into_iter().collect();
3829 let retained_a_set: HashSet<usize> = retained_pairs.iter().map(|&(a, _)| a).collect();
3830 let retained_b_set: HashSet<usize> = retained_pairs.iter().map(|&(_, b)| b).collect();
3831 let mut current_indices = self
3832 .indices
3833 .iter()
3834 .enumerate()
3835 .filter(|(axis, _)| !retained_a_set.contains(axis) && !contracting_a_set.contains(axis))
3836 .map(|(_, index)| index.clone())
3837 .chain(
3838 other
3839 .indices
3840 .iter()
3841 .enumerate()
3842 .filter(|(axis, _)| {
3843 !retained_b_set.contains(axis) && !contracting_b_set.contains(axis)
3844 })
3845 .map(|(_, index)| index.clone()),
3846 )
3847 .collect::<Vec<_>>();
3848 current_indices.extend(
3849 retained_pairs
3850 .iter()
3851 .map(|&(pos_a, _)| self.indices[pos_a].clone()),
3852 );
3853
3854 let mut desired_indices = self
3855 .indices
3856 .iter()
3857 .enumerate()
3858 .filter(|(axis, _)| !contracting_a_set.contains(axis))
3859 .map(|(_, index)| index.clone())
3860 .chain(
3861 other
3862 .indices
3863 .iter()
3864 .enumerate()
3865 .filter(|(axis, _)| {
3866 !contracting_b_set.contains(axis) && !retained_b_set.contains(axis)
3867 })
3868 .map(|(_, index)| index.clone()),
3869 )
3870 .collect::<Vec<_>>();
3871 if desired_indices.is_empty() {
3872 desired_indices = current_indices.clone();
3873 }
3874 let result = if current_indices == desired_indices {
3875 result
3876 } else {
3877 let permutation = desired_indices
3878 .iter()
3879 .map(|desired| {
3880 current_indices
3881 .iter()
3882 .position(|current| current == desired)
3883 .ok_or_else(|| {
3884 anyhow::anyhow!(
3885 "retained pairwise result is missing index {:?}",
3886 desired
3887 )
3888 })
3889 })
3890 .collect::<Result<Vec<_>>>()?;
3891 result.transpose(&permutation)?
3892 };
3893 Self::from_inner(desired_indices, result)
3894 }
3895
3896 pub(crate) fn try_tensordot_pairwise_explicit(
3897 &self,
3898 other: &Self,
3899 pairs: &[(DynIndex, DynIndex)],
3900 ) -> Result<Self> {
3901 use crate::index_ops::ContractionError;
3902
3903 let self_dims = Self::expected_dims_from_indices(&self.indices);
3904 let other_dims = Self::expected_dims_from_indices(&other.indices);
3905 let spec = prepare_contraction_pairs(
3906 &self.indices,
3907 &self_dims,
3908 &other.indices,
3909 &other_dims,
3910 pairs,
3911 )
3912 .map_err(|e| match e {
3913 ContractionError::NoCommonIndices => {
3914 anyhow::anyhow!("tensordot: No pairs specified for contraction")
3915 }
3916 ContractionError::BatchContractionNotImplemented => anyhow::anyhow!(
3917 "tensordot: Common index found but not in contraction pairs. \
3918 Batch contraction is not yet implemented."
3919 ),
3920 ContractionError::IndexNotFound { tensor } => {
3921 anyhow::anyhow!("tensordot: Index not found in {} tensor", tensor)
3922 }
3923 ContractionError::DimensionMismatch {
3924 pos_a,
3925 pos_b,
3926 dim_a,
3927 dim_b,
3928 } => anyhow::anyhow!(
3929 "tensordot: Dimension mismatch: self[{}]={} != other[{}]={}",
3930 pos_a,
3931 dim_a,
3932 pos_b,
3933 dim_b
3934 ),
3935 ContractionError::DuplicateAxis { tensor, pos } => {
3936 anyhow::anyhow!("tensordot: Duplicate axis {} in {} tensor", pos, tensor)
3937 }
3938 })?;
3939 let result_axis_classes = Self::binary_contraction_axis_classes(
3940 self.storage.axis_classes(),
3941 &spec.axes_a,
3942 other.storage.axis_classes(),
3943 &spec.axes_b,
3944 )?;
3945
3946 if self.should_use_structured_payload_contract(other) {
3947 return self.contract_structured_payloads(
3948 other,
3949 spec.result_indices.into_vec(),
3950 &spec.axes_a,
3951 &spec.axes_b,
3952 );
3953 }
3954
3955 if self.indices.is_empty() && other.indices.is_empty() {
3956 let result = self
3957 .try_materialized_inner()?
3958 .mul(other.try_materialized_inner()?)
3959 .map_err(|e| anyhow::anyhow!("tensordot scalar multiply failed: {e}"))?;
3960 return Self::from_inner(spec.result_indices.into_vec(), result);
3961 }
3962
3963 let self_dtype = self.try_materialized_inner()?.dtype();
3964 let other_dtype = other.try_materialized_inner()?.dtype();
3965 if self_dtype != other_dtype {
3966 let self_native = self.try_materialized_inner()?.duplicate_value()?;
3967 let other_native = other.try_materialized_inner()?.duplicate_value()?;
3968 let result_native =
3969 contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)?;
3970 return Self::from_untracked_native_with_axis_classes(
3971 spec.result_indices.into_vec(),
3972 result_native,
3973 result_axis_classes,
3974 );
3975 }
3976
3977 let subscripts = Self::build_binary_einsum_subscripts(
3978 self.indices.len(),
3979 &spec.axes_a,
3980 other.indices.len(),
3981 &spec.axes_b,
3982 )?;
3983 let result = [
3984 self.try_materialized_inner()?,
3985 other.try_materialized_inner()?,
3986 ]
3987 .einsum_subscripts(&subscripts)
3988 .map_err(|e| anyhow::anyhow!("tensordot failed: {e}"))?;
3989 Self::from_inner_with_axis_classes(
3990 spec.result_indices.into_vec(),
3991 result,
3992 result_axis_classes,
3993 )
3994 }
3995
3996 pub(crate) fn try_outer_product_pairwise(&self, other: &Self) -> Result<Self> {
3997 use anyhow::Context;
3998
3999 let common_positions = common_ind_positions(&self.indices, &other.indices);
4001 if !common_positions.is_empty() {
4002 let common_ids: Vec<_> = common_positions
4003 .iter()
4004 .map(|(pos_a, _)| self.indices[*pos_a].id())
4005 .collect();
4006 return Err(anyhow::anyhow!(
4007 "outer_product: tensors have common indices {:?}. \
4008 Use tensordot to contract common indices, or use sim() to replace \
4009 indices with fresh IDs before computing outer product.",
4010 common_ids
4011 ))
4012 .context("outer_product: common indices found");
4013 }
4014
4015 let mut result_indices = self.indices.clone();
4017 result_indices.extend(other.indices.iter().cloned());
4018 let result_axis_classes = Self::binary_contraction_axis_classes(
4019 self.storage.axis_classes(),
4020 &[],
4021 other.storage.axis_classes(),
4022 &[],
4023 )?;
4024 if self.should_use_structured_payload_contract(other) {
4025 return self.contract_structured_payloads(other, result_indices, &[], &[]);
4026 }
4027 let self_dtype = self.try_materialized_inner()?.dtype();
4028 let other_dtype = other.try_materialized_inner()?.dtype();
4029 if self_dtype != other_dtype {
4030 let self_native = self.try_materialized_inner()?.duplicate_value()?;
4031 let other_native = other.try_materialized_inner()?.duplicate_value()?;
4032 let result_native = contract_native_tensor(&self_native, &[], &other_native, &[])?;
4033 return Self::from_untracked_native_with_axis_classes(
4034 result_indices,
4035 result_native,
4036 result_axis_classes,
4037 );
4038 }
4039
4040 let subscripts = Self::build_binary_einsum_subscripts(
4041 self.indices.len(),
4042 &[],
4043 other.indices.len(),
4044 &[],
4045 )?;
4046 let result = [
4047 self.try_materialized_inner()?,
4048 other.try_materialized_inner()?,
4049 ]
4050 .einsum_subscripts(&subscripts)
4051 .map_err(|e| anyhow::anyhow!("outer_product failed: {e}"))?;
4052 Self::from_inner_with_axis_classes(result_indices, result, result_axis_classes)
4053 }
4054}
4055
4056impl IdxTensor {
4061 pub fn random<T: RandomScalar, R: Rng>(
4091 rng: &mut R,
4092 indices: Vec<DynIndex>,
4093 ) -> std::result::Result<Self, IdxTensorError> {
4094 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
4095 let size = checked_product(&dims)?;
4096 let data: Vec<T> = (0..size).map(|_| T::random_value(rng)).collect();
4097 Self::from_dense(indices, data)
4098 }
4099}
4100
4101impl IdxTensor {
4102 pub fn add(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
4140 if self.indices.len() != other.indices.len() {
4142 return Err(anyhow::anyhow!(
4143 "Index count mismatch: self has {} indices, other has {}",
4144 self.indices.len(),
4145 other.indices.len()
4146 )
4147 .into());
4148 }
4149
4150 let self_set: HashSet<_> = self.indices.iter().collect();
4152 let other_set: HashSet<_> = other.indices.iter().collect();
4153
4154 if self_set != other_set {
4155 return Err(
4156 anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
4157 );
4158 }
4159
4160 let other_aligned = other.permute_indices(&self.indices)?;
4162
4163 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
4165 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
4166 if self_expected_dims != other_expected_dims {
4167 use crate::TagSetLike;
4168 let fmt = |indices: &[DynIndex]| -> Vec<String> {
4169 indices
4170 .iter()
4171 .map(|idx| {
4172 let tags: Vec<String> = idx.tags().iter().collect();
4173 format!("{:?}(dim={},tags={:?})", idx.id(), idx.dim(), tags)
4174 })
4175 .collect()
4176 };
4177 return Err(anyhow::anyhow!(
4178 "Dimension mismatch after alignment.\n\
4179 self: dims={:?}, indices(order)={:?}\n\
4180 other_aligned: dims={:?}, indices(order)={:?}",
4181 self_expected_dims,
4182 fmt(&self.indices),
4183 other_expected_dims,
4184 fmt(&other_aligned.indices)
4185 )
4186 .into());
4187 }
4188
4189 self.axpby(
4190 AnyScalar::new_real(1.0),
4191 &other_aligned,
4192 AnyScalar::new_real(1.0),
4193 )
4194 }
4195
4196 pub fn axpby(
4221 &self,
4222 a: AnyScalar,
4223 other: &Self,
4224 b: AnyScalar,
4225 ) -> std::result::Result<Self, IdxTensorError> {
4226 if self.indices.len() != other.indices.len() {
4228 return Err(anyhow::anyhow!(
4229 "Index count mismatch: self has {} indices, other has {}",
4230 self.indices.len(),
4231 other.indices.len()
4232 )
4233 .into());
4234 }
4235
4236 let self_set: HashSet<_> = self.indices.iter().collect();
4238 let other_set: HashSet<_> = other.indices.iter().collect();
4239 if self_set != other_set {
4240 return Err(
4241 anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
4242 );
4243 }
4244
4245 let other_aligned = other.permute_indices(&self.indices)?;
4247
4248 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
4250 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
4251 if self_expected_dims != other_expected_dims {
4252 return Err(anyhow::anyhow!(
4253 "Dimension mismatch after alignment: self={:?}, other_aligned={:?}",
4254 self_expected_dims,
4255 other_expected_dims
4256 )
4257 .into());
4258 }
4259
4260 let axis_classes = if self.storage.axis_classes() == other_aligned.storage.axis_classes() {
4261 self.storage.axis_classes().to_vec()
4262 } else {
4263 Self::dense_axis_classes(self.indices.len())
4264 };
4265
4266 let same_compact_layout = self.storage.payload_dims()
4267 == other_aligned.storage.payload_dims()
4268 && self.storage.payload_strides_vec() == other_aligned.storage.payload_strides_vec()
4269 && self.storage.axis_classes() == other_aligned.storage.axis_classes();
4270 if same_compact_layout
4271 && matches!(&self.storage, IdxTensorStorage::Materialized(_))
4272 && matches!(&other_aligned.storage, IdxTensorStorage::Materialized(_))
4273 && !self.tracks_grad()
4274 && !other_aligned.tracks_grad()
4275 && !a.tracks_grad()
4276 && !b.tracks_grad()
4277 {
4278 let lhs_storage = self.storage.materialize(self.indices.len())?;
4279 let rhs_storage = other_aligned
4280 .storage
4281 .materialize(other_aligned.indices.len())?;
4282 let combined = lhs_storage
4283 .axpby(
4284 &a.to_backend_scalar(),
4285 rhs_storage.as_ref(),
4286 &b.to_backend_scalar(),
4287 )
4288 .map_err(|e| anyhow::anyhow!("storage axpby failed: {e}"))?;
4289 return Self::from_storage(self.indices.clone(), Arc::new(combined));
4290 }
4291
4292 let lhs = self.scale(a)?;
4293 let rhs = other_aligned.scale(b)?;
4294 let combined = lhs
4295 .try_materialized_inner()?
4296 .add(rhs.try_materialized_inner()?)
4297 .map_err(|e| anyhow::anyhow!("tensor addition failed: {e}"))?;
4298 Self::from_inner_with_axis_classes(self.indices.clone(), combined, axis_classes)
4299 .map_err(IdxTensorError::from)
4300 }
4301
4302 pub fn scale(&self, scalar: AnyScalar) -> std::result::Result<Self, IdxTensorError> {
4321 if matches!(
4322 &self.storage,
4323 IdxTensorStorage::Eager { .. }
4324 | IdxTensorStorage::Compact(_)
4325 | IdxTensorStorage::Materialized(_)
4326 ) {
4327 let storage = self.storage.scale_eager_payload(&scalar)?;
4332 return Ok(Self {
4333 indices: self.indices.clone(),
4334 storage,
4335 eager_cache: Self::empty_eager_cache(),
4336 });
4337 }
4338
4339 let self_dtype = self.try_materialized_inner()?.dtype();
4340 let scalar_dtype = scalar.as_tensor()?.try_materialized_inner()?.dtype();
4341 if self_dtype != scalar_dtype {
4342 let target_dtype = Self::scale_target_dtype(self_dtype, scalar_dtype)?;
4343 let self_inner = self.try_materialized_inner()?;
4344 let self_inner = if self_inner.dtype() == target_dtype {
4345 self_inner.clone()
4346 } else {
4347 self_inner.cast(target_dtype)?
4348 };
4349 let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
4350 let scalar_inner = if scalar_inner.dtype() == target_dtype {
4351 scalar_inner.clone()
4352 } else {
4353 scalar_inner.cast(target_dtype)?
4354 };
4355 let scaled = if self.indices.is_empty() {
4356 self_inner
4357 .mul(&scalar_inner)
4358 .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4359 } else {
4360 let subscripts = Self::scale_subscripts(self.indices.len())?;
4361 [&self_inner, &scalar_inner]
4362 .einsum_subscripts(&subscripts)
4363 .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4364 };
4365 return Self::from_inner_with_axis_classes(
4366 self.indices.clone(),
4367 scaled,
4368 self.storage.axis_classes().to_vec(),
4369 )
4370 .map_err(IdxTensorError::from);
4371 }
4372 let scaled = if self.indices.is_empty() {
4373 self.try_materialized_inner()?
4374 .mul(scalar.as_tensor()?.try_materialized_inner()?)
4375 .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4376 } else {
4377 let subscripts = Self::scale_subscripts(self.indices.len())?;
4378 [
4379 self.try_materialized_inner()?,
4380 scalar.as_tensor()?.try_materialized_inner()?,
4381 ]
4382 .einsum_subscripts(&subscripts)
4383 .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4384 };
4385 Self::from_inner_with_axis_classes(
4386 self.indices.clone(),
4387 scaled,
4388 self.storage.axis_classes().to_vec(),
4389 )
4390 .map_err(IdxTensorError::from)
4391 }
4392
4393 pub fn inner_product(&self, other: &Self) -> std::result::Result<AnyScalar, IdxTensorError> {
4414 if self.indices.len() == other.indices.len() {
4415 let self_set: HashSet<_> = self.indices.iter().collect();
4416 let other_set: HashSet<_> = other.indices.iter().collect();
4417 if self_set == other_set {
4418 let other_aligned = other.permute_indices(&self.indices)?;
4419 let result = super::contract::contract_pair_with_operand_options(
4420 self,
4421 &other_aligned,
4422 PairwiseContractionOptions::new().with_lhs_conj(true),
4423 )?;
4424 return result.sum();
4425 }
4426 }
4427
4428 let result = super::contract::contract_pair_with_operand_options(
4430 self,
4431 other,
4432 PairwiseContractionOptions::new().with_lhs_conj(true),
4433 )?;
4434 result.sum()
4436 }
4437}
4438
4439impl IdxTensor {
4444 pub fn replaceind(
4479 &self,
4480 old_index: &DynIndex,
4481 new_index: &DynIndex,
4482 ) -> std::result::Result<Self, IdxTensorError> {
4483 if old_index.dim() != new_index.dim() {
4485 return Err(IdxTensorError::ShapeMismatch {
4486 operation: "replaceind",
4487 expected: format!("dimension {}", old_index.dim()),
4488 actual: format!("dimension {}", new_index.dim()),
4489 });
4490 }
4491
4492 let new_indices: Vec<_> = self
4493 .indices
4494 .iter()
4495 .map(|idx| {
4496 if *idx == *old_index {
4497 new_index.clone()
4498 } else {
4499 idx.clone()
4500 }
4501 })
4502 .collect();
4503
4504 Ok(Self {
4505 indices: new_indices,
4506 storage: self.storage.clone(),
4507 eager_cache: Arc::clone(&self.eager_cache),
4508 })
4509 }
4510
4511 pub fn replace_indices(
4551 &self,
4552 old_indices: &[DynIndex],
4553 new_indices: &[DynIndex],
4554 ) -> std::result::Result<Self, IdxTensorError> {
4555 if old_indices.len() != new_indices.len() {
4556 return Err(IdxTensorError::ShapeMismatch {
4557 operation: "replace_indices",
4558 expected: format!("{} indices", old_indices.len()),
4559 actual: format!("{} indices", new_indices.len()),
4560 });
4561 }
4562
4563 for (old, new) in old_indices.iter().zip(new_indices.iter()) {
4565 if old.dim() != new.dim() {
4566 return Err(IdxTensorError::ShapeMismatch {
4567 operation: "replace_indices",
4568 expected: format!("dimension {}", old.dim()),
4569 actual: format!("dimension {}", new.dim()),
4570 });
4571 }
4572 }
4573
4574 let replacement_map: std::collections::HashMap<_, _> =
4576 old_indices.iter().zip(new_indices.iter()).collect();
4577
4578 let new_indices_vec: Vec<_> = self
4579 .indices
4580 .iter()
4581 .map(|idx| {
4582 if let Some(new_idx) = replacement_map.get(idx) {
4583 (*new_idx).clone()
4584 } else {
4585 idx.clone()
4586 }
4587 })
4588 .collect();
4589
4590 Ok(Self {
4591 indices: new_indices_vec,
4592 storage: self.storage.clone(),
4593 eager_cache: Arc::clone(&self.eager_cache),
4594 })
4595 }
4596}
4597
4598impl IdxTensor {
4603 pub fn conj(&self) -> Self {
4632 self.conj_with(&conjugate_eager)
4633 }
4634
4635 fn conj_with<F>(&self, conjugate: &F) -> Self
4636 where
4637 F: Fn(
4638 &EagerTensor,
4639 ) -> std::result::Result<
4640 EagerTensor,
4641 Arc<dyn std::error::Error + Send + Sync + 'static>,
4642 >,
4643 {
4644 let new_indices: Vec<DynIndex> = self.indices.iter().map(|idx| idx.conj()).collect();
4648 let mut storage = match self.storage.conjugate_with(conjugate) {
4649 Ok(storage) => storage,
4650 Err(error) => self.storage.clone().with_deferred_error(error),
4651 };
4652 let mut eager_cache = if storage.deferred_error().is_some() {
4653 Arc::clone(&self.eager_cache)
4654 } else {
4655 Self::empty_eager_cache()
4656 };
4657
4658 if storage.deferred_error().is_none() {
4659 if let Some(inner) = self.eager_cache.get() {
4660 match conjugate(inner.as_ref()) {
4661 Ok(conjugated) => eager_cache = Self::eager_cache_with(conjugated),
4662 Err(source) => {
4663 storage =
4664 storage.with_deferred_error(TensorStorageError::Conjugation { source });
4665 eager_cache = Arc::clone(&self.eager_cache);
4668 }
4669 }
4670 }
4671 }
4672
4673 Self {
4674 indices: new_indices,
4675 storage,
4676 eager_cache,
4677 }
4678 }
4679}
4680
4681#[derive(Debug, Default)]
4682struct Lassq {
4683 scale: f64,
4684 sumsq: f64,
4685 infinite: bool,
4686}
4687
4688impl Lassq {
4689 fn add_component(&mut self, value: f64) {
4690 let value = value.abs();
4691 if value == 0.0 {
4692 return;
4693 }
4694 if value.is_infinite() {
4695 self.infinite = true;
4696 return;
4697 }
4698 if self.scale < value {
4699 if self.scale == 0.0 {
4700 self.sumsq = 1.0;
4701 } else {
4702 let ratio = self.scale / value;
4703 self.sumsq = 1.0 + self.sumsq * ratio * ratio;
4704 }
4705 self.scale = value;
4706 } else {
4707 let ratio = value / self.scale;
4708 self.sumsq += ratio * ratio;
4709 }
4710 }
4711
4712 fn add_complex(&mut self, value: Complex64) {
4713 self.add_component(value.re);
4714 self.add_component(value.im);
4715 }
4716
4717 fn add_scaled(&mut self, scale: f64, coefficient: f64) {
4718 if scale == 0.0 || coefficient == 0.0 {
4719 return;
4720 }
4721 if self.scale < scale {
4722 if self.scale == 0.0 {
4723 self.sumsq = coefficient * coefficient;
4724 } else {
4725 let ratio = self.scale / scale;
4726 self.sumsq = coefficient * coefficient + self.sumsq * ratio * ratio;
4727 }
4728 self.scale = scale;
4729 } else {
4730 let ratio = scale / self.scale * coefficient;
4731 self.sumsq += ratio * ratio;
4732 }
4733 }
4734
4735 fn add_component_difference(&mut self, lhs: f64, rhs: f64) {
4736 if lhs == rhs {
4737 return;
4738 }
4739 let scale = lhs.abs().max(rhs.abs());
4740 if scale != 0.0 {
4741 self.add_scaled(scale, (lhs / scale - rhs / scale).abs());
4742 }
4743 }
4744
4745 fn add_complex_difference(&mut self, lhs: Complex64, rhs: Complex64) {
4746 self.add_component_difference(lhs.re, rhs.re);
4747 self.add_component_difference(lhs.im, rhs.im);
4748 }
4749
4750 fn is_zero(&self) -> bool {
4751 !self.infinite && self.scale == 0.0
4752 }
4753
4754 fn norm(&self) -> f64 {
4755 if self.infinite {
4756 f64::INFINITY
4757 } else if self.scale == 0.0 {
4758 0.0
4759 } else {
4760 self.scale * self.sumsq.sqrt()
4761 }
4762 }
4763
4764 fn norm_squared(&self) -> f64 {
4765 let norm = self.norm();
4766 norm * norm
4767 }
4768
4769 fn log_norm(&self) -> f64 {
4770 if self.infinite {
4771 f64::INFINITY
4772 } else if self.scale == 0.0 {
4773 f64::NEG_INFINITY
4774 } else {
4775 self.scale.ln() + 0.5 * self.sumsq.ln()
4776 }
4777 }
4778}
4779
4780impl IdxTensor {
4785 pub fn norm_squared(&self) -> std::result::Result<f64, IdxTensorError> {
4811 let dtype = self
4812 .scalar_dtype()
4813 .map_err(IdxTensorError::scalar_extraction)?;
4814 if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4815 return Err(IdxTensorError::ScalarTypeMismatch {
4816 expected: "f32, f64, c32, or c64",
4817 actual: Self::dtype_name(dtype).to_string(),
4818 });
4819 }
4820 let (has_nan, _) = self
4821 .compact_nonfinite_flags()
4822 .map_err(IdxTensorError::materialization)?;
4823 if has_nan {
4824 return Err(IdxTensorError::NaNInput {
4825 operation: "norm_squared",
4826 });
4827 }
4828
4829 let mut norm = Lassq::default();
4830 self.storage
4831 .for_each_payload_value(|value| norm.add_complex(value))
4832 .map_err(IdxTensorError::materialization)?;
4833 let value = norm.norm_squared();
4834 if value.is_nan() {
4835 return Err(IdxTensorError::NaNInput {
4836 operation: "norm_squared",
4837 });
4838 }
4839 Ok(value)
4840 }
4841
4842 pub fn norm(&self) -> std::result::Result<f64, IdxTensorError> {
4860 Ok(self.norm_squared()?.sqrt())
4861 }
4862
4863 pub fn maxabs(&self) -> std::result::Result<f64, IdxTensorError> {
4879 if let Some(error) = self.storage.deferred_error() {
4880 return Err(IdxTensorError::Storage {
4881 source: error.clone(),
4882 });
4883 }
4884 let dtype = self
4885 .storage
4886 .dtype()
4887 .ok_or_else(|| IdxTensorError::ScalarTypeMismatch {
4888 expected: "f32, f64, c32, or c64",
4889 actual: "unknown".to_string(),
4890 })?;
4891 if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4892 return Err(IdxTensorError::ScalarTypeMismatch {
4893 expected: "f32, f64, c32, or c64",
4894 actual: Self::dtype_name(dtype).to_string(),
4895 });
4896 }
4897 let (has_nan, _) = self
4898 .compact_nonfinite_flags()
4899 .map_err(IdxTensorError::materialization)?;
4900 if has_nan {
4901 return Err(IdxTensorError::NaNInput {
4902 operation: "maxabs",
4903 });
4904 }
4905 let mut value = 0.0_f64;
4906 self.storage
4907 .for_each_payload_value(|scalar| {
4908 let magnitude = scalar.re.hypot(scalar.im);
4909 value = value.max(magnitude);
4910 })
4911 .map_err(IdxTensorError::materialization)?;
4912 Ok(value)
4913 }
4914
4915 fn native_complex_payload_value_at(
4916 native: &EagerTensor,
4917 payload_coords: &[usize],
4918 ) -> Result<Complex64> {
4919 if !(payload_coords.len() == native.shape().len()) {
4920 return Err(anyhow::anyhow!(
4921 "payload coordinate rank {} does not match payload rank {}",
4922 payload_coords.len(),
4923 native.shape().len()
4924 ));
4925 };
4926 for (&coordinate, &dim) in payload_coords.iter().zip(native.shape().iter()) {
4927 if coordinate >= dim {
4928 return Err(anyhow::anyhow!(
4929 "payload coordinate {coordinate} is out of bounds for dim {dim}"
4930 ));
4931 }
4932 }
4933 let value = native.value()?;
4934 match value.as_tensor_view() {
4935 TensorView::F32(view) => view
4936 .get(payload_coords)
4937 .copied()
4938 .map(|value| Complex64::new(f64::from(value), 0.0))
4939 .ok_or_else(|| anyhow::anyhow!("failed to read f32 payload value")),
4940 TensorView::F64(view) => view
4941 .get(payload_coords)
4942 .copied()
4943 .map(|value| Complex64::new(value, 0.0))
4944 .ok_or_else(|| anyhow::anyhow!("failed to read f64 payload value")),
4945 TensorView::C32(view) => view
4946 .get(payload_coords)
4947 .copied()
4948 .map(|value| Complex64::new(f64::from(value.re), f64::from(value.im)))
4949 .ok_or_else(|| anyhow::anyhow!("failed to read c32 payload value")),
4950 TensorView::C64(view) => view
4951 .get(payload_coords)
4952 .copied()
4953 .ok_or_else(|| anyhow::anyhow!("failed to read c64 payload value")),
4954 view => Err(anyhow::anyhow!(
4955 "unsupported payload dtype {:?}",
4956 view.dtype()
4957 )),
4958 }
4959 }
4960
4961 fn native_sum_scalar(native: &EagerTensor) -> Result<AnyScalar> {
4962 let value = native.value()?;
4963 match native.dtype() {
4964 DType::F32 => Ok(AnyScalar::from_value(
4965 value.as_slice::<f32>()?.iter().copied().sum::<f32>(),
4966 )),
4967 DType::F64 => Ok(AnyScalar::from_value(
4968 value.as_slice::<f64>()?.iter().copied().sum::<f64>(),
4969 )),
4970 DType::C32 => Ok(AnyScalar::from_value(
4971 value
4972 .as_slice::<Complex32>()?
4973 .iter()
4974 .copied()
4975 .sum::<Complex32>(),
4976 )),
4977 DType::C64 => Ok(AnyScalar::from_value(
4978 value
4979 .as_slice::<Complex64>()?
4980 .iter()
4981 .copied()
4982 .sum::<Complex64>(),
4983 )),
4984 dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
4985 }
4986 }
4987
4988 fn native_nonfinite_flags_typed<T: TensorElement>(
4989 native: &EagerTensor,
4990 classify: impl Fn(T) -> (bool, bool),
4991 ) -> Result<(bool, bool)> {
4992 let value = native.value()?;
4993 if let Ok(values) = value.as_slice::<T>() {
4994 return Ok(values.iter().copied().map(&classify).fold(
4995 (false, false),
4996 |(has_nan, has_infinity), (is_nan, is_infinite)| {
4997 (has_nan | is_nan, has_infinity | is_infinite)
4998 },
4999 ));
5000 }
5001 drop(value);
5002 let value = IdxTensorStorage::materialize_eager_payload(native)?;
5003 Ok(value.as_slice::<T>()?.iter().copied().map(classify).fold(
5004 (false, false),
5005 |(has_nan, has_infinity), (is_nan, is_infinite)| {
5006 (has_nan | is_nan, has_infinity | is_infinite)
5007 },
5008 ))
5009 }
5010
5011 fn native_nonfinite_flags(native: &EagerTensor) -> Result<(bool, bool)> {
5012 match native.dtype() {
5013 DType::F32 => Self::native_nonfinite_flags_typed(native, |value: f32| {
5014 (value.is_nan(), value.is_infinite())
5015 }),
5016 DType::F64 => Self::native_nonfinite_flags_typed(native, |value: f64| {
5017 (value.is_nan(), value.is_infinite())
5018 }),
5019 DType::C32 => Self::native_nonfinite_flags_typed(native, |value: Complex32| {
5020 (
5021 value.re.is_nan() || value.im.is_nan(),
5022 value.re.is_infinite() || value.im.is_infinite(),
5023 )
5024 }),
5025 DType::C64 => Self::native_nonfinite_flags_typed(native, |value: Complex64| {
5026 (
5027 value.re.is_nan() || value.im.is_nan(),
5028 value.re.is_infinite() || value.im.is_infinite(),
5029 )
5030 }),
5031 dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
5032 }
5033 }
5034
5035 fn compact_nonfinite_flags(&self) -> Result<(bool, bool)> {
5036 self.storage.nonfinite_flags()
5037 }
5038
5039 pub fn sub(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
5049 self.axpby(AnyScalar::new_real(1.0), other, AnyScalar::new_real(-1.0))
5050 }
5051
5052 pub fn neg(&self) -> std::result::Result<Self, IdxTensorError> {
5059 self.scale(AnyScalar::new_real(-1.0))
5060 }
5061
5062 pub fn isapprox(
5074 &self,
5075 other: &Self,
5076 atol: f64,
5077 rtol: f64,
5078 ) -> std::result::Result<bool, IdxTensorError> {
5079 for (name, value) in [("atol", atol), ("rtol", rtol)] {
5080 if !value.is_finite() || value < 0.0 {
5081 return Err(IdxTensorError::InvalidTolerance { name, value });
5082 }
5083 }
5084 if self.indices.len() != other.indices.len() {
5085 return Err(IdxTensorError::ShapeMismatch {
5086 operation: "isapprox",
5087 expected: format!("indices {:?}", self.indices),
5088 actual: format!("indices {:?}", other.indices),
5089 });
5090 }
5091
5092 let other_axis_by_index = other
5093 .indices
5094 .iter()
5095 .cloned()
5096 .enumerate()
5097 .map(|(axis, index)| (index, axis))
5098 .collect::<HashMap<_, _>>();
5099 let other_positions = self
5100 .indices
5101 .iter()
5102 .map(|index| {
5103 other_axis_by_index.get(index).copied().ok_or_else(|| {
5104 IdxTensorError::ShapeMismatch {
5105 operation: "isapprox",
5106 expected: format!("indices {:?}", self.indices),
5107 actual: format!("indices {:?}", other.indices),
5108 }
5109 })
5110 })
5111 .collect::<std::result::Result<Vec<_>, _>>()?;
5112 let self_dims = self.dims();
5113 let other_dims = other.dims();
5114 for (axis, &other_axis) in other_positions.iter().enumerate() {
5115 if self_dims[axis] != other_dims[other_axis] {
5116 return Err(IdxTensorError::ShapeMismatch {
5117 operation: "isapprox",
5118 expected: format!("dims {:?}", self_dims),
5119 actual: format!("dims {:?}", other_dims),
5120 });
5121 }
5122 }
5123 for tensor in [self, other] {
5124 if tensor
5125 .compact_nonfinite_flags()
5126 .map_err(IdxTensorError::materialization)?
5127 .0
5128 {
5129 return Err(IdxTensorError::NaNInput {
5130 operation: "isapprox",
5131 });
5132 }
5133 }
5134
5135 let exact = atol == 0.0 && rtol == 0.0;
5136 let lhs_payload_dims = self.storage.payload_dims().to_vec();
5137 let rhs_payload_dims = other.storage.payload_dims().to_vec();
5138 let lhs_payload_len =
5139 checked_product(&lhs_payload_dims).map_err(IdxTensorError::materialization)?;
5140 let rhs_payload_len =
5141 checked_product(&rhs_payload_dims).map_err(IdxTensorError::materialization)?;
5142 let lhs_axis_classes = self.storage.axis_classes();
5143 let rhs_axis_classes = other.storage.axis_classes();
5144 let self_to_other = other_positions.clone();
5145 let mut other_to_self = vec![0usize; self_to_other.len()];
5146 for (self_axis, &other_axis) in self_to_other.iter().enumerate() {
5147 other_to_self[other_axis] = self_axis;
5148 }
5149
5150 let mut diff = Lassq::default();
5151 let mut lhs_norm = Lassq::default();
5152 let mut rhs_norm = Lassq::default();
5153 let mut compare = |lhs: Complex64, rhs: Complex64| -> bool {
5154 if exact {
5155 return lhs == rhs;
5156 }
5157 let lhs_infinite = lhs.re.is_infinite() || lhs.im.is_infinite();
5158 let rhs_infinite = rhs.re.is_infinite() || rhs.im.is_infinite();
5159 if lhs_infinite || rhs_infinite {
5160 return lhs == rhs;
5161 }
5162 lhs_norm.add_complex(lhs);
5163 rhs_norm.add_complex(rhs);
5164 diff.add_complex_difference(lhs, rhs);
5165 true
5166 };
5167
5168 let mut lhs_coords = vec![0usize; lhs_payload_dims.len()];
5172 let mut rhs_from_lhs = vec![0usize; rhs_payload_dims.len()];
5173 let mut rhs_seen = vec![false; rhs_payload_dims.len()];
5174 for _ in 0..lhs_payload_len {
5175 let lhs = self
5176 .storage
5177 .payload_value_at(&lhs_coords)
5178 .map_err(IdxTensorError::materialization)?;
5179 if map_payload_support_coordinate(
5180 lhs_axis_classes,
5181 rhs_axis_classes,
5182 &self_to_other,
5183 &lhs_coords,
5184 &mut rhs_from_lhs,
5185 &mut rhs_seen,
5186 ) {
5187 let rhs = other
5188 .storage
5189 .payload_value_at(&rhs_from_lhs)
5190 .map_err(IdxTensorError::materialization)?;
5191 if !compare(lhs, rhs) {
5192 return Ok(false);
5193 }
5194 } else if !compare(lhs, Complex64::new(0.0, 0.0)) {
5195 return Ok(false);
5196 }
5197
5198 increment_col_major_coordinate(&mut lhs_coords, &lhs_payload_dims);
5199 }
5200
5201 let mut rhs_coords = vec![0usize; rhs_payload_dims.len()];
5205 let mut lhs_from_rhs = vec![0usize; lhs_payload_dims.len()];
5206 let mut lhs_seen = vec![false; lhs_payload_dims.len()];
5207 for _ in 0..rhs_payload_len {
5208 let rhs = other
5209 .storage
5210 .payload_value_at(&rhs_coords)
5211 .map_err(IdxTensorError::materialization)?;
5212 if !map_payload_support_coordinate(
5213 rhs_axis_classes,
5214 lhs_axis_classes,
5215 &other_to_self,
5216 &rhs_coords,
5217 &mut lhs_from_rhs,
5218 &mut lhs_seen,
5219 ) && !compare(Complex64::new(0.0, 0.0), rhs)
5220 {
5221 return Ok(false);
5222 }
5223 increment_col_major_coordinate(&mut rhs_coords, &rhs_payload_dims);
5224 }
5225
5226 if exact {
5227 return Ok(true);
5228 }
5229 let absolute_ok = if diff.is_zero() {
5230 true
5231 } else if atol == 0.0 || diff.infinite {
5232 false
5233 } else {
5234 diff.log_norm() <= atol.ln()
5235 };
5236 let relative_ok = if rtol == 0.0 || diff.is_zero() {
5237 diff.is_zero()
5238 } else if diff.infinite {
5239 lhs_norm.infinite || rhs_norm.infinite
5240 } else if lhs_norm.infinite || rhs_norm.infinite {
5241 true
5242 } else {
5243 let reference_log = lhs_norm.log_norm().max(rhs_norm.log_norm());
5244 diff.log_norm() <= rtol.ln() + reference_log
5245 };
5246 Ok(absolute_ok || relative_ok)
5247 }
5248
5249 pub fn diagonal(
5256 input_index: &DynIndex,
5257 output_index: &DynIndex,
5258 ) -> std::result::Result<Self, IdxTensorError> {
5259 <Self as TensorConstructionLike>::diagonal(input_index, output_index)
5260 }
5261
5262 pub fn delta(
5268 input_indices: &[DynIndex],
5269 output_indices: &[DynIndex],
5270 ) -> std::result::Result<Self, IdxTensorError> {
5271 <Self as TensorConstructionLike>::delta(input_indices, output_indices)
5272 }
5273
5274 pub fn scalar_one() -> std::result::Result<Self, IdxTensorError> {
5281 <Self as TensorConstructionLike>::scalar_one()
5282 }
5283
5284 pub fn ones(indices: &[DynIndex]) -> std::result::Result<Self, IdxTensorError> {
5291 <Self as TensorConstructionLike>::ones(indices)
5292 }
5293
5294 pub fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, IdxTensorError> {
5301 <Self as TensorConstructionLike>::onehot(index_vals)
5302 }
5303
5304 pub fn mask_index(
5343 &self,
5344 index: &DynIndex,
5345 position: usize,
5346 ) -> std::result::Result<Self, IdxTensorError> {
5347 if !(self.indices.iter().any(|candidate| candidate == index)) {
5348 return Err(anyhow::anyhow!("mask_index: index is not present in tensor").into());
5349 };
5350 if !(position < index.dim()) {
5351 return Err(anyhow::anyhow!(
5352 "mask_index: position {position} is out of range for dimension {}",
5353 index.dim()
5354 )
5355 .into());
5356 };
5357
5358 let mask = match self.scalar_dtype()? {
5363 DType::F32 => Self::from_dense(
5364 vec![index.clone()],
5365 (0..index.dim())
5366 .map(|value| if value == position { 1.0_f32 } else { 0.0 })
5367 .collect(),
5368 ),
5369 DType::F64 => Self::from_dense(
5370 vec![index.clone()],
5371 (0..index.dim())
5372 .map(|value| if value == position { 1.0_f64 } else { 0.0 })
5373 .collect(),
5374 ),
5375 DType::C32 => Self::from_dense(
5376 vec![index.clone()],
5377 (0..index.dim())
5378 .map(|value| {
5379 if value == position {
5380 num_complex::Complex32::new(1.0, 0.0)
5381 } else {
5382 num_complex::Complex32::new(0.0, 0.0)
5383 }
5384 })
5385 .collect(),
5386 ),
5387 DType::C64 => Self::from_dense(
5388 vec![index.clone()],
5389 (0..index.dim())
5390 .map(|value| {
5391 if value == position {
5392 Complex64::new(1.0, 0.0)
5393 } else {
5394 Complex64::new(0.0, 0.0)
5395 }
5396 })
5397 .collect(),
5398 ),
5399 dtype => {
5400 return Err(anyhow::anyhow!("mask_index does not support dtype {dtype:?}").into())
5401 }
5402 }?;
5403 super::contract::contract_pair_with_options(
5404 self,
5405 &mask,
5406 super::contract::ContractionOptions::new()
5407 .with_retain_indices(std::slice::from_ref(index)),
5408 )
5409 }
5410
5411 pub fn distance(&self, other: &Self) -> std::result::Result<f64, IdxTensorError> {
5446 let norm_self = self.norm()?;
5447
5448 let neg_other = other.scale(AnyScalar::new_real(-1.0))?;
5450 let diff = self.add(&neg_other)?;
5451 let norm_diff = diff.norm()?;
5452
5453 if norm_self > 0.0 {
5454 Ok(norm_diff / norm_self)
5455 } else {
5456 Ok(norm_diff)
5457 }
5458 }
5459}
5460
5461impl std::fmt::Debug for IdxTensor {
5462 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5463 f.debug_struct("IdxTensor")
5464 .field("indices", &self.indices)
5465 .field("dims", &self.dims())
5466 .field("is_diag", &self.is_diag())
5467 .finish()
5468 }
5469}
5470
5471pub fn diag_idx_tensor(
5495 indices: Vec<DynIndex>,
5496 diag_data: Vec<f64>,
5497) -> std::result::Result<IdxTensor, IdxTensorError> {
5498 IdxTensor::from_diag(indices, diag_data)
5499}
5500
5501#[allow(clippy::type_complexity)]
5502pub(crate) type UnfoldSplitInnerResult = (
5503 EagerTensor,
5504 usize,
5505 usize,
5506 usize,
5507 Vec<DynIndex>,
5508 Vec<DynIndex>,
5509);
5510
5511#[allow(clippy::type_complexity)]
5551pub fn unfold_split(
5552 t: &IdxTensor,
5553 left_inds: &[DynIndex],
5554) -> std::result::Result<
5555 (
5556 NativeTensor,
5557 usize,
5558 usize,
5559 usize,
5560 Vec<DynIndex>,
5561 Vec<DynIndex>,
5562 ),
5563 IdxTensorError,
5564> {
5565 let (matrix_inner, left_len, m, n, left_indices, right_indices) =
5566 unfold_split_inner(t, left_inds)?;
5567
5568 Ok((
5569 matrix_inner.duplicate_value()?,
5570 left_len,
5571 m,
5572 n,
5573 left_indices,
5574 right_indices,
5575 ))
5576}
5577
5578pub(crate) fn unfold_split_inner(
5579 t: &IdxTensor,
5580 left_inds: &[DynIndex],
5581) -> Result<UnfoldSplitInnerResult> {
5582 let rank = t.indices.len();
5583
5584 if !(rank >= 2) {
5586 return Err(anyhow::anyhow!(
5587 "Tensor must have rank >= 2, got rank {}",
5588 rank
5589 ));
5590 };
5591
5592 let left_len = left_inds.len();
5593
5594 if !(left_len > 0 && left_len < rank) {
5596 return Err(anyhow::anyhow!("Left indices must be a non-empty proper subset of tensor indices (0 < left_len < rank), got left_len={}, rank={}",
5597 left_len,
5598 rank));
5599 };
5600
5601 let tensor_set: HashSet<_> = t.indices.iter().collect();
5603 let mut left_set = HashSet::new();
5604
5605 for left_idx in left_inds {
5606 if !(tensor_set.contains(left_idx)) {
5607 return Err(anyhow::anyhow!("Index in left_inds not found in tensor"));
5608 };
5609 if !(left_set.insert(left_idx)) {
5610 return Err(anyhow::anyhow!("Duplicate index in left_inds"));
5611 };
5612 }
5613
5614 let mut right_inds = Vec::new();
5616 for idx in &t.indices {
5617 if !left_set.contains(idx) {
5618 right_inds.push(idx.clone());
5619 }
5620 }
5621
5622 let mut new_indices = Vec::with_capacity(rank);
5624 new_indices.extend_from_slice(left_inds);
5625 new_indices.extend_from_slice(&right_inds);
5626
5627 let unfolded = t.permute_indices(&new_indices)?;
5629
5630 let unfolded_dims = unfolded.dims();
5632 let m = checked_product(&unfolded_dims[..left_len])?;
5633 let n = checked_product(&unfolded_dims[left_len..])?;
5634
5635 let matrix_tensor = unfolded.try_materialized_inner()?.reshape(&[m, n])?;
5636
5637 Ok((
5638 matrix_tensor,
5639 left_len,
5640 m,
5641 n,
5642 left_inds.to_vec(),
5643 right_inds,
5644 ))
5645}
5646
5647use crate::tensor_index::TensorIndex;
5652
5653impl TensorIndex for IdxTensor {
5654 type Index = DynIndex;
5655 type Error = IdxTensorError;
5656
5657 fn external_indices(&self) -> Vec<DynIndex> {
5658 self.indices.clone()
5660 }
5661
5662 fn num_external_indices(&self) -> usize {
5663 self.indices.len()
5664 }
5665
5666 fn replaceind(
5667 &self,
5668 old_index: &DynIndex,
5669 new_index: &DynIndex,
5670 ) -> std::result::Result<Self, Self::Error> {
5671 IdxTensor::replaceind(self, old_index, new_index)
5673 }
5674
5675 fn replace_indices(
5676 &self,
5677 old_indices: &[DynIndex],
5678 new_indices: &[DynIndex],
5679 ) -> std::result::Result<Self, Self::Error> {
5680 IdxTensor::replace_indices(self, old_indices, new_indices)
5682 }
5683}
5684
5685use crate::tensor_like::{
5690 FactorizeError, FactorizeOptions, FactorizeResult, IncrementalQrState, TensorConstructionLike,
5691 TensorContractionLike, TensorFactorizationLike, TensorVectorSpace,
5692};
5693
5694impl TensorVectorSpace for IdxTensor {
5695 fn norm_squared(&self) -> std::result::Result<f64, Self::Error> {
5696 IdxTensor::norm_squared(self)
5697 }
5698
5699 fn maxabs(&self) -> std::result::Result<f64, Self::Error> {
5700 IdxTensor::maxabs(self)
5701 }
5702
5703 fn isapprox(
5704 &self,
5705 other: &Self,
5706 atol: f64,
5707 rtol: f64,
5708 ) -> std::result::Result<bool, Self::Error> {
5709 IdxTensor::isapprox(self, other, atol, rtol)
5710 }
5711
5712 fn axpby(
5713 &self,
5714 a: crate::AnyScalar,
5715 other: &Self,
5716 b: crate::AnyScalar,
5717 ) -> std::result::Result<Self, Self::Error> {
5718 IdxTensor::axpby(self, a, other, b)
5719 }
5720
5721 fn scale(&self, scalar: crate::AnyScalar) -> std::result::Result<Self, Self::Error> {
5722 IdxTensor::scale(self, scalar)
5723 }
5724
5725 fn scale_in(
5726 &self,
5727 factor: f64,
5728 context: &ExecutionContext,
5729 ) -> std::result::Result<Self, Self::Error> {
5730 IdxTensor::scale_in(self, factor, context)
5731 }
5732
5733 fn norm_in(&self, context: &ExecutionContext) -> std::result::Result<f64, Self::Error> {
5734 IdxTensor::norm_in(self, context)
5735 }
5736
5737 fn inner_product(&self, other: &Self) -> std::result::Result<crate::AnyScalar, Self::Error> {
5738 IdxTensor::inner_product(self, other)
5739 }
5740}
5741
5742impl TensorFactorizationLike for IdxTensor {
5743 fn factorize(
5744 &self,
5745 left_inds: &[DynIndex],
5746 options: &FactorizeOptions,
5747 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5748 crate::factorize::factorize(self, left_inds, options)
5749 }
5750
5751 fn factorize_in(
5752 &self,
5753 left_inds: &[DynIndex],
5754 options: &FactorizeOptions,
5755 context: &ExecutionContext,
5756 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5757 crate::factorize::factorize_in(self, left_inds, options, context)
5758 }
5759
5760 fn factorize_auto(
5761 &self,
5762 left_inds: &[DynIndex],
5763 options: &FactorizeOptions,
5764 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5765 crate::factorize::factorize_auto(self, left_inds, options)
5766 }
5767
5768 fn factorize_full_rank(
5769 &self,
5770 left_inds: &[DynIndex],
5771 alg: crate::FactorizeAlg,
5772 canonical: crate::Canonical,
5773 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5774 crate::factorize::factorize_full_rank(self, left_inds, alg, canonical)
5775 }
5776
5777 fn factorize_full_rank_in(
5778 &self,
5779 left_inds: &[DynIndex],
5780 alg: crate::FactorizeAlg,
5781 canonical: crate::Canonical,
5782 context: &ExecutionContext,
5783 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5784 crate::factorize::factorize_full_rank_in(self, left_inds, alg, canonical, context)
5785 }
5786
5787 fn factorize_probe_columns_incremental(
5788 previous: Option<&FactorizeResult<Self>>,
5789 all_columns: &[&Self],
5790 appended_columns: &[&Self],
5791 left_inds: &[DynIndex],
5792 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5793 factorize_probe_columns_incremental(previous, all_columns, appended_columns, left_inds)
5794 }
5795
5796 fn factorize_probe_batch_incremental(
5797 previous: Option<&FactorizeResult<IdxTensor>>,
5798 batch_tensor: &IdxTensor,
5799 batch_index: &DynIndex,
5800 left_inds: &[DynIndex],
5801 ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5802 factorize_probe_batch_incremental_impl(previous, batch_tensor, batch_index, left_inds)
5803 }
5804
5805 fn factorize_probe_batch_incremental_in(
5806 previous: Option<&FactorizeResult<IdxTensor>>,
5807 batch_tensor: &IdxTensor,
5808 batch_index: &DynIndex,
5809 left_inds: &[DynIndex],
5810 context: &ExecutionContext,
5811 ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5812 if let Some(previous) = previous {
5813 previous
5814 .left
5815 .validate_context(context)
5816 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5817 previous
5818 .right
5819 .validate_context(context)
5820 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5821 }
5822 batch_tensor
5823 .validate_context(context)
5824 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5825 if context.is_global_default_cpu() {
5830 return factorize_probe_batch_incremental_impl(
5831 previous,
5832 batch_tensor,
5833 batch_index,
5834 left_inds,
5835 );
5836 }
5837 Self::resident_probe_batch_qr(previous, batch_tensor, batch_index, left_inds, context)
5838 }
5839
5840 fn src_error_estimate(
5841 &self,
5842 ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5843 if self.is_cuda_resident() {
5844 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5845 "CUDA-resident QR factor requires src_error_estimate_in with the owning execution context"
5846 )));
5847 }
5848 let indices = self.indices();
5849 if indices.len() != 2 {
5850 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5851 "SRC QR factor must have rank 2, got {}",
5852 indices.len()
5853 )));
5854 }
5855 let nrows = indices[0].dim();
5856 let ncols = indices[1].dim();
5857 let result = if self.is_f64() {
5858 let matrix = Matrix::from_col_major_vec(
5859 nrows,
5860 ncols,
5861 self.to_vec::<f64>()
5862 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5863 );
5864 backend_src_error_estimate(&matrix)
5865 } else if self.is_c64() {
5866 let matrix = Matrix::from_col_major_vec(
5867 nrows,
5868 ncols,
5869 self.to_vec::<Complex64>()
5870 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5871 );
5872 backend_src_error_estimate(&matrix)
5873 } else {
5874 return Err(FactorizeError::UnsupportedStorage(
5875 "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
5876 ));
5877 };
5878 result.map_err(|error| FactorizeError::ComputationError(anyhow::anyhow!(error)))
5879 }
5880
5881 fn src_error_estimate_in(
5882 &self,
5883 context: &ExecutionContext,
5884 ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5885 self.validate_context(context)
5886 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5887 #[cfg(feature = "tenferro-cuda")]
5888 if matches!(context, ExecutionContext::Cuda(_)) {
5889 return self.resident_src_error_estimate(context);
5890 }
5891 if context.is_global_default_cpu() {
5892 return TensorFactorizationLike::src_error_estimate(self);
5893 }
5894 self.host_src_error_estimate_general()
5895 }
5896}
5897
5898impl IdxTensor {
5899 fn host_src_error_estimate_general(
5901 &self,
5902 ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5903 let indices = self.indices();
5904 if indices.len() != 2 {
5905 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5906 "SRC QR factor must have rank 2, got {}",
5907 indices.len()
5908 )));
5909 }
5910 let nrows = indices[0].dim();
5911 let ncols = indices[1].dim();
5912 let result = if self.is_f64() {
5913 let matrix = Matrix::from_col_major_vec(
5914 nrows,
5915 ncols,
5916 self.to_vec::<f64>()
5917 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5918 );
5919 backend_src_error_estimate_general(&matrix)
5920 } else if self.is_c64() {
5921 let matrix = Matrix::from_col_major_vec(
5922 nrows,
5923 ncols,
5924 self.to_vec::<Complex64>()
5925 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5926 );
5927 backend_src_error_estimate_general(&matrix)
5928 } else {
5929 return Err(FactorizeError::UnsupportedStorage(
5930 "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
5931 ));
5932 };
5933 result.map_err(|error| FactorizeError::ComputationError(anyhow::anyhow!(error)))
5934 }
5935
5936 fn resident_probe_batch_qr(
5947 previous: Option<&FactorizeResult<IdxTensor>>,
5948 batch_tensor: &IdxTensor,
5949 batch_index: &DynIndex,
5950 left_inds: &[DynIndex],
5951 context: &ExecutionContext,
5952 ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5953 use crate::defaults::idx_tensor::unfold_split_inner;
5954
5955 if !batch_tensor.indices().contains(batch_index) {
5959 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5960 "resident probe batch is missing its batch index"
5961 )));
5962 }
5963 let (appended_inner, _, m, b, _, _) = unfold_split_inner(batch_tensor, left_inds)
5964 .context("resident probe batch unfold failed")
5965 .map_err(FactorizeError::ComputationError)?;
5966 if !matches!(appended_inner.dtype(), DType::F64 | DType::C64) {
5967 return Err(FactorizeError::UnsupportedStorage(
5968 "incremental SRC factorization currently supports f64 and Complex64 tensors",
5969 ));
5970 }
5971 if b == 0 {
5972 if let Some(previous) = previous {
5973 return Ok(previous.clone());
5974 }
5975 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5976 "incremental SRC factorization requires at least one column"
5977 )));
5978 }
5979 let full_inner = if let Some(previous) = previous {
5980 let (q_inner, _, previous_rows, _, _, _) =
5981 unfold_split_inner(&previous.left, left_inds)
5982 .context("resident previous Q unfold failed")
5983 .map_err(FactorizeError::ComputationError)?;
5984 if previous_rows != m {
5985 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5986 "previous and appended sketch row dimensions differ: {previous_rows} vs {m}"
5987 )));
5988 }
5989 let (r_inner, _, _, _, _, _) =
5990 unfold_split_inner(&previous.right, std::slice::from_ref(&previous.bond_index))
5991 .context("resident previous R unfold failed")
5992 .map_err(FactorizeError::ComputationError)?;
5993 let reconstructed = q_inner.matmul(&r_inner).map_err(|error| {
5994 FactorizeError::ComputationError(
5995 anyhow::Error::new(error).context("resident sketch reconstruction failed"),
5996 )
5997 })?;
5998 EagerTensor::concatenate(&[&reconstructed, &appended_inner], 1).map_err(|error| {
5999 FactorizeError::ComputationError(
6000 anyhow::Error::new(error).context("resident sketch concatenation failed"),
6001 )
6002 })?
6003 } else {
6004 appended_inner
6005 };
6006 let total_width = *full_inner.shape().get(1).ok_or_else(|| {
6007 FactorizeError::ComputationError(anyhow::anyhow!(
6008 "resident sketch factor is not rank-2"
6009 ))
6010 })?;
6011 let rank_tolerance = 32.0 * f64::EPSILON * m.max(total_width) as f64;
6012 let decomposition = full_inner
6013 .rank_revealing_qr(RankRevealingQrOptions::default().rtol(rank_tolerance))
6014 .map_err(|error| {
6015 FactorizeError::ComputationError(
6016 anyhow::Error::new(error).context("resident probe batch RRQR failed"),
6017 )
6018 })?;
6019 let rank = Self::read_resident_rank(&decomposition.rank, context)?;
6020 let maximum_rank = m.min(total_width);
6021 if rank > maximum_rank {
6022 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6023 "resident RRQR returned rank {rank}, exceeding maximum {maximum_rank}"
6024 )));
6025 }
6026 if let Some(previous) = previous {
6027 if rank < previous.rank {
6028 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6029 "resident SRC RRQR rank decreased from {} to {rank}",
6030 previous.rank
6031 )));
6032 }
6033 }
6034 let q_full = decomposition
6035 .q
6036 .slice_axis(1, 0..rank)
6037 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6038 let q_adjoint = q_full
6042 .transpose(&[1, 0])
6043 .and_then(|transposed| transposed.conj())
6044 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6045 let r_full = q_adjoint.matmul(&full_inner).map_err(|error| {
6046 FactorizeError::ComputationError(
6047 anyhow::Error::new(error).context("resident RRQR right-factor restoration failed"),
6048 )
6049 })?;
6050 let cap = DynIndex::new_bond(rank)
6051 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6052 let batch = DynIndex::new_link(total_width)
6053 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6054 let mut q_indices = left_inds.to_vec();
6055 q_indices.push(cap.clone());
6056 let q_dims: Vec<usize> = q_indices.iter().map(|index| index.dim()).collect();
6057 let left = Self::from_inner(
6058 q_indices,
6059 q_full.reshape(&q_dims).map_err(|error| {
6060 FactorizeError::ComputationError(
6061 anyhow::Error::new(error).context("resident probe batch Q reshape failed"),
6062 )
6063 })?,
6064 )
6065 .map_err(FactorizeError::ComputationError)?;
6066 let right = Self::from_inner(
6067 vec![cap.clone(), batch],
6068 r_full.reshape(&[rank, total_width]).map_err(|error| {
6069 FactorizeError::ComputationError(
6070 anyhow::Error::new(error).context("resident probe batch R reshape failed"),
6071 )
6072 })?,
6073 )
6074 .map_err(FactorizeError::ComputationError)?;
6075 Ok(FactorizeResult::new(left, right, cap, None, rank))
6076 }
6077
6078 #[cfg(feature = "tenferro-cuda")]
6085 fn resident_src_error_estimate(
6086 &self,
6087 context: &ExecutionContext,
6088 ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
6089 use tensor4all_tensorbackend::SrcErrorEstimate;
6090
6091 let indices = self.indices();
6092 if indices.len() != 2 {
6093 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6094 "SRC QR factor must have rank 2, got {}",
6095 indices.len()
6096 )));
6097 }
6098 let nrows = indices[0].dim();
6099 let ncols = indices[1].dim();
6100 if nrows != ncols {
6101 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6102 "SRC estimator requires a square R, got {nrows}x{ncols}"
6103 )));
6104 }
6105 if nrows == 0 {
6106 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6107 "SRC estimator requires a non-empty R"
6108 )));
6109 }
6110 let inner = self.cuda_eager_inner().ok_or_else(|| {
6111 FactorizeError::ComputationError(anyhow::anyhow!(
6112 "SRC estimator requires a resident eager QR factor"
6113 ))
6114 })?;
6115 let ExecutionContext::Cuda(cuda) = context else {
6116 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6117 "SRC device estimator requires a CUDA execution context"
6118 )));
6119 };
6120 let adjoint = inner
6123 .transpose(&[1, 0])
6124 .and_then(|transposed| transposed.conj())
6125 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6126 let identity = Self::resident_identity(cuda, nrows, inner.dtype())?;
6127 let solved = adjoint.solve(&identity).map_err(|error| {
6129 FactorizeError::ComputationError(
6130 anyhow::Error::new(error).context("SRC inverse-adjoint general solve failed"),
6131 )
6132 })?;
6133 let column_sums = solved
6134 .abs()
6135 .and_then(|magnitudes| magnitudes.reduce_sum_squares(&[0]))
6136 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6137 let total = inner
6138 .abs()
6139 .and_then(|magnitudes| magnitudes.reduce_sum_squares(&[0, 1]))
6140 .and_then(|scalar| scalar.reshape(&[1]))
6141 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6142 let column_norms_sq = Self::read_resident_vector(&column_sums, ncols, context)?;
6143 let total_norm_sq = Self::read_resident_vector(&total, 1, context)?;
6144 Self::src_estimate_from_decision_scalars(&column_norms_sq, total_norm_sq[0], ncols)
6145 .map(|(error, norm)| SrcErrorEstimate { error, norm })
6146 .map_err(FactorizeError::ComputationError)
6147 }
6148
6149 fn read_resident_rank(
6151 rank: &EagerTensor,
6152 context: &ExecutionContext,
6153 ) -> std::result::Result<usize, FactorizeError> {
6154 if !rank.shape().is_empty() {
6155 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6156 "resident RRQR rank metadata has shape {:?}, expected []",
6157 rank.shape()
6158 )));
6159 }
6160 let resident = rank.to_tensor().map_err(|error| {
6161 FactorizeError::ComputationError(
6162 anyhow::Error::new(error).context("resident RRQR rank materialization failed"),
6163 )
6164 })?;
6165 #[cfg(feature = "tenferro-cuda")]
6166 let host = match context {
6167 ExecutionContext::Cpu(_) => resident,
6168 ExecutionContext::Cuda(cuda) => cuda.download(&resident).map_err(|error| {
6169 FactorizeError::ComputationError(
6170 anyhow::Error::new(error).context("resident RRQR rank download failed"),
6171 )
6172 })?,
6173 };
6174 #[cfg(not(feature = "tenferro-cuda"))]
6175 let host = match context {
6176 ExecutionContext::Cpu(_) => resident,
6177 };
6178 let values = host.as_slice::<i64>().map_err(|error| {
6179 FactorizeError::ComputationError(
6180 anyhow::Error::new(error).context("resident RRQR rank decode failed"),
6181 )
6182 })?;
6183 let value = *values.first().ok_or_else(|| {
6184 FactorizeError::ComputationError(anyhow::anyhow!(
6185 "resident RRQR rank metadata is empty"
6186 ))
6187 })?;
6188 usize::try_from(value).map_err(|error| {
6189 FactorizeError::ComputationError(
6190 anyhow::Error::new(error).context("resident RRQR rank is negative"),
6191 )
6192 })
6193 }
6194
6195 #[cfg(feature = "tenferro-cuda")]
6201 fn resident_identity(
6202 cuda: &Arc<tensor4all_tensorbackend::CudaExecutionContext>,
6203 n: usize,
6204 dtype: DType,
6205 ) -> std::result::Result<EagerTensor, FactorizeError> {
6206 let native = match dtype {
6207 DType::F64 => {
6208 let mut data = vec![0.0_f64; n * n];
6209 for diagonal in 0..n {
6210 data[diagonal + diagonal * n] = 1.0;
6211 }
6212 NativeTensor::from_vec_col_major(vec![n, n], data)
6213 }
6214 DType::C64 => {
6215 let mut data = vec![Complex64::new(0.0, 0.0); n * n];
6216 for diagonal in 0..n {
6217 data[diagonal + diagonal * n] = Complex64::new(1.0, 0.0);
6218 }
6219 NativeTensor::from_vec_col_major(vec![n, n], data)
6220 }
6221 _other => {
6222 return Err(FactorizeError::UnsupportedStorage(
6223 "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
6224 ));
6225 }
6226 }
6227 .map_err(|error| {
6228 FactorizeError::ComputationError(
6229 anyhow::Error::new(error).context("SRC estimator identity construction failed"),
6230 )
6231 })?;
6232 let uploaded = cuda.upload_cuda(&native).map_err(|error| {
6233 FactorizeError::ComputationError(
6234 anyhow::Error::new(error).context("SRC estimator identity upload failed"),
6235 )
6236 })?;
6237 let runtime = cuda.eager_runtime().map_err(|error| {
6238 FactorizeError::ComputationError(
6239 anyhow::Error::new(error).context("SRC estimator eager runtime failed"),
6240 )
6241 })?;
6242 EagerTensor::from_tensor_in(uploaded, runtime).map_err(|error| {
6245 FactorizeError::ComputationError(
6246 anyhow::Error::new(error).context("SRC estimator identity wrapping failed"),
6247 )
6248 })
6249 }
6250
6251 #[cfg(feature = "tenferro-cuda")]
6253 fn read_resident_vector(
6254 eager: &EagerTensor,
6255 len: usize,
6256 context: &ExecutionContext,
6257 ) -> std::result::Result<Vec<f64>, FactorizeError> {
6258 if eager.shape() != [len] {
6259 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6260 "device SRC decision vector has shape {:?}, expected [{len}]",
6261 eager.shape()
6262 )));
6263 }
6264 let decision = Self::from_inner(vec![DynIndex::new_dyn(len)], eager.clone())
6265 .map_err(FactorizeError::ComputationError)?;
6266 decision
6267 .read_decision_data(context)
6268 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))
6269 }
6270
6271 #[cfg(feature = "tenferro-cuda")]
6276 fn src_estimate_from_decision_scalars(
6277 column_norms_sq: &[f64],
6278 total_norm_sq: f64,
6279 ncols: usize,
6280 ) -> std::result::Result<(f64, f64), anyhow::Error> {
6281 if !total_norm_sq.is_finite() {
6282 return Err(anyhow::anyhow!(
6283 "SRC estimator requires finite entries in R"
6284 ));
6285 }
6286 let mut inverse_column_error_sq = 0.0_f64;
6287 for (col, &column_norm_sq) in column_norms_sq.iter().enumerate() {
6288 if !column_norm_sq.is_finite() || column_norm_sq == 0.0 {
6289 return Err(anyhow::anyhow!(
6290 "SRC inverse-adjoint solve returned an invalid column norm at column {col}"
6291 ));
6292 }
6293 inverse_column_error_sq += 1.0 / column_norm_sq;
6294 }
6295 let sketch_width = ncols as f64;
6296 let error_sq = inverse_column_error_sq / sketch_width;
6297 let norm_estimate_sq = total_norm_sq / sketch_width;
6298 if !error_sq.is_finite() || !norm_estimate_sq.is_finite() {
6299 return Err(anyhow::anyhow!(
6300 "SRC estimator produced a non-finite estimate"
6301 ));
6302 }
6303 Ok((error_sq.sqrt(), norm_estimate_sq.sqrt()))
6304 }
6305}
6306
6307fn factorize_probe_columns_incremental(
6312 previous: Option<&FactorizeResult<IdxTensor>>,
6313 all_columns: &[&IdxTensor],
6314 appended_columns: &[&IdxTensor],
6315 left_inds: &[DynIndex],
6316) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
6317 let first = all_columns.first().ok_or_else(|| {
6318 FactorizeError::ComputationError(anyhow::anyhow!(
6319 "incremental SRC factorization requires at least one column"
6320 ))
6321 })?;
6322 if first.is_f64() {
6323 incremental_probe_factorize_typed::<f64>(previous, all_columns, appended_columns, left_inds)
6324 } else if first.is_c64() {
6325 incremental_probe_factorize_typed::<Complex64>(
6326 previous,
6327 all_columns,
6328 appended_columns,
6329 left_inds,
6330 )
6331 } else {
6332 Err(FactorizeError::UnsupportedStorage(
6333 "incremental SRC factorization currently supports f64 and Complex64 tensors",
6334 ))
6335 }
6336}
6337
6338trait IncrementalQrStateScalar: IncrementalQrScalar + TensorElement {
6339 fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>>;
6340
6341 fn store(state: IncrementalQr<Self>) -> IncrementalQrState;
6342}
6343
6344impl IncrementalQrStateScalar for f64 {
6345 fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>> {
6346 match previous.incremental_qr_state()? {
6347 IncrementalQrState::F64(state) => Some(state.clone()),
6348 IncrementalQrState::C64(_) => None,
6349 }
6350 }
6351
6352 fn store(state: IncrementalQr<Self>) -> IncrementalQrState {
6353 IncrementalQrState::F64(state)
6354 }
6355}
6356
6357impl IncrementalQrStateScalar for Complex64 {
6358 fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>> {
6359 match previous.incremental_qr_state()? {
6360 IncrementalQrState::F64(_) => None,
6361 IncrementalQrState::C64(state) => Some(state.clone()),
6362 }
6363 }
6364
6365 fn store(state: IncrementalQr<Self>) -> IncrementalQrState {
6366 IncrementalQrState::C64(state)
6367 }
6368}
6369
6370fn incremental_probe_factorize_from_matrices<S>(
6371 previous: Option<&FactorizeResult<IdxTensor>>,
6372 appended_is_empty: bool,
6373 left_inds: &[DynIndex],
6374 initial_matrix: impl FnOnce() -> std::result::Result<Matrix<S>, FactorizeError>,
6375 appended_matrix: impl FnOnce() -> std::result::Result<Matrix<S>, FactorizeError>,
6376) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6377where
6378 S: IncrementalQrStateScalar,
6379{
6380 let (mut state, previous_left, previous_cap, previous_rank) = if let Some(previous) = previous {
6381 if appended_is_empty {
6382 return Ok(previous.clone());
6383 }
6384 let previous_rank = previous.rank;
6385 let resumed = S::resume(previous);
6386 let reused_previous_q = resumed.is_some();
6387 let state = if let Some(state) = resumed {
6388 state
6389 } else {
6390 let q = previous.left.clone();
6391 let cap = previous.bond_index.clone();
6392 let mut q_indices = left_inds.to_vec();
6393 q_indices.push(cap.clone());
6394 let q = q
6395 .permute_indices(&q_indices)
6396 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6397 let q_dims = q_indices.iter().map(IndexLike::dim).collect::<Vec<_>>();
6398 let q_rows = q_dims[..q_dims.len() - 1]
6399 .iter()
6400 .try_fold(1usize, |size, &dim| size.checked_mul(dim))
6401 .ok_or_else(|| {
6402 FactorizeError::ComputationError(anyhow::anyhow!(
6403 "incremental SRC Q row dimension overflows usize"
6404 ))
6405 })?;
6406 let q_matrix = Matrix::from_col_major_vec(
6407 q_rows,
6408 cap.dim(),
6409 q.to_vec::<S>()
6410 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6411 );
6412
6413 let batch = previous.right.indices().last().cloned().ok_or_else(|| {
6414 FactorizeError::ComputationError(anyhow::anyhow!(
6415 "incremental SRC R factor has no batch index"
6416 ))
6417 })?;
6418 let r = previous
6419 .right
6420 .permute_indices(&[cap.clone(), batch.clone()])
6421 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6422 let r_matrix = Matrix::from_col_major_vec(
6423 cap.dim(),
6424 batch.dim(),
6425 r.to_vec::<S>()
6426 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6427 );
6428 IncrementalQr::from_factors(q_matrix, r_matrix)
6429 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6430 };
6431 (
6432 state,
6433 if reused_previous_q {
6434 Some(previous.left.clone())
6435 } else {
6436 None
6437 },
6438 if reused_previous_q {
6439 Some(previous.bond_index.clone())
6440 } else {
6441 None
6442 },
6443 previous_rank,
6444 )
6445 } else {
6446 let initial_matrix = initial_matrix()?;
6447 (
6448 IncrementalQr::new(initial_matrix)
6449 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6450 None,
6451 None,
6452 0,
6453 )
6454 };
6455 if !appended_is_empty && previous.is_some() {
6456 let appended_matrix = appended_matrix()?;
6457 state
6458 .append(&appended_matrix)
6459 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6460 }
6461
6462 let rank = state.rank();
6463 if previous_left.is_some() && rank < previous_rank {
6464 return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6465 "incremental SRC QR rank decreased from {previous_rank} to {rank}"
6466 )));
6467 }
6468 let r = state.r();
6469 let cap = if rank == previous_rank {
6470 match previous_cap.as_ref() {
6471 Some(cap) => cap.clone(),
6472 None => DynIndex::new_bond(rank)
6473 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6474 }
6475 } else {
6476 DynIndex::new_bond(rank)
6477 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6478 };
6479 let left = if let Some(previous_left) = previous_left {
6480 if rank == previous_rank {
6481 previous_left
6482 } else {
6483 let appended_rank = rank - previous_rank;
6484 let appended_cap = DynIndex::new_bond(appended_rank)
6485 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6486 let appended_q = state
6487 .q_columns(previous_rank, appended_rank)
6488 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6489 let mut appended_indices = left_inds.to_vec();
6490 appended_indices.push(appended_cap.clone());
6491 let appended_left =
6492 IdxTensor::from_dense(appended_indices, appended_q.as_col_major_slice().to_vec())
6493 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6494 IdxTensor::concatenate_along_new_index(
6495 &[&previous_left, &appended_left],
6496 &[
6497 previous_cap.ok_or_else(|| {
6498 FactorizeError::ComputationError(anyhow::anyhow!(
6499 "incremental SRC previous QR state has no bond index"
6500 ))
6501 })?,
6502 appended_cap,
6503 ],
6504 cap.clone(),
6505 )
6506 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6507 }
6508 } else {
6509 let q = state.q();
6510 let mut q_indices = left_inds.to_vec();
6511 q_indices.push(cap.clone());
6512 IdxTensor::from_dense(q_indices, q.as_col_major_slice().to_vec())
6513 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6514 };
6515 let batch = DynIndex::new_link(r.ncols())
6516 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6517 let right = IdxTensor::from_dense(vec![cap.clone(), batch], r.as_col_major_slice().to_vec())
6518 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6519 Ok(FactorizeResult::new(left, right, cap, None, rank)
6520 .with_incremental_qr_state(S::store(state)))
6521}
6522
6523fn incremental_probe_factorize_typed<S>(
6524 previous: Option<&FactorizeResult<IdxTensor>>,
6525 all_columns: &[&IdxTensor],
6526 appended_columns: &[&IdxTensor],
6527 left_inds: &[DynIndex],
6528) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6529where
6530 S: IncrementalQrStateScalar,
6531{
6532 incremental_probe_factorize_from_matrices::<S>(
6533 previous,
6534 appended_columns.is_empty(),
6535 left_inds,
6536 || probe_columns_matrix::<S>(all_columns, left_inds),
6537 || probe_columns_matrix::<S>(appended_columns, left_inds),
6538 )
6539}
6540
6541fn incremental_probe_factorize_batch_typed<S>(
6542 previous: Option<&FactorizeResult<IdxTensor>>,
6543 batch_tensor: &IdxTensor,
6544 batch_index: &DynIndex,
6545 left_inds: &[DynIndex],
6546) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6547where
6548 S: IncrementalQrStateScalar,
6549{
6550 incremental_probe_factorize_from_matrices::<S>(
6551 previous,
6552 batch_index.dim() == 0,
6553 left_inds,
6554 || probe_batch_matrix::<S>(batch_tensor, batch_index, left_inds),
6555 || probe_batch_matrix::<S>(batch_tensor, batch_index, left_inds),
6556 )
6557}
6558
6559fn factorize_probe_batch_incremental_impl(
6560 previous: Option<&FactorizeResult<IdxTensor>>,
6561 batch_tensor: &IdxTensor,
6562 batch_index: &DynIndex,
6563 left_inds: &[DynIndex],
6564) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
6565 if batch_tensor.is_f64() {
6566 incremental_probe_factorize_batch_typed::<f64>(
6567 previous,
6568 batch_tensor,
6569 batch_index,
6570 left_inds,
6571 )
6572 } else if batch_tensor.is_c64() {
6573 incremental_probe_factorize_batch_typed::<Complex64>(
6574 previous,
6575 batch_tensor,
6576 batch_index,
6577 left_inds,
6578 )
6579 } else {
6580 Err(FactorizeError::UnsupportedStorage(
6581 "incremental SRC factorization currently supports f64 and Complex64 tensors",
6582 ))
6583 }
6584}
6585
6586fn probe_columns_matrix<S>(
6587 columns: &[&IdxTensor],
6588 left_inds: &[DynIndex],
6589) -> std::result::Result<Matrix<S>, FactorizeError>
6590where
6591 S: TensorElement,
6592{
6593 let batch = DynIndex::new_link(columns.len())
6594 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6595 let stacked = IdxTensor::stack_along_new_index(columns, batch.clone(), -1)
6596 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6597 let mut ordered_indices = left_inds.to_vec();
6598 ordered_indices.push(batch);
6599 let ordered = stacked
6600 .permute_indices(&ordered_indices)
6601 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6602 let nrows = left_inds
6603 .iter()
6604 .try_fold(1usize, |size, index| size.checked_mul(index.dim()))
6605 .ok_or_else(|| {
6606 FactorizeError::ComputationError(anyhow::anyhow!(
6607 "incremental SRC sketch row dimension overflows usize"
6608 ))
6609 })?;
6610 Ok(Matrix::from_col_major_vec(
6611 nrows,
6612 columns.len(),
6613 ordered
6614 .to_vec::<S>()
6615 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6616 ))
6617}
6618
6619fn probe_batch_matrix<S>(
6624 batch_tensor: &IdxTensor,
6625 batch_index: &DynIndex,
6626 left_inds: &[DynIndex],
6627) -> std::result::Result<Matrix<S>, FactorizeError>
6628where
6629 S: TensorElement,
6630{
6631 let mut ordered_indices = left_inds.to_vec();
6632 ordered_indices.push(batch_index.clone());
6633 let ordered = batch_tensor
6634 .permute_indices(&ordered_indices)
6635 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6636 let nrows = left_inds
6637 .iter()
6638 .try_fold(1usize, |size, index| size.checked_mul(index.dim()))
6639 .ok_or_else(|| {
6640 FactorizeError::ComputationError(anyhow::anyhow!(
6641 "incremental SRC sketch row dimension overflows usize"
6642 ))
6643 })?;
6644 Ok(Matrix::from_col_major_vec(
6645 nrows,
6646 batch_index.dim(),
6647 ordered
6648 .to_vec::<S>()
6649 .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6650 ))
6651}
6652
6653impl TensorContractionLike for IdxTensor {
6654 fn conj(&self) -> Self {
6655 IdxTensor::conj(self)
6657 }
6658
6659 fn direct_sum(
6660 &self,
6661 other: &Self,
6662 pairs: &[(DynIndex, DynIndex)],
6663 ) -> std::result::Result<crate::tensor_like::DirectSumResult<Self>, Self::Error> {
6664 let (tensor, new_indices) = crate::direct_sum::direct_sum(self, other, pairs)?;
6665 Ok(crate::tensor_like::DirectSumResult {
6666 tensor,
6667 new_indices,
6668 })
6669 }
6670
6671 fn outer_product(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
6672 super::contract::outer_product(self, other)
6673 }
6674
6675 fn permuteinds(&self, new_order: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
6676 IdxTensor::permute_indices(self, new_order)
6678 }
6679
6680 fn fuse_indices(
6681 &self,
6682 old_indices: &[DynIndex],
6683 new_index: DynIndex,
6684 order: LinearizationOrder,
6685 ) -> std::result::Result<Self, Self::Error> {
6686 IdxTensor::fuse_indices(self, old_indices, new_index, order)
6687 }
6688
6689 fn contract(tensors: &[&Self]) -> std::result::Result<Self, Self::Error> {
6690 super::contract::contract(tensors)
6691 }
6692
6693 fn contract_retaining_indices(
6694 tensors: &[&Self],
6695 retained_indices: &[DynIndex],
6696 ) -> std::result::Result<Self, Self::Error> {
6697 if tensors.len() == 2 {
6698 return tensors[0]
6699 .try_contract_pairwise_retaining(tensors[1], retained_indices)
6700 .map_err(IdxTensorError::from);
6701 }
6702 let options = ContractionOptions::new().with_retain_indices(retained_indices);
6703 super::contract::contract_with_options(tensors, options)
6704 }
6705
6706 fn contract_pair(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
6707 super::contract::contract_pair(self, other)
6708 }
6709}
6710
6711impl TensorConstructionLike for IdxTensor {
6712 fn select_indices(
6713 &self,
6714 selected_indices: &[DynIndex],
6715 positions: &[usize],
6716 ) -> std::result::Result<Self, Self::Error> {
6717 IdxTensor::select_indices(self, selected_indices, positions)
6718 }
6719
6720 fn diagonal(
6721 input_index: &DynIndex,
6722 output_index: &DynIndex,
6723 ) -> std::result::Result<Self, Self::Error> {
6724 let dim = input_index.dim();
6725 if dim != output_index.dim() {
6726 return Err(anyhow::anyhow!(
6727 "Dimension mismatch: input index has dim {}, output has dim {}",
6728 dim,
6729 output_index.dim(),
6730 )
6731 .into());
6732 }
6733
6734 IdxTensor::from_diag(
6735 vec![input_index.clone(), output_index.clone()],
6736 vec![1.0_f64; dim],
6737 )
6738 }
6739
6740 fn scalar_one() -> std::result::Result<Self, Self::Error> {
6741 IdxTensor::from_dense(vec![], vec![1.0_f64])
6742 }
6743
6744 fn ones(indices: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
6745 if indices.is_empty() {
6746 return <Self as TensorConstructionLike>::scalar_one();
6747 }
6748 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
6749 let total_size = checked_total_size(&dims)?;
6750 IdxTensor::from_dense(indices.to_vec(), vec![1.0_f64; total_size])
6751 }
6752
6753 fn ones_in(
6754 context: &tensor4all_tensorbackend::ExecutionContext,
6755 indices: &[DynIndex],
6756 ) -> std::result::Result<Self, Self::Error> {
6757 IdxTensor::ones_in(context, indices)
6758 }
6759
6760 fn validate_context(
6761 &self,
6762 context: &tensor4all_tensorbackend::ExecutionContext,
6763 ) -> std::result::Result<(), Self::Error> {
6764 IdxTensor::validate_context(self, context)
6765 }
6766
6767 fn from_dense_any(
6768 indices: Vec<DynIndex>,
6769 data: Vec<AnyScalar>,
6770 ) -> std::result::Result<Self, Self::Error> {
6771 IdxTensor::from_dense_any(indices, data)
6772 }
6773
6774 fn from_dense<T: TensorElement + Into<AnyScalar>>(
6775 indices: Vec<DynIndex>,
6776 data: Vec<T>,
6777 ) -> std::result::Result<Self, Self::Error> {
6778 IdxTensor::from_dense(indices, data)
6779 }
6780
6781 fn from_dense_in<T: TensorElement + Into<AnyScalar>>(
6782 context: &tensor4all_tensorbackend::ExecutionContext,
6783 indices: Vec<DynIndex>,
6784 data: Vec<T>,
6785 ) -> std::result::Result<Self, Self::Error> {
6786 IdxTensor::from_dense_in(context, indices, data)
6787 }
6788
6789 fn stack_along_new_index(
6790 tensors: &[&Self],
6791 new_index: DynIndex,
6792 axis: isize,
6793 ) -> std::result::Result<Self, Self::Error> {
6794 IdxTensor::stack_along_new_index(tensors, new_index, axis)
6795 }
6796
6797 fn concatenate_along_new_index(
6798 tensors: &[&Self],
6799 source_indices: &[DynIndex],
6800 new_index: DynIndex,
6801 ) -> std::result::Result<Self, Self::Error> {
6802 let first = tensors
6803 .first()
6804 .ok_or_else(|| anyhow::anyhow!("concatenate requires at least one tensor"))?;
6805 if tensors.len() != source_indices.len() {
6806 return Err(anyhow::anyhow!(
6807 "concatenate tensor count {} does not match source-index count {}",
6808 tensors.len(),
6809 source_indices.len()
6810 )
6811 .into());
6812 }
6813 let first_axis = first
6814 .indices
6815 .iter()
6816 .position(|index| index == &source_indices[0])
6817 .ok_or_else(|| anyhow::anyhow!("concatenate source index is not present"))?;
6818 let mut total_dim = 0usize;
6819 let first_indices = first.indices();
6820 let mut inners = Vec::with_capacity(tensors.len());
6821 for (tensor, source) in tensors.iter().zip(source_indices) {
6822 let axis = tensor
6823 .indices
6824 .iter()
6825 .position(|index| index == source)
6826 .ok_or_else(|| anyhow::anyhow!("concatenate source index is not present"))?;
6827 if axis != first_axis {
6828 return Err(
6829 anyhow::anyhow!("concatenate source axes must have the same position").into(),
6830 );
6831 }
6832 if tensor
6833 .indices()
6834 .iter()
6835 .zip(first_indices)
6836 .enumerate()
6837 .any(|(position, (actual, expected))| position != first_axis && actual != expected)
6838 {
6839 return Err(anyhow::anyhow!(
6840 "concatenate tensors must match away from the source axis"
6841 )
6842 .into());
6843 }
6844 total_dim = total_dim
6845 .checked_add(source.dim())
6846 .ok_or_else(|| anyhow::anyhow!("concatenate dimension overflow"))?;
6847 tensor.ensure_shape_packing_preserves_ad("concatenate_along_new_index")?;
6848 inners.push(tensor.try_materialized_inner()?);
6849 }
6850 if new_index.dim() != total_dim {
6851 return Err(anyhow::anyhow!(
6852 "concatenate output dimension {} does not match source dimension sum {}",
6853 new_index.dim(),
6854 total_dim
6855 )
6856 .into());
6857 }
6858 let concatenated = EagerTensor::concatenate(&inners, first_axis)?;
6859 let mut indices = first_indices.to_vec();
6860 indices[first_axis] = new_index;
6861 Self::from_inner(indices, concatenated).map_err(IdxTensorError::from)
6862 }
6863
6864 fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, Self::Error> {
6865 if index_vals.is_empty() {
6866 return <Self as TensorConstructionLike>::scalar_one();
6867 }
6868 let indices: Vec<DynIndex> = index_vals.iter().map(|(idx, _)| idx.clone()).collect();
6869 let vals: Vec<usize> = index_vals.iter().map(|(_, v)| *v).collect();
6870 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
6871
6872 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
6873 if v >= d {
6874 return Err(anyhow::anyhow!(
6875 "onehot: value {} at position {} is >= dimension {}",
6876 v,
6877 k,
6878 d
6879 )
6880 .into());
6881 }
6882 }
6883
6884 let total_size = checked_total_size(&dims).map_err(Self::Error::from)?;
6885 let mut data = vec![0.0_f64; total_size];
6886
6887 let offset = column_major_offset(&dims, &vals).map_err(Self::Error::from)?;
6888 data[offset] = 1.0;
6889
6890 Self::from_dense(indices, data)
6891 }
6892
6893 }
6895
6896fn checked_total_size(dims: &[usize]) -> Result<usize> {
6897 dims.iter().try_fold(1_usize, |acc, &d| {
6898 if d == 0 {
6899 return Err(anyhow::anyhow!("invalid dimension 0"));
6900 }
6901 acc.checked_mul(d)
6902 .ok_or_else(|| anyhow::anyhow!("tensor size overflow"))
6903 })
6904}
6905
6906fn column_major_offset(dims: &[usize], vals: &[usize]) -> Result<usize> {
6907 if dims.len() != vals.len() {
6908 return Err(anyhow::anyhow!(
6909 "column_major_offset: dims.len() != vals.len()"
6910 ));
6911 }
6912 checked_total_size(dims)?;
6913
6914 let mut offset = 0usize;
6915 let mut stride = 1usize;
6916 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
6917 if d == 0 {
6918 return Err(anyhow::anyhow!("invalid dimension 0 at position {}", k));
6919 }
6920 if v >= d {
6921 return Err(anyhow::anyhow!(
6922 "column_major_offset: value {} at position {} is >= dimension {}",
6923 v,
6924 k,
6925 d
6926 ));
6927 }
6928 let term = v
6929 .checked_mul(stride)
6930 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6931 offset = offset
6932 .checked_add(term)
6933 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6934 stride = stride
6935 .checked_mul(d)
6936 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6937 }
6938 Ok(offset)
6939}
6940
6941impl IdxTensor {
6946 fn any_scalar_payload_to_complex(data: Vec<AnyScalar>) -> Vec<Complex64> {
6947 data.into_iter()
6948 .map(|value| {
6949 value
6950 .as_c64()
6951 .unwrap_or_else(|| Complex64::new(value.real(), 0.0))
6952 })
6953 .collect()
6954 }
6955
6956 fn any_scalar_payload_to_real(data: Vec<AnyScalar>) -> Vec<f64> {
6957 data.into_iter().map(|value| value.real()).collect()
6958 }
6959
6960 fn validate_dense_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
6961 let expected_len = checked_total_size(dims)?;
6962 if !(data_len == expected_len) {
6963 return Err(anyhow::anyhow!(
6964 "dense payload length {} does not match dims {:?} (expected {})",
6965 data_len,
6966 dims,
6967 expected_len
6968 ));
6969 };
6970 Ok(())
6971 }
6972
6973 fn validate_diag_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
6974 if !(!dims.is_empty()) {
6975 return Err(anyhow::anyhow!(
6976 "diagonal tensor construction requires at least one index"
6977 ));
6978 };
6979 Self::validate_diag_dims(dims)?;
6980 if !(data_len == dims[0]) {
6981 return Err(anyhow::anyhow!(
6982 "diagonal payload length {} does not match diagonal dimension {}",
6983 data_len,
6984 dims[0]
6985 ));
6986 };
6987 Ok(())
6988 }
6989
6990 pub fn from_dense<T: TensorElement>(
7017 indices: Vec<DynIndex>,
7018 data: Vec<T>,
7019 ) -> std::result::Result<Self, IdxTensorError> {
7020 let dims = Self::expected_dims_from_indices(&indices);
7021 Self::validate_indices(&indices)?;
7022 Self::validate_dense_payload_len(data.len(), &dims)?;
7023 let native = dense_native_tensor_from_col_major_owned(data, &dims)?;
7024 Self::from_native(indices, native).map_err(IdxTensorError::from)
7025 }
7026
7027 pub fn from_dense_any(
7056 indices: Vec<DynIndex>,
7057 data: Vec<AnyScalar>,
7058 ) -> std::result::Result<Self, IdxTensorError> {
7059 if data.iter().any(AnyScalar::is_complex) {
7060 Self::from_dense(indices, Self::any_scalar_payload_to_complex(data))
7061 } else {
7062 Self::from_dense(indices, Self::any_scalar_payload_to_real(data))
7063 }
7064 }
7065
7066 pub fn from_diag<T: TensorElement>(
7098 indices: Vec<DynIndex>,
7099 data: Vec<T>,
7100 ) -> std::result::Result<Self, IdxTensorError> {
7101 let dims = Self::expected_dims_from_indices(&indices);
7102 Self::validate_indices(&indices)?;
7103 Self::validate_diag_payload_len(data.len(), &dims)?;
7104 let native = diag_native_tensor_from_col_major(&data, dims.len())?;
7105 Self::from_native_with_axis_classes(indices, native, Self::diag_axis_classes(dims.len()))
7106 .map_err(IdxTensorError::from)
7107 }
7108
7109 pub fn from_diag_any(
7135 indices: Vec<DynIndex>,
7136 data: Vec<AnyScalar>,
7137 ) -> std::result::Result<Self, IdxTensorError> {
7138 if data.iter().any(AnyScalar::is_complex) {
7139 Self::from_diag(indices, Self::any_scalar_payload_to_complex(data))
7140 } else {
7141 Self::from_diag(indices, Self::any_scalar_payload_to_real(data))
7142 }
7143 }
7144
7145 pub fn copy_tensor(
7169 indices: Vec<DynIndex>,
7170 value: AnyScalar,
7171 ) -> std::result::Result<Self, IdxTensorError> {
7172 if indices.is_empty() {
7173 return Self::from_dense_any(vec![], vec![value]);
7174 }
7175 let dim = indices[0].dim();
7176 let data = vec![value; dim];
7177 Self::from_diag_any(indices, data)
7178 }
7179
7180 pub fn fuse_indices(
7237 &self,
7238 old_indices: &[DynIndex],
7239 new_index: DynIndex,
7240 order: LinearizationOrder,
7241 ) -> std::result::Result<Self, IdxTensorError> {
7242 if !(!old_indices.is_empty()) {
7243 return Err(anyhow::anyhow!("fuse_indices requires at least one index to fuse").into());
7244 };
7245
7246 let old_dims = self.dims();
7247 let mut seen_indices = HashSet::new();
7248 let mut old_axes = Vec::with_capacity(old_indices.len());
7249 for old_index in old_indices {
7250 if !(seen_indices.insert(old_index)) {
7251 return Err(anyhow::anyhow!("duplicate index in old_indices").into());
7252 };
7253 let axis = self
7254 .indices
7255 .iter()
7256 .position(|idx| idx == old_index)
7257 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
7258 if !(old_index.dim() == old_dims[axis]) {
7259 return Err(anyhow::anyhow!(
7260 "old index dimension does not match tensor axis dimension"
7261 )
7262 .into());
7263 };
7264 old_axes.push(axis);
7265 }
7266
7267 let fused_dims: Vec<usize> = old_axes.iter().map(|&axis| old_dims[axis]).collect();
7268 let fused_product = checked_product(&fused_dims)?;
7269 if !(fused_product == new_index.dim()) {
7270 return Err(anyhow::anyhow!(
7271 "product of old index dimensions must match the replacement index dimension"
7272 )
7273 .into());
7274 };
7275
7276 let insertion_axis =
7277 old_axes.iter().copied().min().ok_or_else(|| {
7278 anyhow::anyhow!("fuse_indices requires at least one index to fuse")
7279 })?;
7280 let old_axis_set: HashSet<usize> = old_axes.iter().copied().collect();
7281
7282 let mut result_indices =
7283 Vec::with_capacity(self.indices.len() - old_indices.len() + 1usize);
7284 for (axis, index) in self.indices.iter().enumerate() {
7285 if axis == insertion_axis {
7286 result_indices.push(new_index.clone());
7287 }
7288 if !old_axis_set.contains(&axis) {
7289 result_indices.push(index.clone());
7290 }
7291 }
7292 let mut result_seen = HashSet::new();
7293 for index in &result_indices {
7294 if !(result_seen.insert(index)) {
7295 return Err(
7296 anyhow::anyhow!("fuse_indices result would contain duplicate index").into(),
7297 );
7298 };
7299 }
7300 Self::validate_indices(&result_indices)?;
7301
7302 let mut new_dims = Vec::with_capacity(old_dims.len() - old_indices.len() + 1usize);
7303 for (axis, dim) in old_dims.iter().copied().enumerate() {
7304 if axis == insertion_axis {
7305 new_dims.push(new_index.dim());
7306 }
7307 if !old_axis_set.contains(&axis) {
7308 new_dims.push(dim);
7309 }
7310 }
7311
7312 self.ensure_shape_packing_preserves_ad("fuse_indices")?;
7313
7314 let mut grouped_axes = old_axes.clone();
7315 if matches!(order, LinearizationOrder::RowMajor) {
7316 grouped_axes.reverse();
7317 }
7318 let mut perm = Vec::with_capacity(self.indices.len());
7319 perm.extend((0..insertion_axis).filter(|axis| !old_axis_set.contains(axis)));
7320 perm.extend(grouped_axes);
7321 perm.extend(
7322 ((insertion_axis + 1)..self.indices.len()).filter(|axis| !old_axis_set.contains(axis)),
7323 );
7324 debug_assert_eq!(perm.len(), self.indices.len());
7325
7326 let packed = self.permute(&perm)?;
7327 let reshaped = packed.try_materialized_inner()?.reshape(&new_dims)?;
7328 Self::from_inner(result_indices, reshaped).map_err(IdxTensorError::from)
7329 }
7330
7331 pub fn unfuse_index(
7356 &self,
7357 old_index: &DynIndex,
7358 new_indices: &[DynIndex],
7359 order: LinearizationOrder,
7360 ) -> std::result::Result<Self, IdxTensorError> {
7361 if !(!new_indices.is_empty()) {
7362 return Err(
7363 anyhow::anyhow!("unfuse_index requires at least one replacement index").into(),
7364 );
7365 };
7366
7367 let axis = self
7368 .indices
7369 .iter()
7370 .position(|idx| idx == old_index)
7371 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
7372
7373 let replacement_dims: Vec<usize> = new_indices.iter().map(DynIndex::dim).collect();
7374 let replacement_product = checked_product(&replacement_dims)?;
7375 if !(replacement_product == old_index.dim()) {
7376 return Err(anyhow::anyhow!(
7377 "product of new index dimensions must match the replaced index dimension"
7378 )
7379 .into());
7380 };
7381
7382 let mut result_indices =
7383 Vec::with_capacity(self.indices.len() - 1usize + new_indices.len());
7384 result_indices.extend_from_slice(&self.indices[..axis]);
7385 result_indices.extend(new_indices.iter().cloned());
7386 result_indices.extend_from_slice(&self.indices[axis + 1..]);
7387 Self::validate_indices(&result_indices)?;
7388
7389 let old_dims = self.dims();
7390 let mut new_dims = Vec::with_capacity(old_dims.len() - 1usize + replacement_dims.len());
7391 new_dims.extend_from_slice(&old_dims[..axis]);
7392 new_dims.extend_from_slice(&replacement_dims);
7393 new_dims.extend_from_slice(&old_dims[axis + 1..]);
7394
7395 self.ensure_shape_packing_preserves_ad("unfuse_index")?;
7396
7397 let mut grouped_indices = new_indices.to_vec();
7398 let mut grouped_dims = replacement_dims.clone();
7399 if matches!(order, LinearizationOrder::RowMajor) {
7400 grouped_indices.reverse();
7401 grouped_dims.reverse();
7402 }
7403 let mut packed_indices =
7404 Vec::with_capacity(self.indices.len() - 1usize + grouped_indices.len());
7405 packed_indices.extend_from_slice(&self.indices[..axis]);
7406 packed_indices.extend(grouped_indices);
7407 packed_indices.extend_from_slice(&self.indices[axis + 1..]);
7408
7409 let mut packed_dims = Vec::with_capacity(old_dims.len() - 1usize + grouped_dims.len());
7410 packed_dims.extend_from_slice(&old_dims[..axis]);
7411 packed_dims.extend_from_slice(&grouped_dims);
7412 packed_dims.extend_from_slice(&old_dims[axis + 1..]);
7413
7414 let reshaped = self.try_materialized_inner()?.reshape(&packed_dims)?;
7415 let packed = Self::from_inner(packed_indices, reshaped)?;
7416 if matches!(order, LinearizationOrder::ColumnMajor) {
7417 Ok(packed)
7418 } else {
7419 packed.permute_indices(&result_indices)
7420 }
7421 }
7422
7423 pub fn scalar<T: TensorElement>(value: T) -> std::result::Result<Self, IdxTensorError> {
7437 Self::from_dense(vec![], vec![value])
7438 }
7439
7440 pub fn zeros<T: TensorElement + Zero + Clone>(
7456 indices: Vec<DynIndex>,
7457 ) -> std::result::Result<Self, IdxTensorError> {
7458 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
7459 let size = checked_product(&dims)?;
7460 Self::from_dense(indices, vec![T::zero(); size])
7461 }
7462}
7463
7464impl IdxTensor {
7469 pub fn to_vec<T: TensorElement>(&self) -> std::result::Result<Vec<T>, IdxTensorError> {
7493 self.as_inner()?
7494 .duplicate_value()?
7495 .as_slice::<T>()
7496 .map(|values| values.to_vec())
7497 .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))
7498 }
7499
7500 pub fn with_dense_slice<T: TensorElement, R>(
7539 &self,
7540 read: impl FnOnce(&[T]) -> R,
7541 ) -> std::result::Result<R, IdxTensorError> {
7542 let inner = self.as_inner()?;
7543 let value = inner.value()?;
7544 if let Ok(values) = value.as_slice::<T>() {
7545 return Ok(read(values));
7546 }
7547
7548 drop(value);
7553 let values = inner.duplicate_value()?;
7554 let values = values
7555 .as_slice::<T>()
7556 .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))?;
7557 Ok(read(values))
7558 }
7559
7560 pub fn into_dense_col_major_parts<T: TensorElement>(
7595 self,
7596 ) -> std::result::Result<(Vec<DynIndex>, Vec<T>), IdxTensorError> {
7597 if !(!self.tracks_grad()) {
7598 return Err(anyhow::anyhow!("IdxTensor::into_dense_col_major_parts cannot consume tensors with tracked autodiff state").into());
7599 };
7600 let data = self.to_vec::<T>()?;
7601 Ok((self.indices, data))
7602 }
7603
7604 pub fn is_f64(&self) -> bool {
7617 self.storage.dtype() == Some(DType::F64)
7618 }
7619
7620 pub fn is_f32(&self) -> bool {
7635 self.storage.dtype() == Some(DType::F32)
7636 }
7637
7638 pub fn is_c32(&self) -> bool {
7654 self.storage.dtype() == Some(DType::C32)
7655 }
7656
7657 pub fn is_c64(&self) -> bool {
7673 self.storage.is_c64()
7674 }
7675
7676 pub fn is_diag(&self) -> bool {
7703 self.storage.is_diag()
7704 }
7705
7706 pub fn is_complex(&self) -> bool {
7725 self.storage.is_complex()
7726 }
7727 pub fn from_dense_in<T: TensorElement>(
7754 context: &ExecutionContext,
7755 indices: Vec<DynIndex>,
7756 data: Vec<T>,
7757 ) -> std::result::Result<Self, IdxTensorError> {
7758 let dims = Self::expected_dims_from_indices(&indices);
7759 Self::validate_indices(&indices)?;
7760 Self::validate_dense_payload_len(data.len(), &dims)?;
7761 let native = dense_native_tensor_from_col_major(&data, &dims)
7762 .map_err(|error| IdxTensorError::operation("context-scoped construction", error))?;
7763 let inner = match context {
7764 ExecutionContext::Cpu(context) => {
7765 let runtime = context.eager_runtime().map_err(|error| {
7766 IdxTensorError::operation("CPU context construction", anyhow::Error::new(error))
7767 })?;
7768 EagerTensor::from_tensor_in(native, runtime).map_err(|error| {
7769 IdxTensorError::operation(
7770 "CPU context eager wrapping",
7771 anyhow::Error::new(error),
7772 )
7773 })?
7774 }
7775 #[cfg(feature = "tenferro-cuda")]
7776 ExecutionContext::Cuda(context) => {
7777 let uploaded = context.upload_cuda(&native).map_err(|error| {
7778 IdxTensorError::operation(
7779 "CUDA context construction",
7780 anyhow::Error::new(error),
7781 )
7782 })?;
7783 EagerTensor::from_tensor_in(
7784 uploaded,
7785 context.eager_runtime().map_err(|error| {
7786 IdxTensorError::operation(
7787 "CUDA context eager runtime",
7788 anyhow::Error::new(error),
7789 )
7790 })?,
7791 )
7792 .map_err(|error| {
7793 IdxTensorError::operation(
7794 "CUDA context eager wrapping",
7795 anyhow::Error::new(error),
7796 )
7797 })?
7798 }
7799 };
7800 Self::from_inner(indices, inner).map_err(IdxTensorError::from)
7801 }
7802
7803 pub fn ones_in(
7827 context: &ExecutionContext,
7828 indices: &[DynIndex],
7829 ) -> std::result::Result<Self, IdxTensorError> {
7830 let dims = Self::expected_dims_from_indices(indices);
7831 let total_size = checked_total_size(&dims)?;
7832 Self::from_dense_in(context, indices.to_vec(), vec![1.0_f64; total_size])
7833 }
7834
7835 pub fn validate_context(
7859 &self,
7860 context: &ExecutionContext,
7861 ) -> std::result::Result<(), IdxTensorError> {
7862 let Some(inner) = self.storage.eager() else {
7863 if matches!(context, ExecutionContext::Cpu(_)) {
7864 return Ok(());
7865 }
7866 return Err(IdxTensorError::operation(
7867 "context validation",
7868 anyhow::anyhow!("tensor has no CUDA eager runtime"),
7869 ));
7870 };
7871 match context {
7872 ExecutionContext::Cpu(context) => {
7873 let expected = context.eager_runtime().map_err(|error| {
7874 IdxTensorError::operation("CPU context validation", anyhow::Error::new(error))
7875 })?;
7876 if inner.ctx_id() != expected.id() {
7877 return Err(IdxTensorError::operation(
7878 "context validation",
7879 anyhow::anyhow!("tensor belongs to a different CPU eager context"),
7880 ));
7881 }
7882 }
7883 #[cfg(feature = "tenferro-cuda")]
7884 ExecutionContext::Cuda(context) => {
7885 self.validate_cuda_residency(context).map_err(|error| {
7886 IdxTensorError::operation("CUDA context validation", anyhow::Error::new(error))
7887 })?;
7888 }
7889 }
7890 Ok(())
7891 }
7892
7893 pub fn read_decision_data(
7936 &self,
7937 context: &ExecutionContext,
7938 ) -> std::result::Result<Vec<f64>, IdxTensorError> {
7939 self.validate_context(context)?;
7940 #[cfg(feature = "tenferro-cuda")]
7941 if let ExecutionContext::Cuda(cuda) = context {
7942 let host = self.download(cuda).map_err(|error| {
7943 IdxTensorError::operation("decision readback download", anyhow::Error::new(error))
7944 })?;
7945 return Self::decision_values(&host);
7946 }
7947 Self::decision_values(self)
7948 }
7949
7950 fn decision_values(host: &Self) -> std::result::Result<Vec<f64>, IdxTensorError> {
7951 if host.is_f64() {
7952 host.to_vec::<f64>().map_err(|error| {
7953 IdxTensorError::operation("decision readback", anyhow::Error::new(error))
7954 })
7955 } else if host.is_c64() {
7956 host.to_vec::<Complex64>()
7957 .map(|values| values.into_iter().map(|value| value.re).collect())
7958 .map_err(|error| {
7959 IdxTensorError::operation("decision readback", anyhow::Error::new(error))
7960 })
7961 } else {
7962 Err(IdxTensorError::operation(
7963 "decision readback",
7964 anyhow::anyhow!("decision payloads support f64 and Complex64 storage only"),
7965 ))
7966 }
7967 }
7968
7969 pub fn scale_in(
8008 &self,
8009 factor: f64,
8010 context: &ExecutionContext,
8011 ) -> std::result::Result<Self, IdxTensorError> {
8012 self.validate_context(context)?;
8013 let adopted_storage;
8014 let operand: &EagerTensor = match self.eager_runtime() {
8015 Some(_) => self.as_inner()?,
8016 None => {
8019 adopted_storage = Self::adopt_host_into_context(self, context)?;
8020 adopted_storage.as_inner()?
8021 }
8022 };
8023 let scalar = Self::context_scalar_in(operand, factor, context)?;
8024 let scaled = operand.mul(&scalar).map_err(|error| {
8025 IdxTensorError::operation(
8026 "context-scoped scaling",
8027 anyhow::Error::new(error).context("eager multiplication failed"),
8028 )
8029 })?;
8030 Self::from_inner(self.indices.clone(), scaled).map_err(IdxTensorError::from)
8031 }
8032
8033 fn adopt_host_into_context(
8035 tensor: &Self,
8036 context: &ExecutionContext,
8037 ) -> std::result::Result<Self, IdxTensorError> {
8038 if tensor.is_f64() {
8039 let values = tensor.to_vec::<f64>().map_err(|error| {
8040 IdxTensorError::operation("context adoption", anyhow::Error::new(error))
8041 })?;
8042 Self::from_dense_in(context, tensor.indices.clone(), values)
8043 } else if tensor.is_c64() {
8044 let values = tensor.to_vec::<Complex64>().map_err(|error| {
8045 IdxTensorError::operation("context adoption", anyhow::Error::new(error))
8046 })?;
8047 Self::from_dense_in(context, tensor.indices.clone(), values)
8048 } else {
8049 Err(IdxTensorError::operation(
8050 "context adoption",
8051 anyhow::anyhow!("adoption supports f64 and Complex64 storage only"),
8052 ))
8053 }
8054 }
8055
8056 fn context_scalar_in(
8058 operand: &EagerTensor,
8059 factor: f64,
8060 context: &ExecutionContext,
8061 ) -> std::result::Result<EagerTensor, IdxTensorError> {
8062 let native: NativeTensor = match operand.dtype() {
8063 DType::F64 => NativeTensor::from_vec_col_major(vec![], vec![factor]),
8064 DType::C64 => {
8065 NativeTensor::from_vec_col_major(vec![], vec![Complex64::new(factor, 0.0)])
8066 }
8067 other => {
8068 return Err(IdxTensorError::operation(
8069 "context-scoped scaling",
8070 anyhow::anyhow!("unsupported scalar dtype {other:?}"),
8071 ));
8072 }
8073 }
8074 .map_err(|error| {
8075 IdxTensorError::operation("context-scoped scaling", anyhow::anyhow!("{error}"))
8076 })?;
8077 let runtime = Self::context_eager_runtime(context)?;
8078 let resident = match context {
8079 ExecutionContext::Cpu(_) => native,
8080 #[cfg(feature = "tenferro-cuda")]
8081 ExecutionContext::Cuda(cuda) => cuda.upload_cuda(&native).map_err(|error| {
8082 IdxTensorError::operation(
8083 "context-scoped scaling",
8084 anyhow::Error::new(error).context("scalar upload failed"),
8085 )
8086 })?,
8087 };
8088 EagerTensor::from_tensor_in(resident, runtime).map_err(|error| {
8089 IdxTensorError::operation(
8090 "context-scoped scaling",
8091 anyhow::Error::new(error).context("scalar wrapping failed"),
8092 )
8093 })
8094 }
8095
8096 fn context_eager_runtime(
8098 context: &ExecutionContext,
8099 ) -> std::result::Result<Arc<EagerRuntime>, IdxTensorError> {
8100 match context {
8101 ExecutionContext::Cpu(context) => context.eager_runtime().map_err(|error| {
8102 IdxTensorError::operation("context-scoped scaling", anyhow::Error::new(error))
8103 }),
8104 #[cfg(feature = "tenferro-cuda")]
8105 ExecutionContext::Cuda(context) => context.eager_runtime().map_err(|error| {
8106 IdxTensorError::operation("context-scoped scaling", anyhow::Error::new(error))
8107 }),
8108 }
8109 }
8110
8111 pub fn norm_in(&self, context: &ExecutionContext) -> std::result::Result<f64, IdxTensorError> {
8123 self.validate_context(context)?;
8124 #[cfg(feature = "tenferro-cuda")]
8125 if matches!(context, ExecutionContext::Cuda(_)) {
8126 let inner = self.cuda_eager_inner().ok_or_else(|| {
8127 IdxTensorError::operation(
8128 "context-scoped norm",
8129 anyhow::anyhow!("tensor has no resident eager value"),
8130 )
8131 })?;
8132 let rank = inner.shape().len();
8133 let axes: Vec<usize> = (0..rank).collect();
8134 let squared = inner
8135 .abs()
8136 .and_then(|magnitudes| magnitudes.reduce_sum_squares(&axes))
8137 .map_err(|error| {
8138 IdxTensorError::operation("context-scoped norm", anyhow::anyhow!("{error}"))
8139 })?;
8140 let scalar = squared.reshape(&[1]).map_err(|error| {
8141 IdxTensorError::operation("context-scoped norm", anyhow::anyhow!("{error}"))
8142 })?;
8143 let values = Self::read_resident_vector(&scalar, 1, context).map_err(|error| {
8144 IdxTensorError::operation("context-scoped norm", anyhow::Error::new(error))
8145 })?;
8146 return Ok(values[0].sqrt());
8147 }
8148 #[cfg(not(feature = "tenferro-cuda"))]
8149 let _ = context;
8150 self.norm()
8151 }
8152}
8153
8154fn checked_product(dims: &[usize]) -> Result<usize> {
8155 dims.iter().try_fold(1usize, |acc, &dim| {
8156 acc.checked_mul(dim)
8157 .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
8158 })
8159}
8160
8161fn increment_col_major_coordinate(coords: &mut [usize], dims: &[usize]) {
8162 let mut carry = true;
8163 for (coordinate, &dim) in coords.iter_mut().zip(dims.iter()) {
8164 if !carry {
8165 break;
8166 }
8167 *coordinate += 1;
8168 if *coordinate == dim {
8169 *coordinate = 0;
8170 } else {
8171 carry = false;
8172 }
8173 }
8174}
8175
8176fn map_payload_support_coordinate(
8177 source_axis_classes: &[usize],
8178 target_axis_classes: &[usize],
8179 source_to_target_axes: &[usize],
8180 source_coords: &[usize],
8181 target_coords: &mut [usize],
8182 target_seen: &mut [bool],
8183) -> bool {
8184 if source_axis_classes.len() != source_to_target_axes.len()
8185 || target_coords.len() != target_seen.len()
8186 {
8187 return false;
8188 }
8189 target_seen.fill(false);
8190 for (source_axis, &target_axis) in source_to_target_axes.iter().enumerate() {
8191 let Some(&source_class) = source_axis_classes.get(source_axis) else {
8192 return false;
8193 };
8194 let Some(&target_class) = target_axis_classes.get(target_axis) else {
8195 return false;
8196 };
8197 let Some(&source_value) = source_coords.get(source_class) else {
8198 return false;
8199 };
8200 let Some(target_value) = target_coords.get_mut(target_class) else {
8201 return false;
8202 };
8203 if target_seen[target_class] {
8204 if *target_value != source_value {
8205 return false;
8206 }
8207 } else {
8208 *target_value = source_value;
8209 target_seen[target_class] = true;
8210 }
8211 }
8212 target_seen.iter().all(|&seen| seen)
8213}
8214
8215fn decode_col_major_linear(linear: usize, dims: &[usize]) -> Result<Vec<usize>> {
8216 let total = checked_product(dims)?;
8217 if !(linear < total) {
8218 return Err(anyhow::anyhow!(
8219 "linear offset {} out of bounds for dims {:?}",
8220 linear,
8221 dims
8222 ));
8223 };
8224 let mut remaining = linear;
8225 let mut out = Vec::with_capacity(dims.len());
8226 for &dim in dims {
8227 out.push(remaining % dim);
8228 remaining /= dim;
8229 }
8230 Ok(out)
8231}
8232
8233fn encode_col_major_linear(indices: &[usize], dims: &[usize]) -> Result<usize> {
8234 if !(indices.len() == dims.len()) {
8235 return Err(anyhow::anyhow!(
8236 "index rank {} does not match dims {:?}",
8237 indices.len(),
8238 dims
8239 ));
8240 };
8241 let mut linear = 0usize;
8242 let mut stride = 1usize;
8243 for (&index, &dim) in indices.iter().zip(dims.iter()) {
8244 if !(index < dim) {
8245 return Err(anyhow::anyhow!(
8246 "index {} out of bounds for dimension {}",
8247 index,
8248 dim
8249 ));
8250 };
8251 let term = index
8252 .checked_mul(stride)
8253 .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
8254 linear = linear
8255 .checked_add(term)
8256 .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
8257 stride = stride
8258 .checked_mul(dim)
8259 .ok_or_else(|| anyhow::anyhow!("stride overflow"))?;
8260 }
8261 Ok(linear)
8262}
8263
8264#[cfg(test)]
8265mod tests {
8266 use super::*;
8267 use num_complex::{Complex32, Complex64};
8268 use std::cell::Cell;
8269 use tensor4all_tensorbackend::StorageError;
8270
8271 #[test]
8272 fn structured_contraction_does_not_install_logical_dense_cache() {
8273 let n = 8;
8274 let left = DynIndex::new_dyn(n);
8275 let site = DynIndex::new_dyn(3);
8276 let right = DynIndex::new_dyn(n);
8277 let far = DynIndex::new_dyn(n);
8278 let end = DynIndex::new_dyn(n);
8279 let a =
8280 IdxTensor::from_copy_selector(left, site.clone(), right.clone(), 1, 1.0_f64).unwrap();
8281 let b =
8282 IdxTensor::from_copy_selector(right, site.clone(), far.clone(), 1, 2.0_f64).unwrap();
8283 let c = IdxTensor::from_copy_selector(far, site.clone(), end, 1, 3.0_f64).unwrap();
8284 let result = crate::defaults::contract::contract_with_options(
8285 &[&a, &b, &c],
8286 crate::defaults::contract::ContractionOptions::new()
8287 .with_retain_indices(std::slice::from_ref(&site)),
8288 )
8289 .unwrap();
8290
8291 assert_eq!(result.storage_kind(), StorageKind::Structured);
8292 assert!(result.eager_cache.get().is_none());
8293 assert_eq!(result.storage().unwrap().payload_len(), n * 3);
8294 }
8295
8296 #[test]
8297 fn binary_contraction_axis_classes_preserve_uncontracted_order() {
8298 let axis_classes =
8299 IdxTensor::binary_contraction_axis_classes(&[0, 1, 0], &[1], &[0, 1], &[0]).unwrap();
8300
8301 assert_eq!(axis_classes, vec![0, 0, 1]);
8302 }
8303
8304 #[test]
8305 fn structured_metrics_use_authoritative_compact_payload_for_all_dtypes() {
8306 fn check(tensor: IdxTensor, expected_sum: f64, expected_norm_squared: f64) {
8307 assert!(matches!(tensor.storage, IdxTensorStorage::Compact(_)));
8308 assert!(tensor.eager_cache.get().is_none());
8309 assert!((tensor.sum().unwrap().real() - expected_sum).abs() < 1.0e-6);
8310 assert!((tensor.norm_squared().unwrap() - expected_norm_squared).abs() < 1.0e-6);
8311 assert!((tensor.maxabs().unwrap() - 2.0).abs() < 1.0e-6);
8312 assert!(tensor.isapprox(&tensor, 0.0, 0.0).unwrap());
8313 assert!(tensor.eager_cache.get().is_none());
8314 }
8315
8316 let indices = || vec![DynIndex::new_dyn(2), DynIndex::new_dyn(2)];
8317 check(
8318 IdxTensor::from_diag(indices(), vec![1.0_f32, 2.0]).unwrap(),
8319 3.0,
8320 5.0,
8321 );
8322 check(
8323 IdxTensor::from_diag(indices(), vec![1.0_f64, 2.0]).unwrap(),
8324 3.0,
8325 5.0,
8326 );
8327 check(
8328 IdxTensor::from_diag(
8329 indices(),
8330 vec![Complex32::new(1.0, 0.0), Complex32::new(2.0, 0.0)],
8331 )
8332 .unwrap(),
8333 3.0,
8334 5.0,
8335 );
8336 check(
8337 IdxTensor::from_diag(
8338 indices(),
8339 vec![Complex64::new(1.0, 0.0), Complex64::new(2.0, 0.0)],
8340 )
8341 .unwrap(),
8342 3.0,
8343 5.0,
8344 );
8345 }
8346
8347 #[test]
8348 fn payload_storage_error_retains_typed_source() {
8349 let storage = Storage::from_dense_col_major(vec![1.0_f64, 2.0], &[2]).unwrap();
8350 let error = storage.scalar_at(&[2]).unwrap_err();
8351 assert!(matches!(error, StorageError::InvalidStructuredStorage(_)));
8352 }
8353
8354 #[test]
8355 fn materialization_error_retains_backend_source() {
8356 let native = NativeTensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
8357 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx().unwrap()).unwrap();
8358 let storage = IdxTensorStorage::Eager {
8359 inner: Arc::new(inner),
8360 axis_classes: vec![0, 0],
8361 };
8362
8363 let error = storage.materialize(2).unwrap_err();
8364 assert!(matches!(error, TensorStorageError::Materialization { .. }));
8365 assert!(std::error::Error::source(&error).is_some());
8366 }
8367
8368 fn conjugate_with_injected_failure(
8369 tensor: &IdxTensor,
8370 target: *const EagerTensor,
8371 message: &'static str,
8372 ) -> IdxTensor {
8373 let calls = Cell::new(0usize);
8374 let conjugated = tensor.conj_with(&|inner| {
8375 calls.set(calls.get() + 1);
8376 if std::ptr::eq(inner, target) {
8377 Err(Arc::new(std::io::Error::other(message)) as _)
8378 } else {
8379 conjugate_eager(inner)
8380 }
8381 });
8382 assert!(calls.get() > 0, "injected closure was not reached");
8383 conjugated
8384 }
8385
8386 fn assert_unwrapped_conjugation_error(
8387 tensor: IdxTensor,
8388 target: *const EagerTensor,
8389 message: &'static str,
8390 ) -> IdxTensor {
8391 let conjugated = conjugate_with_injected_failure(&tensor, target, message);
8392 let error = conjugated.to_storage().unwrap_err();
8393 assert!(matches!(error, TensorStorageError::Conjugation { .. }));
8394 let source = std::error::Error::source(&error).unwrap();
8395 assert_eq!(source.to_string(), message);
8396 assert!(
8397 source.source().is_none(),
8398 "source was wrapped more than once"
8399 );
8400 conjugated
8401 }
8402
8403 #[test]
8404 fn authoritative_storage_conjugation_failure_is_deferred_without_detaching() {
8405 let i = DynIndex::new_dyn(2);
8406 let native = NativeTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
8407 let inner = EagerTensor::requires_grad_in(native, default_eager_ctx().unwrap()).unwrap();
8408 let tensor = IdxTensor::from_inner(vec![i], inner).unwrap();
8409 let source = match &tensor.storage {
8410 IdxTensorStorage::Eager { inner, .. } => Arc::clone(inner),
8411 IdxTensorStorage::Compact(payload) => Arc::clone(&payload.payload),
8412 IdxTensorStorage::Materialized(_) | IdxTensorStorage::Deferred { .. } => {
8413 panic!("tracked eager source expected")
8414 }
8415 };
8416
8417 let conjugated = conjugate_with_injected_failure(
8418 &tensor,
8419 Arc::as_ptr(&source),
8420 "forced authoritative eager conjugation failure",
8421 );
8422 assert!(conjugated.tracks_grad());
8423 assert!(conjugated.is_f64());
8424 assert!(!conjugated.is_complex());
8425 assert!(!conjugated.is_diag());
8426 assert_eq!(conjugated.dims(), vec![2]);
8427
8428 let error = conjugated.to_storage().unwrap_err();
8429 let source = std::error::Error::source(&error).unwrap();
8430 assert_eq!(
8431 source.to_string(),
8432 "forced authoritative eager conjugation failure"
8433 );
8434 assert!(conjugated.detach().is_err());
8435 }
8436
8437 #[test]
8438 fn structured_payload_conjugation_failure_retains_graph_and_blocks_detached_primal() {
8439 let i = DynIndex::new_dyn(2);
8440 let j = DynIndex::new_dyn(2);
8441 let tensor = IdxTensor::from_diag(
8442 vec![i, j],
8443 vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, -4.0)],
8444 )
8445 .unwrap()
8446 .enable_grad()
8447 .unwrap();
8448
8449 let target = Arc::as_ptr(
8450 &tensor
8451 .storage
8452 .compact_payload()
8453 .expect("tracked compact payload")
8454 .payload,
8455 );
8456 let conjugated = assert_unwrapped_conjugation_error(
8457 tensor,
8458 target,
8459 "forced structured AD conjugation failure",
8460 );
8461 assert!(conjugated.tracks_grad());
8462 assert!(conjugated.detach().is_err());
8463 assert!(conjugated.clone().enable_grad().is_err());
8464 assert!(conjugated.sum().is_err());
8465 assert!(conjugated.grad().is_err());
8466 assert!(conjugated.clear_grad().is_err());
8467 assert!(conjugated.maxabs().is_err());
8468 assert!(conjugated.norm_squared().is_err());
8469
8470 let twice_conjugated = conjugated.conj();
8471 let error = twice_conjugated.to_storage().unwrap_err();
8472 assert_eq!(
8473 std::error::Error::source(&error).unwrap().to_string(),
8474 "forced structured AD conjugation failure"
8475 );
8476 }
8477
8478 #[test]
8479 fn eager_cache_conjugation_failure_is_deferred_with_original_diagnostic() {
8480 let i = DynIndex::new_dyn(2);
8481 let j = DynIndex::new_dyn(2);
8482 let tensor = IdxTensor::from_diag(vec![i, j], vec![1.0_f64, 2.0]).unwrap();
8483 tensor.as_inner().unwrap();
8484 let target = Arc::as_ptr(tensor.eager_cache.get().unwrap());
8485
8486 let conjugated = assert_unwrapped_conjugation_error(
8487 tensor,
8488 target,
8489 "forced eager cache conjugation failure",
8490 );
8491 assert!(!conjugated.tracks_grad());
8492 assert!(conjugated.detach().is_err());
8493 }
8494
8495 #[test]
8496 fn encode_col_major_linear_rejects_offset_overflow() {
8497 let error =
8498 encode_col_major_linear(&[usize::MAX - 1, usize::MAX - 1], &[usize::MAX, usize::MAX])
8499 .unwrap_err();
8500 assert!(error.to_string().contains("linear offset overflow"));
8501 }
8502
8503 #[test]
8504 fn factorize_probe_batch_incremental_matches_the_column_based_path() {
8505 let row = DynIndex::new_dyn(6);
8506 let batch = DynIndex::new_dyn(4);
8507 let data: Vec<f64> = (0..24).map(|i| i as f64 * 0.37 - 1.5).collect();
8508 let batch_tensor =
8509 IdxTensor::from_dense(vec![row.clone(), batch.clone()], data.clone()).unwrap();
8510
8511 let columns: Vec<IdxTensor> = (0..4)
8512 .map(|position| {
8513 batch_tensor
8514 .select_indices(std::slice::from_ref(&batch), &[position])
8515 .unwrap()
8516 })
8517 .collect();
8518 let column_refs: Vec<&IdxTensor> = columns.iter().collect();
8519
8520 let from_batch = IdxTensor::factorize_probe_batch_incremental(
8521 None,
8522 &batch_tensor,
8523 &batch,
8524 std::slice::from_ref(&row),
8525 )
8526 .unwrap();
8527 let from_columns = IdxTensor::factorize_probe_columns_incremental(
8528 None,
8529 &column_refs,
8530 &column_refs,
8531 &[row],
8532 )
8533 .unwrap();
8534
8535 assert_eq!(from_batch.rank, from_columns.rank);
8536 let batch_data = from_batch.left.to_vec::<f64>().unwrap();
8537 let columns_data = from_columns.left.to_vec::<f64>().unwrap();
8538 assert_eq!(batch_data.len(), columns_data.len());
8539 let max_diff = batch_data
8540 .iter()
8541 .zip(columns_data.iter())
8542 .map(|(a, b)| (a - b).abs())
8543 .fold(0.0_f64, f64::max);
8544 assert!(
8545 max_diff < 1e-12,
8546 "batch-native and column-based factorizations disagree: max diff = {}",
8547 max_diff
8548 );
8549 }
8550}