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::{DType, DotGeneralConfig, Tensor as NativeTensor, TensorRead, TensorView};
18use tenferro_ad::EagerTensor;
19use tenferro_einsum::{EagerEinsumExt, EinsumSubscripts};
20use tenferro_linalg::EagerTensorLinalgExt;
21use tensor4all_tensorbackend::{
22 contract_native_tensor, default_eager_ctx, dense_native_tensor_from_col_major,
23 diag_native_tensor_from_col_major, native_tensor_primal_to_diag,
24 storage_payload_native_read_input, storage_to_native_tensor, NativeTensorReadInput,
25 TensorElement,
26};
27use tensor4all_tensorbackend::{Storage, StorageKind};
28
29use super::contract::PairwiseContractionOptions;
30use super::structured_contraction::{
31 normalize_payload_read_for_roots, storage_from_payload_native, storage_payload_native,
32 OperandLayout, StructuredContractionPlan, StructuredContractionSpec,
33};
34
35fn conjugate_eager(
36 inner: &EagerTensor,
37) -> std::result::Result<EagerTensor, Arc<dyn std::error::Error + Send + Sync + 'static>> {
38 inner.conj().map_err(|source| Arc::new(source) as _)
39}
40
41#[derive(Debug, Default, Clone)]
42struct PairwiseContractProfileEntry {
43 calls: usize,
44 total_time: Duration,
45 total_bytes: usize,
46}
47
48#[derive(Debug, Clone)]
70pub struct TensorHermitianEigendecomposition {
71 pub eigenvalues: Vec<f64>,
73 pub eigenvectors: IdxTensor,
75 pub eigenvector_index: DynIndex,
77}
78
79thread_local! {
80 static PAIRWISE_CONTRACT_PROFILE_STATE: RefCell<HashMap<&'static str, PairwiseContractProfileEntry>> =
81 RefCell::new(HashMap::new());
82}
83
84fn pairwise_contract_profile_enabled() -> bool {
85 static ENABLED: OnceLock<bool> = OnceLock::new();
86 *ENABLED.get_or_init(|| env::var("T4A_PROFILE_PAIRWISE_CONTRACT").is_ok())
87}
88
89fn record_pairwise_contract_profile(section: &'static str, elapsed: Duration) {
90 if !pairwise_contract_profile_enabled() {
91 return;
92 }
93 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
94 let mut state = state.borrow_mut();
95 let entry = state.entry(section).or_default();
96 entry.calls += 1;
97 entry.total_time += elapsed;
98 });
99}
100
101fn record_pairwise_contract_profile_bytes(section: &'static str, bytes: usize) {
102 if !pairwise_contract_profile_enabled() {
103 return;
104 }
105 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
106 let mut state = state.borrow_mut();
107 let entry = state.entry(section).or_default();
108 entry.total_bytes += bytes;
109 });
110}
111
112fn profile_pairwise_contract_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
113 if !pairwise_contract_profile_enabled() {
114 return f();
115 }
116 let started = Instant::now();
117 let result = f();
118 record_pairwise_contract_profile(section, started.elapsed());
119 result
120}
121
122pub fn reset_pairwise_contract_profile() {
124 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| state.borrow_mut().clear());
125}
126
127pub fn print_and_reset_pairwise_contract_profile() {
129 if !pairwise_contract_profile_enabled() {
130 return;
131 }
132 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
133 let mut entries: Vec<_> = state
134 .borrow()
135 .iter()
136 .map(|(section, entry)| (*section, entry.clone()))
137 .collect();
138 state.borrow_mut().clear();
139 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
140
141 eprintln!("=== IdxTensor pairwise contract profile ===");
142 for (section, entry) in entries {
143 let per_call_us = if entry.calls == 0 {
144 0.0
145 } else {
146 entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64
147 };
148 eprintln!(
149 "{section}: calls={} total={:.6}ms per_call={:.3}us bytes={}",
150 entry.calls,
151 entry.total_time.as_secs_f64() * 1.0e3,
152 per_call_us,
153 entry.total_bytes,
154 );
155 }
156 });
157}
158
159fn tensor_profile_bytes(dtype: DType, shape: &[usize]) -> usize {
160 let element_size = match dtype {
161 DType::F32 => 4,
162 DType::F64 => 8,
163 DType::C32 => 8,
164 DType::C64 => 16,
165 DType::I32 => 4,
166 DType::I64 => 8,
167 DType::Bool => 1,
168 };
169 shape
170 .iter()
171 .try_fold(1usize, |bytes, &dim| bytes.checked_mul(dim))
172 .and_then(|elements| elements.checked_mul(element_size))
173 .unwrap_or(usize::MAX)
174}
175
176pub trait RandomScalar: TensorElement {
180 fn random_value<R: Rng>(rng: &mut R) -> Self;
182}
183
184impl RandomScalar for f64 {
185 fn random_value<R: Rng>(rng: &mut R) -> Self {
186 StandardNormal.sample(rng)
187 }
188}
189
190impl RandomScalar for Complex64 {
191 fn random_value<R: Rng>(rng: &mut R) -> Self {
192 Complex64::new(StandardNormal.sample(rng), StandardNormal.sample(rng))
193 }
194}
195
196pub fn compute_permutation_from_indices(
222 original_indices: &[DynIndex],
223 new_indices: &[DynIndex],
224) -> std::result::Result<Vec<usize>, IdxTensorError> {
225 if !(new_indices.len() == original_indices.len()) {
226 return Err(
227 anyhow::anyhow!("new_indices length must match original_indices length").into(),
228 );
229 };
230
231 let mut perm = Vec::with_capacity(new_indices.len());
232 let mut used = std::collections::HashSet::new();
233
234 for new_idx in new_indices {
235 let pos = original_indices
238 .iter()
239 .position(|old_idx| old_idx == new_idx)
240 .ok_or_else(|| {
241 anyhow::anyhow!("new_indices must be a permutation of original_indices")
242 })?;
243
244 if !(used.insert(pos)) {
245 return Err(anyhow::anyhow!("duplicate index in new_indices").into());
246 };
247 perm.push(pos);
248 }
249
250 Ok(perm)
251}
252
253#[derive(Clone)]
258pub(crate) struct StructuredPayload {
259 payload: Arc<EagerTensor>,
260 payload_dims: Vec<usize>,
261 axis_classes: Vec<usize>,
262}
263
264#[derive(Debug, Clone, thiserror::Error)]
283pub enum TensorStorageError {
284 #[error("failed to materialize IdxTensor storage: {source}")]
286 Materialization {
287 #[source]
289 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
290 },
291 #[error(
293 "compact IdxTensor storage does not support dtype {dtype}; the eager payload remains authoritative"
294 )]
295 UnsupportedDtype {
296 dtype: &'static str,
298 },
299 #[error("failed to conjugate IdxTensor storage: {source}")]
302 Conjugation {
303 #[source]
305 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
306 },
307}
308
309#[derive(Debug, Clone, thiserror::Error)]
335pub enum IdxTensorError {
336 #[error("IdxTensor storage operation failed: {source}")]
338 Storage {
339 #[source]
341 source: TensorStorageError,
342 },
343 #[error("IdxTensor materialization failed: {source}")]
345 Materialization {
346 #[source]
348 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
349 },
350 #[error("IdxTensor scalar extraction failed: {source}")]
352 ScalarExtraction {
353 #[source]
355 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
356 },
357 #[error("IdxTensor scalar type mismatch: expected {expected}, got {actual}")]
360 ScalarTypeMismatch {
361 expected: &'static str,
363 actual: String,
365 },
366 #[error("IdxTensor shape mismatch during {operation}: expected {expected}, got {actual}")]
370 ShapeMismatch {
371 operation: &'static str,
373 expected: String,
375 actual: String,
377 },
378 #[error("IdxTensor subtraction failed: {source}")]
380 Subtraction {
381 #[source]
383 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
384 },
385 #[error("IdxTensor {operation} received NaN input")]
388 NaNInput {
389 operation: &'static str,
391 },
392 #[error("IdxTensor tolerance {name} is invalid: {value}")]
394 InvalidTolerance {
395 name: &'static str,
397 value: f64,
399 },
400 #[error("IdxTensor {operation} failed: {source}")]
402 Operation {
403 operation: &'static str,
405 #[source]
407 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
408 },
409}
410
411impl From<anyhow::Error> for IdxTensorError {
412 fn from(source: anyhow::Error) -> Self {
413 Self::operation("IdxTensor", source)
414 }
415}
416
417impl From<TensorStorageError> for IdxTensorError {
418 fn from(source: TensorStorageError) -> Self {
419 Self::Storage { source }
420 }
421}
422
423impl From<tenferro_ad::Error> for IdxTensorError {
424 fn from(source: tenferro_ad::Error) -> Self {
425 Self::Materialization {
426 source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
427 }
428 }
429}
430
431impl From<tensor4all_tensorbackend::EagerContextError> for IdxTensorError {
432 fn from(source: tensor4all_tensorbackend::EagerContextError) -> Self {
433 Self::Materialization {
434 source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
435 }
436 }
437}
438
439impl From<tensor4all_tensorbackend::BridgeError> for IdxTensorError {
440 fn from(source: tensor4all_tensorbackend::BridgeError) -> Self {
441 Self::Materialization {
442 source: Arc::new(source),
443 }
444 }
445}
446
447impl IdxTensorError {
448 fn boxed(error: anyhow::Error) -> Arc<dyn std::error::Error + Send + Sync + 'static> {
449 Arc::from(error.into_boxed_dyn_error())
450 }
451
452 fn materialization(error: anyhow::Error) -> Self {
453 Self::Materialization {
454 source: Self::boxed(error),
455 }
456 }
457
458 fn scalar_extraction(error: anyhow::Error) -> Self {
459 Self::ScalarExtraction {
460 source: Self::boxed(error),
461 }
462 }
463
464 fn operation(operation: &'static str, error: anyhow::Error) -> Self {
465 Self::Operation {
466 operation,
467 source: Self::boxed(error),
468 }
469 }
470}
471
472#[derive(Clone)]
473pub(crate) enum IdxTensorStorage {
474 Materialized(Arc<Storage>),
475 Eager {
476 inner: Arc<EagerTensor>,
477 axis_classes: Vec<usize>,
478 },
479 Compact(Arc<StructuredPayload>),
481 Deferred {
484 source: Box<Self>,
485 error: Arc<TensorStorageError>,
486 },
487}
488
489impl IdxTensorStorage {
490 fn from_storage(storage: Arc<Storage>) -> Self {
491 Self::Materialized(storage)
492 }
493
494 fn from_eager_dense(inner: EagerTensor, rank: usize) -> Self {
495 Self::Eager {
496 inner: Arc::new(inner),
497 axis_classes: IdxTensor::dense_axis_classes(rank),
498 }
499 }
500
501 fn eager(&self) -> Option<&EagerTensor> {
502 match self {
503 Self::Materialized(_) => None,
504 Self::Eager { inner, .. } => Some(inner.as_ref()),
505 Self::Compact(payload) => Some(payload.payload.as_ref()),
506 Self::Deferred { source, .. } => source.eager(),
507 }
508 }
509
510 fn deferred_error(&self) -> Option<&TensorStorageError> {
511 match self {
512 Self::Deferred { error, .. } => Some(error.as_ref()),
513 _ => None,
514 }
515 }
516
517 fn with_deferred_error(self, error: TensorStorageError) -> Self {
518 if self.deferred_error().is_some() {
519 self
520 } else {
521 Self::Deferred {
522 source: Box::new(self),
523 error: Arc::new(error),
524 }
525 }
526 }
527
528 fn axis_classes(&self) -> &[usize] {
529 match self {
530 Self::Materialized(storage) => storage.axis_classes(),
531 Self::Eager { axis_classes, .. } => axis_classes,
532 Self::Compact(payload) => &payload.axis_classes,
533 Self::Deferred { source, .. } => source.axis_classes(),
534 }
535 }
536
537 fn payload_dims(&self) -> &[usize] {
538 match self {
539 Self::Materialized(storage) => storage.payload_dims(),
540 Self::Eager { inner, .. } => inner.shape(),
541 Self::Compact(payload) => &payload.payload_dims,
542 Self::Deferred { source, .. } => source.payload_dims(),
543 }
544 }
545
546 fn payload_strides_vec(&self) -> Vec<isize> {
547 match self {
548 Self::Materialized(storage) => storage.payload_strides().to_vec(),
549 Self::Eager { inner, .. } => {
550 IdxTensor::col_major_strides(inner.shape()).unwrap_or_default()
551 }
552 Self::Compact(payload) => {
553 IdxTensor::col_major_strides(&payload.payload_dims).unwrap_or_default()
554 }
555 Self::Deferred { source, .. } => source.payload_strides_vec(),
556 }
557 }
558
559 fn is_f64(&self) -> bool {
560 match self {
561 Self::Materialized(storage) => storage.is_f64(),
562 Self::Eager { inner, .. } => inner.dtype() == DType::F64,
563 Self::Compact(payload) => payload.payload.dtype() == DType::F64,
564 Self::Deferred { source, .. } => source.is_f64(),
565 }
566 }
567
568 fn is_c64(&self) -> bool {
569 match self {
570 Self::Materialized(storage) => storage.is_c64(),
571 Self::Eager { inner, .. } => inner.dtype() == DType::C64,
572 Self::Compact(payload) => payload.payload.dtype() == DType::C64,
573 Self::Deferred { source, .. } => source.is_c64(),
574 }
575 }
576
577 fn dtype(&self) -> Option<DType> {
578 match self {
579 Self::Materialized(storage) => Some(if storage.is_c64() {
580 DType::C64
581 } else {
582 DType::F64
583 }),
584 Self::Eager { inner, .. } => Some(inner.dtype()),
585 Self::Compact(payload) => Some(payload.payload.dtype()),
586 Self::Deferred { source, .. } => source.dtype(),
587 }
588 }
589
590 fn is_complex(&self) -> bool {
591 match self {
592 Self::Materialized(storage) => storage.is_complex(),
593 Self::Eager { inner, .. } => matches!(inner.dtype(), DType::C32 | DType::C64),
594 Self::Compact(payload) => {
595 matches!(payload.payload.dtype(), DType::C32 | DType::C64)
596 }
597 Self::Deferred { source, .. } => source.is_complex(),
598 }
599 }
600
601 fn is_diag(&self) -> bool {
602 match self {
603 Self::Materialized(storage) => storage.is_diag(),
604 Self::Eager { axis_classes, .. } => IdxTensor::is_diag_axis_classes(axis_classes),
605 Self::Compact(payload) => IdxTensor::is_diag_axis_classes(&payload.axis_classes),
606 Self::Deferred { source, .. } => source.is_diag(),
607 }
608 }
609
610 fn storage_kind(&self) -> StorageKind {
611 match self {
612 Self::Materialized(storage) => storage.storage_kind(),
613 Self::Eager { axis_classes, .. } => {
614 if axis_classes.iter().copied().eq(0..axis_classes.len()) {
615 StorageKind::Dense
616 } else if IdxTensor::is_diag_axis_classes(axis_classes) {
617 StorageKind::Diagonal
618 } else {
619 StorageKind::Structured
620 }
621 }
622 Self::Compact(payload) => {
623 if payload
624 .axis_classes
625 .iter()
626 .copied()
627 .eq(0..payload.axis_classes.len())
628 {
629 StorageKind::Dense
630 } else if IdxTensor::is_diag_axis_classes(&payload.axis_classes) {
631 StorageKind::Diagonal
632 } else {
633 StorageKind::Structured
634 }
635 }
636 Self::Deferred { source, .. } => source.storage_kind(),
637 }
638 }
639
640 fn materialize_eager_payload(inner: &EagerTensor) -> Result<NativeTensor> {
641 let native = inner.duplicate_value()?;
642 if native.is_col_major_contiguous()? {
643 return Ok(native);
644 }
645 let read = inner.tensor_read();
646 Ok(inner.runtime().with_execution_session(|session| {
647 session.to_contiguous_read(TensorRead::from_view(read.tensor_view()))
648 })??)
649 }
650
651 fn materialize(
652 &self,
653 logical_rank: usize,
654 ) -> std::result::Result<Arc<Storage>, TensorStorageError> {
655 match self {
656 Self::Materialized(storage) => Ok(Arc::clone(storage)),
657 Self::Eager {
658 inner,
659 axis_classes,
660 } => {
661 let native = Self::materialize_eager_payload(inner).map_err(|source| {
662 TensorStorageError::Materialization {
663 source: Arc::from(source.into_boxed_dyn_error()),
664 }
665 })?;
666 let dtype = native.dtype();
667 if matches!(dtype, DType::F32 | DType::C32) {
668 return Err(TensorStorageError::UnsupportedDtype {
669 dtype: IdxTensor::dtype_name(dtype),
670 });
671 }
672 IdxTensor::storage_from_native_with_axis_classes(
673 &native,
674 axis_classes,
675 logical_rank,
676 )
677 .map(Arc::new)
678 .map_err(|source| TensorStorageError::Materialization {
679 source: Arc::from(source.into_boxed_dyn_error()),
680 })
681 }
682 Self::Compact(payload) => {
683 let native =
684 Self::materialize_eager_payload(&payload.payload).map_err(|source| {
685 TensorStorageError::Materialization {
686 source: Arc::from(source.into_boxed_dyn_error()),
687 }
688 })?;
689 let dtype = native.dtype();
690 if matches!(dtype, DType::F32 | DType::C32) {
691 return Err(TensorStorageError::UnsupportedDtype {
692 dtype: IdxTensor::dtype_name(dtype),
693 });
694 }
695 IdxTensor::storage_from_native_with_axis_classes(
696 &native,
697 &payload.axis_classes,
698 logical_rank,
699 )
700 .map(Arc::new)
701 .map_err(|source| TensorStorageError::Materialization {
702 source: Arc::from(source.into_boxed_dyn_error()),
703 })
704 }
705 Self::Deferred { error, .. } => Err((**error).clone()),
706 }
707 }
708
709 fn scale_eager_payload(&self, scalar: &AnyScalar) -> Result<Self> {
710 let (payload, payload_dims, axis_classes) = match self {
711 Self::Materialized(storage) => {
712 let native = if storage.is_f64() {
713 let values = storage
714 .payload_f64_col_major_vec()
715 .map_err(anyhow::Error::new)?;
716 dense_native_tensor_from_col_major(&values, storage.payload_dims())?
717 } else {
718 let values = storage
719 .payload_c64_col_major_vec()
720 .map_err(anyhow::Error::new)?;
721 dense_native_tensor_from_col_major(&values, storage.payload_dims())?
722 };
723 (
724 EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
725 storage.payload_dims().to_vec(),
726 storage.axis_classes().to_vec(),
727 )
728 }
729 Self::Eager {
730 inner,
731 axis_classes,
732 } => (
733 (**inner).clone(),
734 inner.shape().to_vec(),
735 axis_classes.clone(),
736 ),
737 Self::Compact(compact) => (
738 (*compact.payload).clone(),
739 compact.payload_dims.clone(),
740 compact.axis_classes.clone(),
741 ),
742 Self::Deferred { error, .. } => return Err(anyhow::Error::new((**error).clone())),
743 };
744 let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
745 let target_dtype = IdxTensor::scale_target_dtype(payload.dtype(), scalar_inner.dtype())?;
746 let payload = if payload.dtype() == target_dtype {
747 payload
748 } else {
749 payload.cast(target_dtype)?
750 };
751 let scalar_inner = if scalar_inner.dtype() == target_dtype {
752 scalar_inner.clone()
753 } else {
754 scalar_inner.cast(target_dtype)?
755 };
756 let scaled = if payload.shape().is_empty() {
757 payload.mul(&scalar_inner)?
758 } else {
759 let subscripts = IdxTensor::scale_subscripts(payload.shape().len())?;
760 [&payload, &scalar_inner].einsum_subscripts(&subscripts)?
761 };
762 match self {
763 Self::Eager { .. } => Ok(Self::Eager {
764 inner: Arc::new(scaled),
765 axis_classes,
766 }),
767 Self::Compact(_) | Self::Materialized(_) => {
768 Ok(Self::Compact(Arc::new(StructuredPayload {
769 payload: Arc::new(scaled),
770 payload_dims,
771 axis_classes,
772 })))
773 }
774 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
775 }
776 }
777
778 fn conjugate_with<F>(&self, conjugate: &F) -> std::result::Result<Self, TensorStorageError>
779 where
780 F: Fn(
781 &EagerTensor,
782 ) -> std::result::Result<
783 EagerTensor,
784 Arc<dyn std::error::Error + Send + Sync + 'static>,
785 >,
786 {
787 match self {
788 Self::Materialized(storage) => Ok(Self::Materialized(Arc::new(storage.conj()))),
789 Self::Eager {
790 inner,
791 axis_classes,
792 } => conjugate(inner)
793 .map(|conjugated| Self::Eager {
794 inner: Arc::new(conjugated),
795 axis_classes: axis_classes.clone(),
796 })
797 .map_err(|source| TensorStorageError::Conjugation { source }),
798 Self::Compact(payload) => conjugate(payload.payload.as_ref())
799 .map(|conjugated| {
800 Self::Compact(Arc::new(StructuredPayload {
801 payload: Arc::new(conjugated),
802 payload_dims: payload.payload_dims.clone(),
803 axis_classes: payload.axis_classes.clone(),
804 }))
805 })
806 .map_err(|source| TensorStorageError::Conjugation { source }),
807 Self::Deferred { error, .. } => Err((**error).clone()),
808 }
809 }
810
811 fn sum_scalar(&self) -> Result<AnyScalar> {
812 match self {
813 Self::Materialized(storage) => {
814 if storage.is_f64() {
815 Ok(AnyScalar::new_real(storage.sum::<f64>()))
816 } else {
817 let value = storage.sum::<Complex64>();
818 Ok(AnyScalar::new_complex(value.re, value.im))
819 }
820 }
821 Self::Eager { inner, .. } => IdxTensor::native_sum_scalar(inner),
822 Self::Compact(payload) => IdxTensor::native_sum_scalar(&payload.payload),
823 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
824 }
825 }
826
827 fn nonfinite_flags(&self) -> Result<(bool, bool)> {
828 match self {
829 Self::Materialized(storage) => Ok(storage.payload_nonfinite_flags()),
830 Self::Eager { inner, .. } => IdxTensor::native_nonfinite_flags(inner),
831 Self::Compact(payload) => IdxTensor::native_nonfinite_flags(&payload.payload),
832 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
833 }
834 }
835
836 fn payload_value_at(&self, payload_coords: &[usize]) -> Result<Complex64> {
837 match self {
838 Self::Materialized(storage) => storage
839 .scalar_at(payload_coords)
840 .map(Complex64::from)
841 .map_err(anyhow::Error::new),
842 Self::Eager { inner, .. } => {
843 IdxTensor::native_complex_payload_value_at(inner, payload_coords)
844 }
845 Self::Compact(payload) => {
846 IdxTensor::native_complex_payload_value_at(&payload.payload, payload_coords)
847 }
848 Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
849 }
850 }
851
852 fn for_each_payload_value(&self, mut f: impl FnMut(Complex64)) -> Result<()> {
853 let payload_dims = self.payload_dims();
854 let payload_len = checked_product(payload_dims)?;
855 let mut payload_coords = vec![0usize; payload_dims.len()];
856 for _ in 0..payload_len {
857 f(self.payload_value_at(&payload_coords)?);
858 let mut carry = true;
859 for (coordinate, &dim) in payload_coords.iter_mut().zip(payload_dims.iter()) {
860 if !carry {
861 break;
862 }
863 *coordinate += 1;
864 if *coordinate == dim {
865 *coordinate = 0;
866 } else {
867 carry = false;
868 }
869 }
870 }
871 Ok(())
872 }
873
874 fn compact_payload(&self) -> Option<&StructuredPayload> {
875 match self {
876 Self::Compact(payload) => Some(payload.as_ref()),
877 Self::Deferred { source, .. } => source.compact_payload(),
878 _ => None,
879 }
880 }
881}
882
883#[derive(Debug, thiserror::Error)]
898pub enum StructuredSelectorError {
899 #[error("copy-selector bond dimensions differ: left={left}, right={right}")]
901 BondDimensionMismatch {
902 left: usize,
904 right: usize,
906 },
907 #[error("copy-selector {axis} dimension must be positive")]
909 ZeroDimension {
910 axis: &'static str,
912 },
913 #[error("selected site value {value} is outside 0..{site_dim}")]
915 SelectedValueOutOfBounds {
916 value: usize,
918 site_dim: usize,
920 },
921 #[error("copy-selector payload size overflows usize for dimensions {bond_dim} x {site_dim}")]
923 PayloadSizeOverflow {
924 bond_dim: usize,
926 site_dim: usize,
928 },
929 #[error("copy-selector bond stride {bond_dim} exceeds isize::MAX")]
931 StrideOverflow {
932 bond_dim: usize,
934 },
935 #[error("could not allocate copy-selector payload with {elements} elements")]
937 AllocationFailed {
938 elements: usize,
940 },
941 #[error("invalid copy-selector storage: {message}")]
943 InvalidStorage {
944 message: String,
946 },
947}
948
949#[derive(Clone)]
993pub struct IdxTensor {
994 pub indices: Vec<DynIndex>,
996 pub(crate) storage: IdxTensorStorage,
1000 pub(crate) eager_cache: Arc<OnceLock<Arc<EagerTensor>>>,
1002}
1003
1004impl IdxTensor {
1005 fn dense_axis_classes(rank: usize) -> Vec<usize> {
1006 (0..rank).collect()
1007 }
1008
1009 fn dtype_name(dtype: DType) -> &'static str {
1010 match dtype {
1011 DType::F32 => "f32",
1012 DType::F64 => "f64",
1013 DType::C32 => "c32",
1014 DType::C64 => "c64",
1015 DType::I32 => "i32",
1016 DType::I64 => "i64",
1017 DType::Bool => "bool",
1018 }
1019 }
1020
1021 fn scalar_dtype(&self) -> Result<DType> {
1022 if let Some(inner) = self.storage.eager() {
1023 return Ok(inner.dtype());
1024 }
1025 if self.storage.is_f64() {
1026 Ok(DType::F64)
1027 } else if self.storage.is_c64() {
1028 Ok(DType::C64)
1029 } else {
1030 Err(anyhow::anyhow!(
1031 "unable to determine IdxTensor scalar dtype"
1032 ))
1033 }
1034 }
1035
1036 fn diag_axis_classes(rank: usize) -> Vec<usize> {
1037 if rank == 0 {
1038 vec![]
1039 } else {
1040 vec![0; rank]
1041 }
1042 }
1043
1044 fn canonicalize_axis_classes(axis_classes: &[usize]) -> Vec<usize> {
1045 let mut map = std::collections::HashMap::new();
1046 let mut next = 0usize;
1047 axis_classes
1048 .iter()
1049 .map(|&class_id| {
1050 *map.entry(class_id).or_insert_with(|| {
1051 let canonical = next;
1052 next += 1;
1053 canonical
1054 })
1055 })
1056 .collect()
1057 }
1058
1059 fn permute_axis_classes(&self, perm: &[usize]) -> Vec<usize> {
1060 let axis_classes = self.storage.axis_classes();
1061 let permuted: Vec<usize> = perm.iter().map(|&index| axis_classes[index]).collect();
1062 Self::canonicalize_axis_classes(&permuted)
1063 }
1064
1065 fn normalize_insert_axis(op: &str, axis: isize, rank: usize) -> Result<usize> {
1066 let normalized = if axis < 0 {
1067 rank as isize + 1 + axis
1068 } else {
1069 axis
1070 };
1071 if !(normalized >= 0 && normalized <= rank as isize) {
1072 return Err(anyhow::anyhow!(
1073 "{op}: axis {axis} is out of bounds for inserting into rank {rank}"
1074 ));
1075 };
1076 Ok(normalized as usize)
1077 }
1078
1079 fn is_diag_axis_classes(axis_classes: &[usize]) -> bool {
1080 axis_classes.len() >= 2 && axis_classes.iter().all(|&class_id| class_id == 0)
1081 }
1082
1083 fn validate_axis_classes(axis_classes: &[usize], rank: usize) -> Result<()> {
1084 if axis_classes.len() != rank {
1085 return Err(anyhow::anyhow!(
1086 "axis-class rank {} does not match tensor rank {rank}",
1087 axis_classes.len()
1088 ));
1089 }
1090 if Self::canonicalize_axis_classes(axis_classes) != axis_classes {
1091 return Err(anyhow::anyhow!(
1092 "axis classes must be canonical first-occurrence labels: {axis_classes:?}"
1093 ));
1094 }
1095 Ok(())
1096 }
1097
1098 fn einsum_subscripts_from_usize_ids(
1099 inputs: &[Vec<usize>],
1100 output: &[usize],
1101 ) -> Result<EinsumSubscripts> {
1102 let input_labels = inputs
1103 .iter()
1104 .map(|ids| {
1105 ids.iter()
1106 .map(|&id| {
1107 u32::try_from(id)
1108 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1109 })
1110 .collect::<Result<Vec<_>>>()
1111 })
1112 .collect::<Result<Vec<_>>>()?;
1113 let output_labels = output
1114 .iter()
1115 .map(|&id| {
1116 u32::try_from(id)
1117 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1118 })
1119 .collect::<Result<Vec<_>>>()?;
1120 let input_refs = input_labels.iter().map(Vec::as_slice).collect::<Vec<_>>();
1121 Ok(EinsumSubscripts::new(&input_refs, &output_labels))
1122 }
1123
1124 fn build_binary_einsum_subscripts(
1125 lhs_rank: usize,
1126 axes_a: &[usize],
1127 rhs_rank: usize,
1128 axes_b: &[usize],
1129 ) -> Result<EinsumSubscripts> {
1130 if !(axes_a.len() == axes_b.len()) {
1131 return Err(anyhow::anyhow!(
1132 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1133 axes_a,
1134 axes_b
1135 ));
1136 };
1137
1138 let mut lhs_ids = vec![usize::MAX; lhs_rank];
1139 let mut rhs_ids = vec![usize::MAX; rhs_rank];
1140 let mut next_id = 0usize;
1141
1142 let mut seen_lhs = vec![false; lhs_rank];
1143 let mut seen_rhs = vec![false; rhs_rank];
1144
1145 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1146 if !(lhs_axis < lhs_rank) {
1147 return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1148 };
1149 if !(rhs_axis < rhs_rank) {
1150 return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1151 };
1152 if !(!seen_lhs[lhs_axis]) {
1153 return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1154 };
1155 if !(!seen_rhs[rhs_axis]) {
1156 return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1157 };
1158 seen_lhs[lhs_axis] = true;
1159 seen_rhs[rhs_axis] = true;
1160 lhs_ids[lhs_axis] = next_id;
1161 rhs_ids[rhs_axis] = next_id;
1162 next_id += 1;
1163 }
1164
1165 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1166 for id in &mut lhs_ids {
1167 if *id == usize::MAX {
1168 *id = next_id;
1169 output_ids.push(next_id);
1170 next_id += 1;
1171 }
1172 }
1173 for id in &mut rhs_ids {
1174 if *id == usize::MAX {
1175 *id = next_id;
1176 output_ids.push(next_id);
1177 next_id += 1;
1178 }
1179 }
1180
1181 Self::einsum_subscripts_from_usize_ids(&[lhs_ids, rhs_ids], &output_ids)
1182 }
1183
1184 fn binary_dot_general_config(axes_a: &[usize], axes_b: &[usize]) -> Result<DotGeneralConfig> {
1185 if !(axes_a.len() == axes_b.len()) {
1186 return Err(anyhow::anyhow!(
1187 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1188 axes_a,
1189 axes_b
1190 ));
1191 };
1192 Ok(DotGeneralConfig {
1193 lhs_contracting_dims: axes_a.to_vec(),
1194 rhs_contracting_dims: axes_b.to_vec(),
1195 lhs_batch_dims: vec![],
1196 rhs_batch_dims: vec![],
1197 })
1198 }
1199
1200 fn binary_contraction_axis_classes(
1201 lhs_axis_classes: &[usize],
1202 axes_a: &[usize],
1203 rhs_axis_classes: &[usize],
1204 axes_b: &[usize],
1205 ) -> Result<Vec<usize>> {
1206 debug_assert_eq!(axes_a.len(), axes_b.len());
1207
1208 fn find(parent: &mut [usize], value: usize) -> usize {
1209 if parent[value] != value {
1210 parent[value] = find(parent, parent[value]);
1211 }
1212 parent[value]
1213 }
1214
1215 fn union(parent: &mut [usize], lhs: usize, rhs: usize) {
1216 let lhs_root = find(parent, lhs);
1217 let rhs_root = find(parent, rhs);
1218 if lhs_root != rhs_root {
1219 parent[rhs_root] = lhs_root;
1220 }
1221 }
1222
1223 let lhs_payload_rank = match lhs_axis_classes.iter().copied().max() {
1224 Some(value) => value
1225 .checked_add(1)
1226 .ok_or_else(|| anyhow::anyhow!("left payload rank overflows usize"))?,
1227 None => 0,
1228 };
1229 let rhs_payload_rank = match rhs_axis_classes.iter().copied().max() {
1230 Some(value) => value
1231 .checked_add(1)
1232 .ok_or_else(|| anyhow::anyhow!("right payload rank overflows usize"))?,
1233 None => 0,
1234 };
1235 let rhs_offset = lhs_payload_rank;
1236 let parent_len = lhs_payload_rank
1237 .checked_add(rhs_payload_rank)
1238 .ok_or_else(|| anyhow::anyhow!("payload rank sum overflows usize"))?;
1239 let mut parent: Vec<usize> = (0..parent_len).collect();
1240
1241 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1242 union(
1243 &mut parent,
1244 lhs_axis_classes[lhs_axis],
1245 rhs_offset
1246 .checked_add(rhs_axis_classes[rhs_axis])
1247 .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?,
1248 );
1249 }
1250
1251 let mut lhs_contracted = vec![false; lhs_axis_classes.len()];
1252 for &axis in axes_a {
1253 lhs_contracted[axis] = true;
1254 }
1255 let mut rhs_contracted = vec![false; rhs_axis_classes.len()];
1256 for &axis in axes_b {
1257 rhs_contracted[axis] = true;
1258 }
1259
1260 let mut root_to_class = std::collections::HashMap::new();
1261 let mut next_class = 0usize;
1262 let mut axis_classes = Vec::new();
1263
1264 for (axis, &class_id) in lhs_axis_classes.iter().enumerate() {
1265 if !lhs_contracted[axis] {
1266 let root = find(&mut parent, class_id);
1267 let class = *root_to_class.entry(root).or_insert_with(|| {
1268 let value = next_class;
1269 next_class += 1;
1270 value
1271 });
1272 axis_classes.push(class);
1273 }
1274 }
1275 for (axis, &class_id) in rhs_axis_classes.iter().enumerate() {
1276 if !rhs_contracted[axis] {
1277 let rhs_class = rhs_offset
1278 .checked_add(class_id)
1279 .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?;
1280 let root = find(&mut parent, rhs_class);
1281 let class = *root_to_class.entry(root).or_insert_with(|| {
1282 let value = next_class;
1283 next_class += 1;
1284 value
1285 });
1286 axis_classes.push(class);
1287 }
1288 }
1289
1290 Ok(axis_classes)
1291 }
1292
1293 fn scale_subscripts(rank: usize) -> Result<EinsumSubscripts> {
1294 let ids: Vec<usize> = (0..rank).collect();
1295 Self::einsum_subscripts_from_usize_ids(&[ids.clone(), Vec::new()], &ids)
1296 }
1297
1298 fn scale_target_dtype(payload: DType, scalar: DType) -> Result<DType> {
1299 let target = match payload {
1300 DType::F32 => match scalar {
1301 DType::C32 | DType::C64 => DType::C32,
1302 DType::F32 | DType::F64 => DType::F32,
1303 dtype => {
1304 return Err(anyhow::anyhow!(
1305 "unsupported scalar dtype {dtype:?} for f32 scaling"
1306 ));
1307 }
1308 },
1309 DType::C32 => match scalar {
1310 DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C32,
1311 dtype => {
1312 return Err(anyhow::anyhow!(
1313 "unsupported scalar dtype {dtype:?} for c32 scaling"
1314 ));
1315 }
1316 },
1317 DType::F64 => match scalar {
1318 DType::C32 | DType::C64 => DType::C64,
1319 DType::F32 | DType::F64 => DType::F64,
1320 dtype => {
1321 return Err(anyhow::anyhow!(
1322 "unsupported scalar dtype {dtype:?} for f64 scaling"
1323 ));
1324 }
1325 },
1326 DType::C64 => match scalar {
1327 DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C64,
1328 dtype => {
1329 return Err(anyhow::anyhow!(
1330 "unsupported scalar dtype {dtype:?} for c64 scaling"
1331 ));
1332 }
1333 },
1334 dtype => {
1335 return Err(anyhow::anyhow!(
1336 "unsupported tensor dtype {dtype:?} for scaling"
1337 ));
1338 }
1339 };
1340 Ok(target)
1341 }
1342
1343 fn validate_indices(indices: &[DynIndex]) -> Result<()> {
1344 let mut seen = HashSet::new();
1345 for idx in indices {
1346 if !(seen.insert(idx.clone())) {
1347 return Err(anyhow::anyhow!("Tensor indices must all be unique"));
1348 };
1349 }
1350 Ok(())
1351 }
1352
1353 fn validate_diag_dims(dims: &[usize]) -> Result<()> {
1354 if !dims.is_empty() {
1355 let first_dim = dims[0];
1356 for (i, &dim) in dims.iter().enumerate() {
1357 if !(dim == first_dim) {
1358 return Err(anyhow::anyhow!("DiagTensor requires all indices to have the same dimension, but dims[{i}] = {dim} != dims[0] = {first_dim}"));
1359 };
1360 }
1361 }
1362 Ok(())
1363 }
1364
1365 fn seed_native_payload(storage: &Storage, dims: &[usize]) -> Result<NativeTensor> {
1366 Ok(storage_to_native_tensor(storage, dims)?)
1367 }
1368
1369 fn empty_eager_cache() -> Arc<OnceLock<Arc<EagerTensor>>> {
1370 Arc::new(OnceLock::new())
1371 }
1372
1373 fn eager_cache_with(inner: EagerTensor) -> Arc<OnceLock<Arc<EagerTensor>>> {
1374 let cache = Arc::new(OnceLock::new());
1375 let _ = cache.set(Arc::new(inner));
1376 cache
1377 }
1378
1379 fn compact_payload_inner(&self) -> Result<EagerTensor> {
1380 self.ensure_storage_ready()?;
1381 if let Some(inner) = self.storage.eager() {
1382 return Ok(inner.clone());
1383 }
1384 Ok(EagerTensor::from_tensor_in(
1385 storage_payload_native(self.storage.materialize(self.indices.len())?.as_ref())?,
1386 default_eager_ctx()?,
1387 )?)
1388 }
1389
1390 fn dense_inner_from_payload(
1391 payload: &EagerTensor,
1392 axis_classes: &[usize],
1393 logical_dims: &[usize],
1394 ) -> Result<EagerTensor> {
1395 let payload_rank = match axis_classes.iter().copied().max() {
1396 Some(class_id) => class_id
1397 .checked_add(1)
1398 .ok_or_else(|| anyhow::anyhow!("structured payload class rank overflows usize"))?,
1399 None => 0,
1400 };
1401 if !(payload.shape().len() == payload_rank) {
1402 return Err(anyhow::anyhow!(
1403 "structured payload rank {} does not match axis classes {:?}",
1404 payload.shape().len(),
1405 axis_classes
1406 ));
1407 };
1408 if !(logical_dims.len() == axis_classes.len()) {
1409 return Err(anyhow::anyhow!(
1410 "logical rank {} does not match axis class rank {}",
1411 logical_dims.len(),
1412 axis_classes.len()
1413 ));
1414 };
1415
1416 if axis_classes == Self::dense_axis_classes(logical_dims.len()) {
1417 if !(payload.shape() == logical_dims) {
1418 return Err(anyhow::anyhow!(
1419 "dense payload dims {:?} do not match logical dims {:?}",
1420 payload.shape(),
1421 logical_dims
1422 ));
1423 };
1424 return Ok(payload.clone());
1425 }
1426
1427 let mut first_axis_by_class = vec![None; payload_rank];
1428 let mut dense = payload.clone();
1429 for (logical_axis, &class_id) in axis_classes.iter().enumerate() {
1430 let first_axis = match first_axis_by_class[class_id] {
1431 Some(first_axis) => first_axis,
1432 None => {
1433 first_axis_by_class[class_id] = Some(logical_axis);
1434 continue;
1435 }
1436 };
1437 dense = dense.embed_diag(first_axis, logical_axis)?;
1438 }
1439 if !(dense.shape() == logical_dims) {
1440 return Err(anyhow::anyhow!(
1441 "expanded structured payload dims {:?} do not match logical dims {:?}",
1442 dense.shape(),
1443 logical_dims
1444 ));
1445 };
1446 Ok(dense)
1447 }
1448
1449 fn tracked_compact_payload_value(&self) -> Option<&StructuredPayload> {
1450 self.storage
1451 .deferred_error()
1452 .is_none()
1453 .then_some(self.storage.compact_payload())
1454 .flatten()
1455 .filter(|value| value.payload.tracks_grad())
1456 }
1457
1458 fn ensure_storage_ready(&self) -> Result<()> {
1459 if let Some(error) = self.storage.deferred_error() {
1460 return Err(anyhow::Error::new(error.clone()));
1461 }
1462 Ok(())
1463 }
1464
1465 fn compact_payload_is_logical_dense(&self, payload_dims: &[usize]) -> bool {
1466 self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len())
1467 && payload_dims == self.dims()
1468 }
1469
1470 fn uses_tracked_compact_storage(&self) -> bool {
1471 self.tracked_compact_payload_value()
1472 .is_some_and(|value| !self.compact_payload_is_logical_dense(&value.payload_dims))
1473 }
1474
1475 fn ensure_shape_packing_preserves_ad(&self, op_name: &str) -> Result<()> {
1476 self.ensure_storage_ready()?;
1477 if !(!self.uses_tracked_compact_storage()) {
1478 return Err(anyhow::anyhow!("{op_name}: structured AD tensors with compact storage are not supported because materializing compact storage would detach gradients"));
1479 };
1480 Ok(())
1481 }
1482
1483 fn operand_indices_for_contraction(&self, conjugate: bool) -> Vec<DynIndex> {
1484 if conjugate {
1485 self.indices.iter().map(|index| index.conj()).collect()
1486 } else {
1487 self.indices.clone()
1488 }
1489 }
1490
1491 fn build_binary_contraction_labels(
1492 lhs_rank: usize,
1493 axes_a: &[usize],
1494 rhs_rank: usize,
1495 axes_b: &[usize],
1496 ) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
1497 if !(axes_a.len() == axes_b.len()) {
1498 return Err(anyhow::anyhow!(
1499 "contract axis length mismatch: lhs {:?}, rhs {:?}",
1500 axes_a,
1501 axes_b
1502 ));
1503 };
1504
1505 let mut lhs_ids = vec![usize::MAX; lhs_rank];
1506 let mut rhs_ids = vec![usize::MAX; rhs_rank];
1507 let mut next_id = 0usize;
1508
1509 let mut seen_lhs = vec![false; lhs_rank];
1510 let mut seen_rhs = vec![false; rhs_rank];
1511
1512 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1513 if !(lhs_axis < lhs_rank) {
1514 return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1515 };
1516 if !(rhs_axis < rhs_rank) {
1517 return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1518 };
1519 if !(!seen_lhs[lhs_axis]) {
1520 return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1521 };
1522 if !(!seen_rhs[rhs_axis]) {
1523 return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1524 };
1525 seen_lhs[lhs_axis] = true;
1526 seen_rhs[rhs_axis] = true;
1527 lhs_ids[lhs_axis] = next_id;
1528 rhs_ids[rhs_axis] = next_id;
1529 next_id += 1;
1530 }
1531
1532 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1533 for id in &mut lhs_ids {
1534 if *id == usize::MAX {
1535 *id = next_id;
1536 output_ids.push(next_id);
1537 next_id += 1;
1538 }
1539 }
1540 for id in &mut rhs_ids {
1541 if *id == usize::MAX {
1542 *id = next_id;
1543 output_ids.push(next_id);
1544 next_id += 1;
1545 }
1546 }
1547
1548 Ok((lhs_ids, rhs_ids, output_ids))
1549 }
1550
1551 fn build_payload_einsum_subscripts(
1552 input_roots: &[Vec<usize>],
1553 output_roots: &[usize],
1554 ) -> Result<EinsumSubscripts> {
1555 Self::einsum_subscripts_from_usize_ids(input_roots, output_roots)
1556 }
1557
1558 fn normalize_eager_payload_for_roots(
1559 payload: &EagerTensor,
1560 roots: &[usize],
1561 ) -> Result<(Option<EagerTensor>, Vec<usize>)> {
1562 if !(payload.shape().len() == roots.len()) {
1563 return Err(anyhow::anyhow!(
1564 "payload rank {} does not match root label count {}",
1565 payload.shape().len(),
1566 roots.len()
1567 ));
1568 };
1569
1570 let mut current_payload = None;
1571 let mut current_roots = roots.to_vec();
1572 while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(¤t_roots) {
1573 let source = current_payload.as_ref().unwrap_or(payload);
1574 current_payload = Some(source.extract_diag(axis_a, axis_b)?);
1575 current_roots.remove(axis_b);
1576 }
1577
1578 Ok((current_payload, current_roots))
1579 }
1580
1581 fn first_duplicate_pair(values: &[usize]) -> Option<(usize, usize)> {
1582 let mut first_axis_by_value = std::collections::HashMap::new();
1583 for (axis, &value) in values.iter().enumerate() {
1584 if let Some(&first_axis) = first_axis_by_value.get(&value) {
1585 return Some((first_axis, axis));
1586 }
1587 first_axis_by_value.insert(value, axis);
1588 }
1589 None
1590 }
1591
1592 fn from_structured_payload_inner(
1593 indices: Vec<DynIndex>,
1594 payload_inner: EagerTensor,
1595 payload_dims: Vec<usize>,
1596 axis_classes: Vec<usize>,
1597 ) -> Result<Self> {
1598 Self::validate_indices(&indices)?;
1599 if payload_inner.shape() != payload_dims {
1600 return Err(anyhow::anyhow!(
1601 "structured payload dims {:?} do not match planned payload dims {:?}",
1602 payload_inner.shape(),
1603 payload_dims
1604 ));
1605 }
1606 if axis_classes == Self::dense_axis_classes(indices.len()) {
1607 return Self::from_inner_with_axis_classes(indices, payload_inner, axis_classes);
1608 }
1609 let structured_payload = Arc::new(StructuredPayload {
1610 payload: Arc::new(payload_inner),
1611 payload_dims,
1612 axis_classes,
1613 });
1614 Ok(Self {
1615 indices,
1616 storage: IdxTensorStorage::Compact(structured_payload),
1617 eager_cache: Self::empty_eager_cache(),
1618 })
1619 }
1620
1621 fn contract_structured_payloads(
1622 &self,
1623 other: &Self,
1624 result_indices: Vec<DynIndex>,
1625 axes_a: &[usize],
1626 axes_b: &[usize],
1627 ) -> Result<Self> {
1628 let (lhs_labels, rhs_labels, output_labels) = Self::build_binary_contraction_labels(
1629 self.indices.len(),
1630 axes_a,
1631 other.indices.len(),
1632 axes_b,
1633 )?;
1634 Self::contract_structured_payloads_nary(
1635 &[self, other],
1636 result_indices,
1637 vec![lhs_labels, rhs_labels],
1638 output_labels,
1639 )
1640 }
1641
1642 pub(crate) fn contract_structured_payloads_nary(
1643 operands: &[&Self],
1644 result_indices: Vec<DynIndex>,
1645 input_labels: Vec<Vec<usize>>,
1646 output_labels: Vec<usize>,
1647 ) -> Result<Self> {
1648 if !(!operands.is_empty()) {
1649 return Err(anyhow::anyhow!("structured contraction needs operands"));
1650 };
1651 for operand in operands {
1652 operand.ensure_storage_ready()?;
1653 }
1654 let layouts = operands
1655 .iter()
1656 .map(|operand| {
1657 OperandLayout::new(operand.dims(), operand.storage.axis_classes().to_vec())
1658 })
1659 .collect::<Result<Vec<_>>>()?;
1660 let spec = StructuredContractionSpec {
1661 input_labels,
1662 output_labels,
1663 retained_labels: Default::default(),
1664 };
1665 let plan = StructuredContractionPlan::new(&layouts, &spec)?;
1666 let any_grad = operands.iter().any(|operand| operand.tracks_grad());
1667
1668 if any_grad {
1669 let dtypes = operands
1670 .iter()
1671 .map(|operand| {
1672 operand.storage.dtype().ok_or_else(|| {
1673 anyhow::anyhow!("structured contraction operand has no scalar dtype")
1674 })
1675 })
1676 .collect::<Result<Vec<_>>>()?;
1677 let target = Self::common_eager_dtype(&dtypes)?;
1678 let mut payloads = Vec::with_capacity(operands.len());
1679 for operand in operands {
1680 let payload = operand.compact_payload_inner()?;
1681 payloads.push(if payload.dtype() == target {
1682 payload
1683 } else {
1684 payload.cast(target)?
1685 });
1686 }
1687
1688 let mut normalized = Vec::with_capacity(payloads.len());
1689 let mut labels = Vec::with_capacity(payloads.len());
1690 for (operand_idx, (payload, operand_plan)) in
1691 payloads.iter().zip(plan.operand_plans.iter()).enumerate()
1692 {
1693 let (payload, roots) =
1694 Self::normalize_eager_payload_for_roots(payload, &operand_plan.class_roots)?;
1695 normalized.push(payload.unwrap_or_else(|| payloads[operand_idx].clone()));
1696 labels.push(roots);
1697 }
1698 let refs = normalized.iter().collect::<Vec<_>>();
1699 let subscripts =
1700 Self::build_payload_einsum_subscripts(&labels, &plan.output_payload_roots)?;
1701 let payload = refs.as_slice().einsum_subscripts(&subscripts)?;
1702 return Self::from_structured_payload_inner(
1703 result_indices,
1704 payload,
1705 plan.output_payload_dims,
1706 plan.output_axis_classes,
1707 );
1708 }
1709
1710 let storage_owners = operands
1714 .iter()
1715 .map(|operand| match &operand.storage {
1716 IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1717 IdxTensorStorage::Deferred { source, .. } => match source.as_ref() {
1718 IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1719 _ => None,
1720 },
1721 _ => None,
1722 })
1723 .collect::<Vec<_>>();
1724 let mut inputs = Vec::with_capacity(operands.len());
1725 for (operand_idx, operand) in operands.iter().enumerate() {
1726 if let Some(storage) = storage_owners[operand_idx].as_ref() {
1727 inputs.push(storage_payload_native_read_input(storage.as_ref())?);
1728 } else {
1729 let inner = operand
1730 .storage
1731 .eager()
1732 .ok_or_else(|| anyhow::anyhow!("structured operand has no compact payload"))?;
1733 inputs.push(NativeTensorReadInput::Borrowed(inner.tensor_read()));
1734 }
1735 }
1736 let mut normalized = Vec::with_capacity(inputs.len());
1737 let mut labels = Vec::with_capacity(inputs.len());
1738 for (input, operand_plan) in inputs.into_iter().zip(plan.operand_plans.iter()) {
1739 let (input, roots) =
1740 normalize_payload_read_for_roots(input, &operand_plan.class_roots)?;
1741 normalized.push(input);
1742 labels.push(roots);
1743 }
1744 let refs = normalized
1745 .iter()
1746 .zip(labels.iter())
1747 .map(|(input, labels)| (input, labels.as_slice()))
1748 .collect::<Vec<_>>();
1749 let payload = tensor4all_tensorbackend::einsum_native_tensor_reads(
1750 &refs,
1751 &plan.output_payload_roots,
1752 )?;
1753 let payload_inner = EagerTensor::from_tensor_in(payload, default_eager_ctx()?)?;
1754 Self::from_structured_payload_inner(
1755 result_indices,
1756 payload_inner,
1757 plan.output_payload_dims,
1758 plan.output_axis_classes,
1759 )
1760 }
1761
1762 fn common_eager_dtype(dtypes: &[DType]) -> Result<DType> {
1763 let target = if dtypes.contains(&DType::C64)
1764 || (dtypes.contains(&DType::C32) && dtypes.contains(&DType::F64))
1765 {
1766 DType::C64
1767 } else if dtypes.contains(&DType::F64) || dtypes.contains(&DType::C32) {
1768 if dtypes.contains(&DType::C32) {
1769 DType::C32
1770 } else {
1771 DType::F64
1772 }
1773 } else {
1774 DType::F32
1775 };
1776 if !(dtypes
1777 .iter()
1778 .all(|dtype| matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64)))
1779 {
1780 return Err(anyhow::anyhow!(
1781 "structured contraction supports only f32, f64, c32, and c64 operands"
1782 ));
1783 };
1784 Ok(target)
1785 }
1786
1787 fn should_use_structured_payload_contract(&self, other: &Self) -> bool {
1788 self.tracks_grad()
1789 || other.tracks_grad()
1790 || self.storage.axis_classes() != Self::dense_axis_classes(self.indices.len())
1791 || other.storage.axis_classes() != Self::dense_axis_classes(other.indices.len())
1792 }
1793
1794 fn storage_from_native_with_axis_classes(
1795 native: &NativeTensor,
1796 axis_classes: &[usize],
1797 logical_rank: usize,
1798 ) -> Result<Storage> {
1799 if matches!(native.dtype(), DType::F32 | DType::C32) {
1800 return Err(anyhow::anyhow!(
1801 "compact IdxTensor storage does not support dtype {:?}; retain the eager payload",
1802 native.dtype()
1803 ));
1804 }
1805 if Self::is_diag_axis_classes(axis_classes) {
1806 match native.dtype() {
1807 DType::F64 | DType::I32 | DType::I64 | DType::Bool => Storage::from_diag_col_major(
1808 native_tensor_primal_to_diag::<f64>(native)?,
1809 logical_rank,
1810 ),
1811 DType::C64 => Storage::from_diag_col_major(
1812 native_tensor_primal_to_diag::<Complex64>(native)?,
1813 logical_rank,
1814 ),
1815 DType::F32 | DType::C32 => Err(anyhow::anyhow!(
1816 "compact IdxTensor storage does not support dtype {:?}",
1817 native.dtype()
1818 )),
1819 }
1820 } else {
1821 storage_from_payload_native(native.duplicate()?, native.shape(), axis_classes.to_vec())
1822 }
1823 }
1824
1825 fn dense_selected_diag_payload<T: TensorElement + Copy + Zero>(
1826 payload: Vec<T>,
1827 kept_dims: &[usize],
1828 selected_positions: &[usize],
1829 ) -> Result<Vec<T>> {
1830 let output_len = checked_product(kept_dims)?;
1831 let mut data = vec![T::zero(); output_len];
1832 if output_len == 0 {
1833 return Ok(data);
1834 }
1835
1836 let Some((&first_position, rest)) = selected_positions.split_first() else {
1837 return Ok(data);
1838 };
1839 if rest.iter().any(|&position| position != first_position) {
1840 return Ok(data);
1841 }
1842
1843 let value = payload[first_position];
1844 if kept_dims.is_empty() {
1845 data[0] = value;
1846 return Ok(data);
1847 }
1848
1849 let mut offset = 0usize;
1850 let mut stride = 1usize;
1851 for &dim in kept_dims {
1852 let term = first_position
1853 .checked_mul(stride)
1854 .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1855 offset = offset
1856 .checked_add(term)
1857 .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1858 stride = stride
1859 .checked_mul(dim)
1860 .ok_or_else(|| anyhow::anyhow!("diagonal selection stride overflow"))?;
1861 }
1862 data[offset] = value;
1863 Ok(data)
1864 }
1865
1866 fn select_diag_indices(
1867 &self,
1868 kept_indices: Vec<DynIndex>,
1869 kept_dims: Vec<usize>,
1870 positions: &[usize],
1871 ) -> Result<Self> {
1872 if self.storage.is_f64() {
1873 let storage = self.storage.materialize(self.indices.len())?;
1874 let payload = storage
1875 .payload_f64_col_major_vec()
1876 .map_err(anyhow::Error::new)?;
1877 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1878 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1879 } else if self.storage.is_c64() {
1880 let storage = self.storage.materialize(self.indices.len())?;
1881 let payload = storage
1882 .payload_c64_col_major_vec()
1883 .map_err(anyhow::Error::new)?;
1884 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1885 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1886 } else if self.storage.dtype() == Some(DType::F32) {
1887 let inner = self
1888 .storage
1889 .eager()
1890 .ok_or_else(|| anyhow::anyhow!("failed to read f32 diagonal payload"))?;
1891 let payload = inner.value()?.as_slice::<f32>()?.to_vec();
1892 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1893 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1894 } else if self.storage.dtype() == Some(DType::C32) {
1895 let inner = self
1896 .storage
1897 .eager()
1898 .ok_or_else(|| anyhow::anyhow!("failed to read c32 diagonal payload"))?;
1899 let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
1900 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1901 Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1902 } else {
1903 Err(anyhow::anyhow!("unsupported diagonal storage scalar type"))
1904 }
1905 }
1906
1907 fn col_major_strides(dims: &[usize]) -> Result<Vec<isize>> {
1908 let mut strides = Vec::with_capacity(dims.len());
1909 let mut stride = 1isize;
1910 for &dim in dims {
1911 strides.push(stride);
1912 let dim = isize::try_from(dim)
1913 .map_err(|_| anyhow::anyhow!("dimension does not fit in isize"))?;
1914 stride = stride
1915 .checked_mul(dim)
1916 .ok_or_else(|| anyhow::anyhow!("column-major stride overflow"))?;
1917 }
1918 Ok(strides)
1919 }
1920
1921 fn zero_structured_selection<T>(
1922 kept_indices: Vec<DynIndex>,
1923 kept_dims: &[usize],
1924 ) -> Result<Self>
1925 where
1926 T: TensorElement + Zero,
1927 {
1928 let output_len = checked_product(kept_dims)?;
1929 Self::from_dense(kept_indices, vec![T::zero(); output_len]).map_err(anyhow::Error::from)
1930 }
1931
1932 fn selected_structured_class_positions(
1933 axis_classes: &[usize],
1934 payload_rank: usize,
1935 selected_axes: &[usize],
1936 positions: &[usize],
1937 ) -> Option<Vec<Option<usize>>> {
1938 let mut selected_class_positions = vec![None; payload_rank];
1939 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
1940 let class_id = axis_classes[axis];
1941 if let Some(existing) = selected_class_positions[class_id] {
1942 if existing != position {
1943 return None;
1944 }
1945 } else {
1946 selected_class_positions[class_id] = Some(position);
1947 }
1948 }
1949 Some(selected_class_positions)
1950 }
1951
1952 fn select_structured_indices_typed<T, F>(
1953 &self,
1954 payload: Vec<T>,
1955 kept_axes: &[usize],
1956 kept_indices: Vec<DynIndex>,
1957 kept_dims: Vec<usize>,
1958 selected: (&[usize], &[usize]),
1959 make_output: F,
1960 ) -> Result<Self>
1961 where
1962 T: TensorElement + Zero,
1963 F: FnOnce(Vec<T>, Vec<usize>, Vec<isize>, Vec<usize>) -> Result<Self>,
1964 {
1965 let (selected_axes, positions) = selected;
1966 let payload_dims = self.storage.payload_dims();
1967 let axis_classes = self.storage.axis_classes();
1968 let payload_rank = payload_dims.len();
1969 let Some(selected_class_positions) = Self::selected_structured_class_positions(
1970 axis_classes,
1971 payload_rank,
1972 selected_axes,
1973 positions,
1974 ) else {
1975 return Self::zero_structured_selection::<T>(kept_indices, &kept_dims);
1976 };
1977
1978 let selected_class_kept = kept_axes
1979 .iter()
1980 .any(|&axis| selected_class_positions[axis_classes[axis]].is_some());
1981 if selected_class_kept {
1982 return self.select_structured_indices_dense(
1983 payload,
1984 kept_axes,
1985 kept_indices,
1986 kept_dims,
1987 &selected_class_positions,
1988 );
1989 }
1990
1991 let mut old_to_new_class = vec![None; payload_rank];
1992 let mut output_payload_dims = Vec::new();
1993 let mut output_axis_classes = Vec::with_capacity(kept_axes.len());
1994 for &axis in kept_axes {
1995 let class_id = axis_classes[axis];
1996 let new_class = match old_to_new_class[class_id] {
1997 Some(new_class) => new_class,
1998 None => {
1999 let new_class = output_payload_dims.len();
2000 old_to_new_class[class_id] = Some(new_class);
2001 output_payload_dims.push(payload_dims[class_id]);
2002 new_class
2003 }
2004 };
2005 output_axis_classes.push(new_class);
2006 }
2007
2008 let output_len = checked_product(&output_payload_dims)?;
2009 let mut output_payload = Vec::with_capacity(output_len);
2010 for linear in 0..output_len {
2011 let output_payload_index = decode_col_major_linear(linear, &output_payload_dims)?;
2012 let mut input_payload_index = vec![0usize; payload_rank];
2013 for class_id in 0..payload_rank {
2014 input_payload_index[class_id] =
2015 if let Some(position) = selected_class_positions[class_id] {
2016 position
2017 } else if let Some(new_class) = old_to_new_class[class_id] {
2018 output_payload_index[new_class]
2019 } else {
2020 return Err(anyhow::anyhow!(
2021 "structured payload class {class_id} is neither selected nor kept"
2022 ));
2023 };
2024 }
2025 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2026 output_payload.push(payload[input_linear]);
2027 }
2028
2029 let output_strides = Self::col_major_strides(&output_payload_dims)?;
2030 make_output(
2031 output_payload,
2032 output_payload_dims,
2033 output_strides,
2034 output_axis_classes,
2035 )
2036 }
2037
2038 fn select_structured_indices_dense<T>(
2039 &self,
2040 payload: Vec<T>,
2041 kept_axes: &[usize],
2042 kept_indices: Vec<DynIndex>,
2043 kept_dims: Vec<usize>,
2044 selected_class_positions: &[Option<usize>],
2045 ) -> Result<Self>
2046 where
2047 T: TensorElement + Zero,
2048 {
2049 let payload_dims = self.storage.payload_dims();
2050 let axis_classes = self.storage.axis_classes();
2051 let output_len = checked_product(&kept_dims)?;
2052 let mut output = Vec::with_capacity(output_len);
2053
2054 for linear in 0..output_len {
2055 let kept_position = decode_col_major_linear(linear, &kept_dims)?;
2056 let mut input_payload_index = selected_class_positions.to_vec();
2057 let mut is_structural_zero = false;
2058
2059 for (&axis, &position) in kept_axes.iter().zip(kept_position.iter()) {
2060 let class_id = axis_classes[axis];
2061 match input_payload_index[class_id] {
2062 Some(existing) if existing != position => {
2063 is_structural_zero = true;
2064 break;
2065 }
2066 Some(_) => {}
2067 None => input_payload_index[class_id] = Some(position),
2068 }
2069 }
2070
2071 if is_structural_zero {
2072 output.push(T::zero());
2073 continue;
2074 }
2075
2076 let input_payload_index = input_payload_index
2077 .into_iter()
2078 .enumerate()
2079 .map(|(class_id, position)| {
2080 position.ok_or_else(|| {
2081 anyhow::anyhow!(
2082 "structured payload class {class_id} is neither selected nor kept"
2083 )
2084 })
2085 })
2086 .collect::<Result<Vec<_>>>()?;
2087 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2088 output.push(payload[input_linear]);
2089 }
2090
2091 Self::from_dense(kept_indices, output).map_err(anyhow::Error::from)
2092 }
2093
2094 fn select_structured_indices(
2095 &self,
2096 kept_axes: &[usize],
2097 kept_indices: Vec<DynIndex>,
2098 kept_dims: Vec<usize>,
2099 selected_axes: &[usize],
2100 positions: &[usize],
2101 ) -> Result<Self> {
2102 if self.storage.is_f64() {
2103 let storage = self.storage.materialize(self.indices.len())?;
2104 let payload = storage
2105 .payload_f64_col_major_vec()
2106 .map_err(anyhow::Error::new)?;
2107 let output_indices = kept_indices.clone();
2108 self.select_structured_indices_typed(
2109 payload,
2110 kept_axes,
2111 kept_indices,
2112 kept_dims,
2113 (selected_axes, positions),
2114 move |payload, dims, strides, classes| {
2115 let storage = Storage::new_structured(payload, dims, strides, classes)?;
2116 Self::from_storage(output_indices, Arc::new(storage))
2117 .map_err(anyhow::Error::from)
2118 },
2119 )
2120 } else if self.storage.is_c64() {
2121 let storage = self.storage.materialize(self.indices.len())?;
2122 let payload = storage
2123 .payload_c64_col_major_vec()
2124 .map_err(anyhow::Error::new)?;
2125 let output_indices = kept_indices.clone();
2126 self.select_structured_indices_typed(
2127 payload,
2128 kept_axes,
2129 kept_indices,
2130 kept_dims,
2131 (selected_axes, positions),
2132 move |payload, dims, strides, classes| {
2133 let storage = Storage::new_structured(payload, dims, strides, classes)?;
2134 Self::from_storage(output_indices, Arc::new(storage))
2135 .map_err(anyhow::Error::from)
2136 },
2137 )
2138 } else if self.storage.dtype() == Some(DType::F32) {
2139 let inner = self
2140 .storage
2141 .eager()
2142 .ok_or_else(|| anyhow::anyhow!("failed to read f32 structured payload"))?;
2143 let payload = inner.value()?.as_slice::<f32>()?.to_vec();
2144 let output_indices = kept_indices.clone();
2145 self.select_structured_indices_typed(
2146 payload,
2147 kept_axes,
2148 kept_indices,
2149 kept_dims,
2150 (selected_axes, positions),
2151 move |payload, dims, _strides, classes| {
2152 let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2153 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2154 Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2155 },
2156 )
2157 } else if self.storage.dtype() == Some(DType::C32) {
2158 let inner = self
2159 .storage
2160 .eager()
2161 .ok_or_else(|| anyhow::anyhow!("failed to read c32 structured payload"))?;
2162 let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
2163 let output_indices = kept_indices.clone();
2164 self.select_structured_indices_typed(
2165 payload,
2166 kept_axes,
2167 kept_indices,
2168 kept_dims,
2169 (selected_axes, positions),
2170 move |payload, dims, _strides, classes| {
2171 let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2172 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2173 Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2174 },
2175 )
2176 } else {
2177 Err(anyhow::anyhow!(
2178 "unsupported structured storage scalar type"
2179 ))
2180 }
2181 }
2182
2183 fn validate_storage_matches_indices(indices: &[DynIndex], storage: &Storage) -> Result<()> {
2184 let dims = Self::expected_dims_from_indices(indices);
2185 let storage_dims = storage.logical_dims();
2186 if storage_dims != dims {
2187 return Err(anyhow::anyhow!(
2188 "storage logical dims {:?} do not match indices dims {:?}",
2189 storage_dims,
2190 dims
2191 ));
2192 }
2193 if storage.is_diag() {
2194 Self::validate_diag_dims(&dims)?;
2195 }
2196 Ok(())
2197 }
2198
2199 fn try_materialized_inner(&self) -> Result<&EagerTensor> {
2200 self.ensure_storage_ready()?;
2201 let logical_dims = self.dims();
2202 if let Some(value) = self.tracked_compact_payload_value() {
2203 if self.compact_payload_is_logical_dense(&value.payload_dims) {
2204 return Ok(value.payload.as_ref());
2205 }
2206 if self.eager_cache.get().is_none() {
2207 let dense = Self::dense_inner_from_payload(
2208 value.payload.as_ref(),
2209 &value.axis_classes,
2210 &logical_dims,
2211 )?;
2212 let _ = self.eager_cache.set(Arc::new(dense));
2213 }
2214 return self
2215 .eager_cache
2216 .get()
2217 .map(|inner| inner.as_ref())
2218 .ok_or_else(|| {
2219 anyhow::anyhow!("IdxTensor structured AD cache was not initialized")
2220 });
2221 }
2222 if let Some(inner) = self.storage.eager() {
2223 if self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len()) {
2224 return Ok(inner);
2225 }
2226 if self.eager_cache.get().is_none() {
2227 let dense = Self::dense_inner_from_payload(
2228 inner,
2229 self.storage.axis_classes(),
2230 &logical_dims,
2231 )?;
2232 let _ = self.eager_cache.set(Arc::new(dense));
2233 }
2234 return self
2235 .eager_cache
2236 .get()
2237 .map(|inner| inner.as_ref())
2238 .ok_or_else(|| {
2239 anyhow::anyhow!("IdxTensor structured eager cache was not initialized")
2240 });
2241 }
2242 if self.eager_cache.get().is_none() {
2243 let native = profile_pairwise_contract_section("materialize_storage_to_native", || {
2244 let storage = self.storage.materialize(self.indices.len())?;
2245 Self::seed_native_payload(storage.as_ref(), &logical_dims)
2246 })
2247 .context("IdxTensor materialization failed")?;
2248 record_pairwise_contract_profile_bytes(
2249 "materialize_storage_to_native",
2250 tensor_profile_bytes(native.dtype(), native.shape()),
2251 );
2252 let _ = self.eager_cache.set(Arc::new(EagerTensor::from_tensor_in(
2253 native,
2254 default_eager_ctx()?,
2255 )?));
2256 }
2257 self.eager_cache
2258 .get()
2259 .map(|inner| inner.as_ref())
2260 .ok_or_else(|| anyhow::anyhow!("IdxTensor materialization cache was not initialized"))
2261 }
2262
2263 pub(crate) fn as_inner(&self) -> Result<&EagerTensor> {
2264 self.try_materialized_inner()
2265 }
2266
2267 #[inline]
2269 fn expected_dims_from_indices(indices: &[DynIndex]) -> Vec<usize> {
2270 indices.iter().map(|idx| idx.dim()).collect()
2271 }
2272
2273 pub fn dims(&self) -> Vec<usize> {
2292 Self::expected_dims_from_indices(&self.indices)
2293 }
2294
2295 pub fn select_indices(
2338 &self,
2339 selected_indices: &[DynIndex],
2340 positions: &[usize],
2341 ) -> std::result::Result<Self, IdxTensorError> {
2342 if selected_indices.len() != positions.len() {
2343 return Err(anyhow::anyhow!(
2344 "selected_indices length {} does not match positions length {}",
2345 selected_indices.len(),
2346 positions.len()
2347 )
2348 .into());
2349 }
2350 if selected_indices.is_empty() {
2351 return Ok(self.clone());
2352 }
2353
2354 let mut selected_axes = Vec::with_capacity(selected_indices.len());
2355 let mut seen_axes = HashSet::with_capacity(selected_indices.len());
2356 for (selected, &position) in selected_indices.iter().zip(positions.iter()) {
2357 let axis = self
2358 .indices
2359 .iter()
2360 .position(|index| index == selected)
2361 .ok_or_else(|| anyhow::anyhow!("selected index is not present in tensor"))?;
2362 if !seen_axes.insert(axis) {
2363 return Err(anyhow::anyhow!("selected index appears more than once").into());
2364 }
2365 let dim = self.indices[axis].dim();
2366 if position >= dim {
2367 return Err(anyhow::anyhow!(
2368 "selected coordinate {position} is out of range for axis {axis} with dim {dim}"
2369 )
2370 .into());
2371 }
2372 selected_axes.push(axis);
2373 }
2374
2375 let kept_axes = self
2376 .indices
2377 .iter()
2378 .enumerate()
2379 .filter(|(axis, _)| !seen_axes.contains(axis))
2380 .map(|(axis, _)| axis)
2381 .collect::<Vec<_>>();
2382 let kept_indices = kept_axes
2383 .iter()
2384 .map(|&axis| self.indices[axis].clone())
2385 .collect::<Vec<_>>();
2386 let kept_dims = kept_axes
2387 .iter()
2388 .map(|&axis| self.indices[axis].dim())
2389 .collect::<Vec<_>>();
2390
2391 if matches!(
2392 self.storage.storage_kind(),
2393 StorageKind::Diagonal | StorageKind::Structured
2394 ) {
2395 self.ensure_shape_packing_preserves_ad("select_indices")?;
2396 }
2397 if self.storage.storage_kind() == StorageKind::Diagonal {
2398 return self
2399 .select_diag_indices(kept_indices, kept_dims, positions)
2400 .map_err(IdxTensorError::from);
2401 }
2402 if self.storage.storage_kind() == StorageKind::Structured {
2403 return self
2404 .select_structured_indices(
2405 &kept_axes,
2406 kept_indices,
2407 kept_dims,
2408 &selected_axes,
2409 positions,
2410 )
2411 .map_err(IdxTensorError::from);
2412 }
2413 if self.storage.storage_kind() != StorageKind::Dense {
2414 return Err(anyhow::anyhow!(
2415 "select_indices got unsupported storage kind {:?}",
2416 self.storage.storage_kind()
2417 )
2418 .into());
2419 }
2420
2421 let rank = self.indices.len();
2422 let mut starts = vec![0_i64; rank];
2423 let mut slice_sizes = self.dims();
2424 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
2425 starts[axis] = i64::try_from(position)
2426 .map_err(|_| anyhow::anyhow!("selected coordinate does not fit in i64"))?;
2427 slice_sizes[axis] = 1;
2428 }
2429
2430 let starts_tensor = EagerTensor::from_tensor_in(
2431 NativeTensor::from_vec_col_major(vec![rank], starts).map_err(anyhow::Error::new)?,
2432 default_eager_ctx()?,
2433 )?;
2434 let sliced = self
2435 .try_materialized_inner()?
2436 .dynamic_slice(&starts_tensor, &slice_sizes)?;
2437 Self::from_inner(kept_indices, sliced.reshape(&kept_dims)?).map_err(IdxTensorError::from)
2438 }
2439
2440 pub fn stack_along_new_index(
2473 tensors: &[&Self],
2474 new_index: DynIndex,
2475 axis: isize,
2476 ) -> std::result::Result<Self, IdxTensorError> {
2477 let first = tensors
2478 .first()
2479 .copied()
2480 .ok_or_else(|| anyhow::anyhow!("stack_along_new_index requires at least one tensor"))?;
2481 if !(new_index.dim() == tensors.len()) {
2482 return Err(anyhow::anyhow!(
2483 "stack_along_new_index: new index dim {} does not match tensor count {}",
2484 new_index.dim(),
2485 tensors.len()
2486 )
2487 .into());
2488 };
2489
2490 let base_indices = first.indices.clone();
2491 for tensor in tensors.iter().copied().skip(1) {
2492 if !(tensor.indices == base_indices) {
2493 return Err(anyhow::anyhow!(
2494 "stack_along_new_index: input tensors must have identical index order"
2495 )
2496 .into());
2497 };
2498 }
2499 for &tensor in tensors {
2500 tensor.ensure_shape_packing_preserves_ad("stack_along_new_index")?;
2501 }
2502
2503 let insert_axis =
2504 Self::normalize_insert_axis("stack_along_new_index", axis, base_indices.len())?;
2505 let mut result_indices = base_indices;
2506 result_indices.insert(insert_axis, new_index);
2507
2508 let inner_refs = tensors
2509 .iter()
2510 .map(|tensor| tensor.try_materialized_inner())
2511 .collect::<Result<Vec<_>>>()?;
2512 let stacked = EagerTensor::stack(&inner_refs, axis)?;
2513 Self::from_inner(result_indices, stacked).map_err(IdxTensorError::from)
2514 }
2515
2516 pub fn index_select(
2546 &self,
2547 source_index: &DynIndex,
2548 target_index: DynIndex,
2549 positions: &[usize],
2550 ) -> std::result::Result<Self, IdxTensorError> {
2551 if !(target_index.dim() == positions.len()) {
2552 return Err(anyhow::anyhow!(
2553 "index_select: target index dim {} does not match position count {}",
2554 target_index.dim(),
2555 positions.len()
2556 )
2557 .into());
2558 };
2559 let axis = self
2560 .indices
2561 .iter()
2562 .position(|index| index == source_index)
2563 .ok_or_else(|| anyhow::anyhow!("index_select: source index is not present"))?;
2564 let source_dim = self.indices[axis].dim();
2565 for &position in positions {
2566 if !(position < source_dim) {
2567 return Err(anyhow::anyhow!(
2568 "index_select: position {position} is out of range for source dim {source_dim}"
2569 )
2570 .into());
2571 };
2572 }
2573 self.ensure_shape_packing_preserves_ad("index_select")?;
2574
2575 let axis = isize::try_from(axis)
2576 .map_err(|_| anyhow::anyhow!("index_select: axis does not fit in isize"))?;
2577 let selected = self
2578 .try_materialized_inner()?
2579 .index_select(axis, positions)?;
2580 let mut result_indices = self.indices.clone();
2581 result_indices[axis as usize] = target_index;
2582 Self::from_inner(result_indices, selected).map_err(IdxTensorError::from)
2583 }
2584
2585 pub fn new(
2604 indices: Vec<DynIndex>,
2605 storage: Arc<Storage>,
2606 ) -> std::result::Result<Self, IdxTensorError> {
2607 Self::from_storage(indices, storage)
2608 }
2609
2610 pub fn from_indices(
2631 indices: Vec<DynIndex>,
2632 storage: Arc<Storage>,
2633 ) -> std::result::Result<Self, IdxTensorError> {
2634 Self::new(indices, storage)
2635 }
2636
2637 pub fn from_storage(
2657 indices: Vec<DynIndex>,
2658 storage: Arc<Storage>,
2659 ) -> std::result::Result<Self, IdxTensorError> {
2660 Self::validate_indices(&indices)?;
2661 Self::validate_storage_matches_indices(&indices, storage.as_ref())?;
2662 Ok(Self {
2663 indices,
2664 storage: IdxTensorStorage::from_storage(storage),
2665 eager_cache: Self::empty_eager_cache(),
2666 })
2667 }
2668
2669 pub fn from_structured_storage(
2691 indices: Vec<DynIndex>,
2692 storage: Arc<Storage>,
2693 ) -> std::result::Result<Self, IdxTensorError> {
2694 Self::from_storage(indices, storage)
2695 }
2696
2697 pub fn from_copy_selector<T>(
2753 left: DynIndex,
2754 site: DynIndex,
2755 right: DynIndex,
2756 selected_value: usize,
2757 scale: T,
2758 ) -> std::result::Result<Self, StructuredSelectorError>
2759 where
2760 T: TensorElement + Copy + Zero,
2761 {
2762 if left.dim == 0 {
2763 return Err(StructuredSelectorError::ZeroDimension { axis: "left" });
2764 }
2765 if site.dim == 0 {
2766 return Err(StructuredSelectorError::ZeroDimension { axis: "site" });
2767 }
2768 if right.dim == 0 {
2769 return Err(StructuredSelectorError::ZeroDimension { axis: "right" });
2770 }
2771 if left.dim != right.dim {
2772 return Err(StructuredSelectorError::BondDimensionMismatch {
2773 left: left.dim,
2774 right: right.dim,
2775 });
2776 }
2777 if selected_value >= site.dim {
2778 return Err(StructuredSelectorError::SelectedValueOutOfBounds {
2779 value: selected_value,
2780 site_dim: site.dim,
2781 });
2782 }
2783
2784 let payload_len =
2785 left.dim
2786 .checked_mul(site.dim)
2787 .ok_or(StructuredSelectorError::PayloadSizeOverflow {
2788 bond_dim: left.dim,
2789 site_dim: site.dim,
2790 })?;
2791 let _site_stride = isize::try_from(left.dim)
2792 .map_err(|_| StructuredSelectorError::StrideOverflow { bond_dim: left.dim })?;
2793 let selected_offset = left.dim.checked_mul(selected_value).ok_or(
2794 StructuredSelectorError::PayloadSizeOverflow {
2795 bond_dim: left.dim,
2796 site_dim: site.dim,
2797 },
2798 )?;
2799
2800 let mut payload = Vec::new();
2801 payload.try_reserve_exact(payload_len).map_err(|_| {
2802 StructuredSelectorError::AllocationFailed {
2803 elements: payload_len,
2804 }
2805 })?;
2806 payload.resize(payload_len, T::zero());
2807 for bond in 0..left.dim {
2808 payload[selected_offset + bond] = scale;
2809 }
2810 let payload_native = dense_native_tensor_from_col_major(&payload, &[left.dim, site.dim])
2811 .map_err(|error| StructuredSelectorError::InvalidStorage {
2812 message: error.to_string(),
2813 })?;
2814 let payload_dtype = payload_native.dtype();
2815 let payload_inner = EagerTensor::from_tensor_in(
2816 payload_native,
2817 default_eager_ctx().map_err(|error| StructuredSelectorError::InvalidStorage {
2818 message: error.to_string(),
2819 })?,
2820 )
2821 .map_err(|error| StructuredSelectorError::InvalidStorage {
2822 message: error.to_string(),
2823 })?;
2824 let payload_dims = vec![left.dim, site.dim];
2825 let indices = vec![left, site, right];
2826 if !matches!(
2827 payload_dtype,
2828 DType::F32 | DType::F64 | DType::C32 | DType::C64
2829 ) {
2830 return Err(StructuredSelectorError::InvalidStorage {
2831 message: format!("unsupported selector dtype {:?}", payload_dtype),
2832 });
2833 }
2834 Self::from_structured_payload_inner(indices, payload_inner, payload_dims, vec![0, 1, 0])
2835 .map_err(|error| StructuredSelectorError::InvalidStorage {
2836 message: error.to_string(),
2837 })
2838 }
2839
2840 pub(crate) fn from_native(indices: Vec<DynIndex>, native: NativeTensor) -> Result<Self> {
2842 let axis_classes = Self::dense_axis_classes(indices.len());
2843 Self::from_native_with_axis_classes(indices, native, axis_classes)
2844 }
2845
2846 pub(crate) fn from_native_with_axis_classes(
2847 indices: Vec<DynIndex>,
2848 native: NativeTensor,
2849 axis_classes: Vec<usize>,
2850 ) -> Result<Self> {
2851 Self::from_inner_with_axis_classes(
2852 indices,
2853 EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
2854 axis_classes,
2855 )
2856 }
2857
2858 pub(crate) fn from_inner(indices: Vec<DynIndex>, inner: EagerTensor) -> Result<Self> {
2859 let axis_classes = Self::dense_axis_classes(indices.len());
2860 Self::from_inner_with_axis_classes(indices, inner, axis_classes)
2861 }
2862
2863 pub fn hermitian_eigendecomposition(
2903 &self,
2904 hermitian_tol: f64,
2905 ) -> std::result::Result<TensorHermitianEigendecomposition, IdxTensorError> {
2906 if !(self.indices.len() == 2) {
2907 return Err(anyhow::anyhow!(
2908 "IdxTensor::hermitian_eigendecomposition requires a rank-2 tensor, got rank {}",
2909 self.indices.len()
2910 )
2911 .into());
2912 };
2913 let dims = self.dims();
2914 if !(dims[0] == dims[1]) {
2915 return Err(anyhow::anyhow!(
2916 "IdxTensor::hermitian_eigendecomposition requires a square matrix, got {}x{}",
2917 dims[0],
2918 dims[1]
2919 )
2920 .into());
2921 };
2922 if !(dims[0] > 0) {
2923 return Err(anyhow::anyhow!(
2924 "IdxTensor::hermitian_eigendecomposition requires a non-empty matrix"
2925 )
2926 .into());
2927 };
2928 if !(hermitian_tol.is_finite() && hermitian_tol >= 0.0) {
2929 return Err(anyhow::anyhow!(
2930 "IdxTensor::hermitian_eigendecomposition requires a finite non-negative tolerance"
2931 )
2932 .into());
2933 };
2934
2935 let input = self.try_materialized_inner()?;
2936 let (values, vectors) = input
2937 .eigh()
2938 .map_err(|source| anyhow::anyhow!("Hermitian eigendecomposition failed: {source}"))?;
2939
2940 let eigenvalue_index = DynIndex::new_dyn(dims[0]);
2941 let eigenvector_index = DynIndex::new_dyn(dims[0]);
2942 let eigenvalue_tensor = Self::from_inner(vec![eigenvalue_index], values)?;
2943 let eigenvalues = Self::read_real_eigenvalues(&eigenvalue_tensor, hermitian_tol)
2944 .with_context(|| {
2945 "IdxTensor::hermitian_eigendecomposition failed to read eigenvalues"
2946 })?;
2947 let eigenvectors = Self::from_inner(
2948 vec![self.indices[0].clone(), eigenvector_index.clone()],
2949 vectors,
2950 )?;
2951
2952 Ok(TensorHermitianEigendecomposition {
2953 eigenvalues,
2954 eigenvectors,
2955 eigenvector_index,
2956 })
2957 }
2958
2959 fn read_real_eigenvalues(values: &Self, hermitian_tol: f64) -> Result<Vec<f64>> {
2960 if values.is_complex() {
2961 values
2962 .to_vec::<Complex64>()?
2963 .into_iter()
2964 .enumerate()
2965 .map(|(index, value)| {
2966 let imaginary = value.im.abs();
2967 let allowed = hermitian_tol * value.norm().max(1.0);
2968 if !matches!(
2969 imaginary.partial_cmp(&allowed),
2970 Some(std::cmp::Ordering::Less) | Some(std::cmp::Ordering::Equal)
2971 ) {
2972 return Err(anyhow::anyhow!("Hermitian eigenvalue {index} has imaginary part {imaginary}, exceeding tolerance {allowed}"));
2973 };
2974 Ok(value.re)
2975 })
2976 .collect()
2977 } else {
2978 values.to_vec::<f64>().map_err(anyhow::Error::from)
2979 }
2980 }
2981
2982 pub(crate) fn from_diag_inner(
2983 indices: Vec<DynIndex>,
2984 payload_inner: EagerTensor,
2985 ) -> Result<Self> {
2986 let dims = Self::expected_dims_from_indices(&indices);
2987 Self::validate_indices(&indices)?;
2988 Self::validate_diag_dims(&dims)?;
2989 let payload_len = checked_product(payload_inner.shape())?;
2990 Self::validate_diag_payload_len(payload_len, &dims)?;
2991 let axis_classes = Self::diag_axis_classes(dims.len());
2992 let diag_inner = payload_inner.embed_diag(0, 1)?;
2993 Self::from_inner_with_axis_classes(indices, diag_inner, axis_classes)
2994 }
2995
2996 fn compact_inner_from_logical(
2997 inner: &EagerTensor,
2998 axis_classes: &[usize],
2999 ) -> Result<EagerTensor> {
3000 let mut payload = inner.clone();
3001 let mut classes = axis_classes.to_vec();
3002 while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(&classes) {
3003 payload = payload.extract_diag(axis_a, axis_b)?;
3004 classes.remove(axis_b);
3005 }
3006 Ok(payload)
3007 }
3008
3009 pub(crate) fn from_inner_with_axis_classes(
3010 indices: Vec<DynIndex>,
3011 inner: EagerTensor,
3012 axis_classes: Vec<usize>,
3013 ) -> Result<Self> {
3014 let dims = profile_pairwise_contract_section("from_inner_expected_dims", || {
3015 Self::expected_dims_from_indices(&indices)
3016 });
3017 profile_pairwise_contract_section("from_inner_validate_indices", || {
3018 Self::validate_indices(&indices)
3019 })?;
3020 Self::validate_axis_classes(&axis_classes, indices.len())?;
3021 if dims != inner.shape() {
3022 return Err(anyhow::anyhow!(
3023 "native payload dims {:?} do not match indices dims {:?}",
3024 inner.shape(),
3025 dims
3026 ));
3027 }
3028 if Self::is_diag_axis_classes(&axis_classes) {
3029 profile_pairwise_contract_section("from_inner_validate_diag_dims", || {
3030 Self::validate_diag_dims(&dims)
3031 })?;
3032 }
3033 let storage = if axis_classes == Self::dense_axis_classes(indices.len()) {
3034 IdxTensorStorage::from_eager_dense(inner, indices.len())
3035 } else {
3036 let payload = Self::compact_inner_from_logical(&inner, &axis_classes)?;
3037 let payload_dims = payload.shape().to_vec();
3038 IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3039 payload: Arc::new(payload),
3040 payload_dims,
3041 axis_classes,
3042 }))
3043 };
3044 Ok(Self {
3045 indices,
3046 storage,
3047 eager_cache: Self::empty_eager_cache(),
3048 })
3049 }
3050
3051 pub fn indices(&self) -> &[DynIndex] {
3053 &self.indices
3054 }
3055
3056 pub(crate) fn axis_classes(&self) -> &[usize] {
3057 self.storage.axis_classes()
3058 }
3059
3060 #[cfg(feature = "tenferro-cuda")]
3061 pub(crate) fn deferred_storage_error(&self) -> Option<&TensorStorageError> {
3062 self.storage.deferred_error()
3063 }
3064
3065 #[cfg(feature = "tenferro-cuda")]
3066 pub(crate) fn cuda_eager_inner(&self) -> Option<&EagerTensor> {
3067 self.storage
3068 .eager()
3069 .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3070 }
3071
3072 #[cfg(feature = "tenferro-cuda")]
3073 pub(crate) fn cuda_duplicate_native(&self) -> Result<NativeTensor> {
3074 Ok(self.try_materialized_inner()?.duplicate_value()?)
3075 }
3076
3077 pub fn enable_grad(self) -> std::result::Result<Self, IdxTensorError> {
3083 self.ensure_storage_ready()?;
3084 let eager_payload = self
3087 .storage
3088 .eager()
3089 .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3090 .filter(|inner| inner.shape() == self.storage.payload_dims());
3091 let payload = match eager_payload {
3092 Some(inner) => inner.duplicate_value()?,
3093 None => {
3094 let materialized = self.storage.materialize(self.indices.len())?;
3095 storage_payload_native(materialized.as_ref())
3096 .context("IdxTensor::enable_grad failed")?
3097 }
3098 };
3099 let payload_dims = self.storage.payload_dims().to_vec();
3100 let axis_classes = self.storage.axis_classes().to_vec();
3101 let tracked = Arc::new(EagerTensor::requires_grad_in(
3102 payload,
3103 default_eager_ctx()?,
3104 )?);
3105 let storage = if axis_classes == Self::dense_axis_classes(self.indices.len()) {
3106 IdxTensorStorage::Eager {
3107 inner: tracked,
3108 axis_classes,
3109 }
3110 } else {
3111 IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3112 payload: tracked,
3113 payload_dims,
3114 axis_classes,
3115 }))
3116 };
3117 Ok(Self {
3118 indices: self.indices,
3119 storage,
3120 eager_cache: Self::empty_eager_cache(),
3121 })
3122 }
3123
3124 pub fn tracks_grad(&self) -> bool {
3126 self.storage.eager().is_some_and(EagerTensor::tracks_grad)
3127 || self
3128 .eager_cache
3129 .get()
3130 .is_some_and(|inner| inner.tracks_grad())
3131 }
3132
3133 pub fn grad(&self) -> std::result::Result<Option<Self>, IdxTensorError> {
3139 if let Some(value) = self.tracked_compact_payload_value() {
3140 let Some(gradient) = value.payload.grad()? else {
3141 return Ok(None);
3142 };
3143 let gradient_shape = gradient.shape().to_vec();
3144 let gradient_tensor = gradient.to_tensor()?;
3145 if self.compact_payload_is_logical_dense(&value.payload_dims) {
3146 return Ok(Some(Self::from_native_with_axis_classes(
3147 self.indices.clone(),
3148 gradient_tensor,
3149 value.axis_classes.clone(),
3150 )?));
3151 }
3152 if gradient_shape != value.payload_dims {
3153 return Err(anyhow::anyhow!(
3154 "gradient payload dims {:?} do not match {:?}",
3155 gradient_shape,
3156 value.payload_dims
3157 )
3158 .into());
3159 }
3160 let gradient = EagerTensor::from_tensor_in(gradient_tensor, default_eager_ctx()?)?;
3161 return Ok(Some(Self::from_structured_payload_inner(
3162 self.indices.clone(),
3163 gradient,
3164 value.payload_dims.clone(),
3165 value.axis_classes.clone(),
3166 )?));
3167 }
3168
3169 let Some(gradient) = self.try_materialized_inner()?.grad()? else {
3170 return Ok(None);
3171 };
3172 Ok(Some(Self::from_native_with_axis_classes(
3173 self.indices.clone(),
3174 gradient.to_tensor()?,
3175 self.storage.axis_classes().to_vec(),
3176 )?))
3177 }
3178
3179 pub fn clear_grad(&self) -> std::result::Result<(), IdxTensorError> {
3185 self.ensure_storage_ready()?;
3186 if let Some(value) = self.tracked_compact_payload_value() {
3187 value.payload.clear_grad()?;
3188 }
3189 if let Some(inner) = self.storage.eager() {
3190 inner.clear_grad()?;
3191 }
3192 if let Some(inner) = self.eager_cache.get() {
3193 inner.clear_grad()?;
3194 }
3195 Ok(())
3196 }
3197
3198 pub fn backward(&self) -> std::result::Result<(), IdxTensorError> {
3204 if let Some(value) = self.tracked_compact_payload_value() {
3205 return value.payload.backward().map(|_| ()).map_err(|e| {
3206 IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}"))
3207 });
3208 }
3209 self.try_materialized_inner()?
3210 .backward()
3211 .map(|_| ())
3212 .map_err(|e| IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}")))
3213 }
3214
3215 pub fn detach(&self) -> std::result::Result<Self, IdxTensorError> {
3221 Self::from_inner_with_axis_classes(
3222 self.indices.clone(),
3223 self.try_materialized_inner()?.detach(),
3224 self.storage.axis_classes().to_vec(),
3225 )
3226 .map_err(IdxTensorError::from)
3227 }
3228
3229 pub fn is_simple(&self) -> bool {
3231 true
3232 }
3233
3234 pub fn to_storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3245 self.storage.materialize(self.indices.len())
3246 }
3247
3248 pub fn storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3267 self.storage.materialize(self.indices.len())
3268 }
3269
3270 pub fn storage_kind(&self) -> StorageKind {
3289 self.storage.storage_kind()
3290 }
3291
3292 #[cfg(feature = "tenferro-cuda")]
3321 pub fn cuda_dtype(&self) -> std::result::Result<DType, IdxTensorError> {
3322 self.scalar_dtype().map_err(IdxTensorError::from)
3323 }
3324
3325 pub fn sum(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3341 self.ensure_storage_ready()?;
3342 if self.indices.is_empty() {
3343 return AnyScalar::from_tensor(self.clone()).map_err(IdxTensorError::from);
3344 }
3345 if let Some(payload) = self.storage.eager().filter(|payload| payload.tracks_grad()) {
3346 let axes: Vec<usize> = (0..payload.shape().len()).collect();
3347 let reduced = payload.reduce_sum(Some(&axes))?;
3348 return AnyScalar::from_tensor(Self::from_inner(Vec::new(), reduced)?)
3349 .map_err(IdxTensorError::from);
3350 }
3351 self.storage.sum_scalar().map_err(IdxTensorError::from)
3352 }
3353
3354 pub fn only(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3378 let dims = self.dims();
3379 let total_size = checked_product(&dims)?;
3380 if !(total_size == 1 || dims.is_empty()) {
3381 return Err(anyhow::anyhow!(
3382 "only() requires a scalar tensor (1 element), got {} elements with dims {:?}",
3383 if dims.is_empty() { 1 } else { total_size },
3384 dims
3385 )
3386 .into());
3387 };
3388 self.sum()
3389 }
3390
3391 pub fn permute_indices(
3426 &self,
3427 new_indices: &[DynIndex],
3428 ) -> std::result::Result<Self, IdxTensorError> {
3429 let perm = compute_permutation_from_indices(&self.indices, new_indices)?;
3431 if perm.iter().copied().eq(0..perm.len()) {
3432 return Ok(Self {
3433 indices: new_indices.to_vec(),
3434 storage: self.storage.clone(),
3435 eager_cache: Arc::clone(&self.eager_cache),
3436 });
3437 }
3438
3439 let permuted = self.try_materialized_inner()?.transpose(&perm)?;
3440 let axis_classes = self.permute_axis_classes(&perm);
3441 Self::from_inner_with_axis_classes(new_indices.to_vec(), permuted, axis_classes)
3442 .map_err(IdxTensorError::from)
3443 }
3444
3445 pub fn permute(&self, perm: &[usize]) -> std::result::Result<Self, IdxTensorError> {
3477 if !(perm.len() == self.indices.len()) {
3478 return Err(anyhow::anyhow!("permutation length must match tensor rank").into());
3479 };
3480 let mut seen = HashSet::new();
3481 for &axis in perm {
3482 if !(axis < self.indices.len()) {
3483 return Err(anyhow::anyhow!("permutation axis {axis} out of range").into());
3484 };
3485 if !(seen.insert(axis)) {
3486 return Err(anyhow::anyhow!("duplicate axis {axis} in permutation").into());
3487 };
3488 }
3489 if perm.iter().copied().eq(0..perm.len()) {
3490 return Ok(self.clone());
3491 }
3492
3493 let new_indices: Vec<DynIndex> = perm.iter().map(|&i| self.indices[i].clone()).collect();
3495 let permuted = self.try_materialized_inner()?.transpose(perm)?;
3496 let axis_classes = self.permute_axis_classes(perm);
3497 Self::from_inner_with_axis_classes(new_indices, permuted, axis_classes)
3498 .map_err(IdxTensorError::from)
3499 }
3500
3501 pub(crate) fn try_contract_pairwise_default(&self, other: &Self) -> Result<Self> {
3502 self.try_contract_pairwise_default_with_options(other, PairwiseContractionOptions::new())
3503 }
3504
3505 pub(crate) fn try_contract_pairwise_default_with_options(
3506 &self,
3507 other: &Self,
3508 options: PairwiseContractionOptions,
3509 ) -> Result<Self> {
3510 let self_indices = profile_pairwise_contract_section("operand_indices", || {
3511 self.operand_indices_for_contraction(options.lhs_conj)
3512 });
3513 let other_indices = profile_pairwise_contract_section("operand_indices", || {
3514 other.operand_indices_for_contraction(options.rhs_conj)
3515 });
3516 let self_dims = profile_pairwise_contract_section("expected_dims", || {
3517 Self::expected_dims_from_indices(&self_indices)
3518 });
3519 let other_dims = profile_pairwise_contract_section("expected_dims", || {
3520 Self::expected_dims_from_indices(&other_indices)
3521 });
3522 let spec = profile_pairwise_contract_section("prepare_contraction", || {
3523 prepare_contraction(&self_indices, &self_dims, &other_indices, &other_dims)
3524 })
3525 .context("contraction preparation failed")?;
3526 let result_axis_classes = profile_pairwise_contract_section("result_axis_classes", || {
3527 Self::binary_contraction_axis_classes(
3528 self.storage.axis_classes(),
3529 &spec.axes_a,
3530 other.storage.axis_classes(),
3531 &spec.axes_b,
3532 )
3533 })?;
3534
3535 if profile_pairwise_contract_section("structured_check", || {
3536 self.should_use_structured_payload_contract(other)
3537 }) {
3538 if options.has_conj() {
3539 let lhs = if options.lhs_conj {
3540 self.conj()
3541 } else {
3542 self.clone()
3543 };
3544 let rhs = if options.rhs_conj {
3545 other.conj()
3546 } else {
3547 other.clone()
3548 };
3549 return profile_pairwise_contract_section("structured_conj_fallback", || {
3550 lhs.try_contract_pairwise_default(&rhs)
3551 });
3552 }
3553 return profile_pairwise_contract_section("structured_payload_contract", || {
3554 self.contract_structured_payloads(
3555 other,
3556 spec.result_indices.into_vec(),
3557 &spec.axes_a,
3558 &spec.axes_b,
3559 )
3560 });
3561 }
3562
3563 if self.indices.is_empty() && other.indices.is_empty() {
3564 if options.has_conj() {
3565 let lhs = if options.lhs_conj {
3566 self.conj()
3567 } else {
3568 self.clone()
3569 };
3570 let rhs = if options.rhs_conj {
3571 other.conj()
3572 } else {
3573 other.clone()
3574 };
3575 return lhs.try_contract_pairwise_default(&rhs);
3576 }
3577 let result = profile_pairwise_contract_section("scalar_mul", || {
3578 Ok::<_, anyhow::Error>(
3579 self.try_materialized_inner()?
3580 .mul(other.try_materialized_inner()?)?,
3581 )
3582 })?;
3583 return profile_pairwise_contract_section("from_inner", || {
3584 Self::from_inner(spec.result_indices.into_vec(), result)
3585 });
3586 }
3587
3588 let self_dtype = self.try_materialized_inner()?.dtype();
3589 let other_dtype = other.try_materialized_inner()?.dtype();
3590 if self_dtype != other_dtype {
3591 if options.has_conj() {
3592 let lhs = if options.lhs_conj {
3593 self.conj()
3594 } else {
3595 self.clone()
3596 };
3597 let rhs = if options.rhs_conj {
3598 other.conj()
3599 } else {
3600 other.clone()
3601 };
3602 return lhs.try_contract_pairwise_default(&rhs);
3603 }
3604 let self_native = self.try_materialized_inner()?.duplicate_value()?;
3605 let other_native = other.try_materialized_inner()?.duplicate_value()?;
3606 let result_native = profile_pairwise_contract_section("native_contract", || {
3607 contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)
3608 })?;
3609 return profile_pairwise_contract_section("from_native", || {
3610 Self::from_native_with_axis_classes(
3611 spec.result_indices.into_vec(),
3612 result_native,
3613 result_axis_classes,
3614 )
3615 });
3616 }
3617
3618 let config = profile_pairwise_contract_section("build_dot_general_config", || {
3619 Self::binary_dot_general_config(&spec.axes_a, &spec.axes_b)
3620 })?;
3621 let result = profile_pairwise_contract_section("dot_general_with_conj", || {
3622 let lhs = profile_pairwise_contract_section("lhs_try_materialized_inner", || {
3623 self.try_materialized_inner()
3624 })?;
3625 let rhs = profile_pairwise_contract_section("rhs_try_materialized_inner", || {
3626 other.try_materialized_inner()
3627 })?;
3628 profile_pairwise_contract_section("dot_general_execute", || {
3629 lhs.dot_general_with_conj(rhs, config, options.lhs_conj, options.rhs_conj)
3630 })
3631 .map_err(anyhow::Error::from)
3632 })?;
3633 record_pairwise_contract_profile_bytes(
3634 "dot_general_output",
3635 tensor_profile_bytes(result.dtype(), result.shape()),
3636 );
3637 profile_pairwise_contract_section("from_inner_axis_classes", || {
3638 Self::from_inner_with_axis_classes(
3639 spec.result_indices.into_vec(),
3640 result,
3641 result_axis_classes,
3642 )
3643 })
3644 }
3645
3646 pub(crate) fn try_tensordot_pairwise_explicit(
3647 &self,
3648 other: &Self,
3649 pairs: &[(DynIndex, DynIndex)],
3650 ) -> Result<Self> {
3651 use crate::index_ops::ContractionError;
3652
3653 let self_dims = Self::expected_dims_from_indices(&self.indices);
3654 let other_dims = Self::expected_dims_from_indices(&other.indices);
3655 let spec = prepare_contraction_pairs(
3656 &self.indices,
3657 &self_dims,
3658 &other.indices,
3659 &other_dims,
3660 pairs,
3661 )
3662 .map_err(|e| match e {
3663 ContractionError::NoCommonIndices => {
3664 anyhow::anyhow!("tensordot: No pairs specified for contraction")
3665 }
3666 ContractionError::BatchContractionNotImplemented => anyhow::anyhow!(
3667 "tensordot: Common index found but not in contraction pairs. \
3668 Batch contraction is not yet implemented."
3669 ),
3670 ContractionError::IndexNotFound { tensor } => {
3671 anyhow::anyhow!("tensordot: Index not found in {} tensor", tensor)
3672 }
3673 ContractionError::DimensionMismatch {
3674 pos_a,
3675 pos_b,
3676 dim_a,
3677 dim_b,
3678 } => anyhow::anyhow!(
3679 "tensordot: Dimension mismatch: self[{}]={} != other[{}]={}",
3680 pos_a,
3681 dim_a,
3682 pos_b,
3683 dim_b
3684 ),
3685 ContractionError::DuplicateAxis { tensor, pos } => {
3686 anyhow::anyhow!("tensordot: Duplicate axis {} in {} tensor", pos, tensor)
3687 }
3688 })?;
3689 let result_axis_classes = Self::binary_contraction_axis_classes(
3690 self.storage.axis_classes(),
3691 &spec.axes_a,
3692 other.storage.axis_classes(),
3693 &spec.axes_b,
3694 )?;
3695
3696 if self.should_use_structured_payload_contract(other) {
3697 return self.contract_structured_payloads(
3698 other,
3699 spec.result_indices.into_vec(),
3700 &spec.axes_a,
3701 &spec.axes_b,
3702 );
3703 }
3704
3705 if self.indices.is_empty() && other.indices.is_empty() {
3706 let result = self
3707 .try_materialized_inner()?
3708 .mul(other.try_materialized_inner()?)
3709 .map_err(|e| anyhow::anyhow!("tensordot scalar multiply failed: {e}"))?;
3710 return Self::from_inner(spec.result_indices.into_vec(), result);
3711 }
3712
3713 let self_dtype = self.try_materialized_inner()?.dtype();
3714 let other_dtype = other.try_materialized_inner()?.dtype();
3715 if self_dtype != other_dtype {
3716 let self_native = self.try_materialized_inner()?.duplicate_value()?;
3717 let other_native = other.try_materialized_inner()?.duplicate_value()?;
3718 let result_native =
3719 contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)?;
3720 return Self::from_native_with_axis_classes(
3721 spec.result_indices.into_vec(),
3722 result_native,
3723 result_axis_classes,
3724 );
3725 }
3726
3727 let subscripts = Self::build_binary_einsum_subscripts(
3728 self.indices.len(),
3729 &spec.axes_a,
3730 other.indices.len(),
3731 &spec.axes_b,
3732 )?;
3733 let result = [
3734 self.try_materialized_inner()?,
3735 other.try_materialized_inner()?,
3736 ]
3737 .einsum_subscripts(&subscripts)
3738 .map_err(|e| anyhow::anyhow!("tensordot failed: {e}"))?;
3739 Self::from_inner_with_axis_classes(
3740 spec.result_indices.into_vec(),
3741 result,
3742 result_axis_classes,
3743 )
3744 }
3745
3746 pub(crate) fn try_outer_product_pairwise(&self, other: &Self) -> Result<Self> {
3747 use anyhow::Context;
3748
3749 let common_positions = common_ind_positions(&self.indices, &other.indices);
3751 if !common_positions.is_empty() {
3752 let common_ids: Vec<_> = common_positions
3753 .iter()
3754 .map(|(pos_a, _)| self.indices[*pos_a].id())
3755 .collect();
3756 return Err(anyhow::anyhow!(
3757 "outer_product: tensors have common indices {:?}. \
3758 Use tensordot to contract common indices, or use sim() to replace \
3759 indices with fresh IDs before computing outer product.",
3760 common_ids
3761 ))
3762 .context("outer_product: common indices found");
3763 }
3764
3765 let mut result_indices = self.indices.clone();
3767 result_indices.extend(other.indices.iter().cloned());
3768 let result_axis_classes = Self::binary_contraction_axis_classes(
3769 self.storage.axis_classes(),
3770 &[],
3771 other.storage.axis_classes(),
3772 &[],
3773 )?;
3774 if self.should_use_structured_payload_contract(other) {
3775 return self.contract_structured_payloads(other, result_indices, &[], &[]);
3776 }
3777 let self_dtype = self.try_materialized_inner()?.dtype();
3778 let other_dtype = other.try_materialized_inner()?.dtype();
3779 if self_dtype != other_dtype {
3780 let self_native = self.try_materialized_inner()?.duplicate_value()?;
3781 let other_native = other.try_materialized_inner()?.duplicate_value()?;
3782 let result_native = contract_native_tensor(&self_native, &[], &other_native, &[])?;
3783 return Self::from_native_with_axis_classes(
3784 result_indices,
3785 result_native,
3786 result_axis_classes,
3787 );
3788 }
3789
3790 let subscripts = Self::build_binary_einsum_subscripts(
3791 self.indices.len(),
3792 &[],
3793 other.indices.len(),
3794 &[],
3795 )?;
3796 let result = [
3797 self.try_materialized_inner()?,
3798 other.try_materialized_inner()?,
3799 ]
3800 .einsum_subscripts(&subscripts)
3801 .map_err(|e| anyhow::anyhow!("outer_product failed: {e}"))?;
3802 Self::from_inner_with_axis_classes(result_indices, result, result_axis_classes)
3803 }
3804}
3805
3806impl IdxTensor {
3811 pub fn random<T: RandomScalar, R: Rng>(
3841 rng: &mut R,
3842 indices: Vec<DynIndex>,
3843 ) -> std::result::Result<Self, IdxTensorError> {
3844 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
3845 let size = checked_product(&dims)?;
3846 let data: Vec<T> = (0..size).map(|_| T::random_value(rng)).collect();
3847 Self::from_dense(indices, data)
3848 }
3849}
3850
3851impl IdxTensor {
3852 pub fn add(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
3890 if self.indices.len() != other.indices.len() {
3892 return Err(anyhow::anyhow!(
3893 "Index count mismatch: self has {} indices, other has {}",
3894 self.indices.len(),
3895 other.indices.len()
3896 )
3897 .into());
3898 }
3899
3900 let self_set: HashSet<_> = self.indices.iter().collect();
3902 let other_set: HashSet<_> = other.indices.iter().collect();
3903
3904 if self_set != other_set {
3905 return Err(
3906 anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
3907 );
3908 }
3909
3910 let other_aligned = other.permute_indices(&self.indices)?;
3912
3913 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
3915 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
3916 if self_expected_dims != other_expected_dims {
3917 use crate::TagSetLike;
3918 let fmt = |indices: &[DynIndex]| -> Vec<String> {
3919 indices
3920 .iter()
3921 .map(|idx| {
3922 let tags: Vec<String> = idx.tags().iter().collect();
3923 format!("{:?}(dim={},tags={:?})", idx.id(), idx.dim(), tags)
3924 })
3925 .collect()
3926 };
3927 return Err(anyhow::anyhow!(
3928 "Dimension mismatch after alignment.\n\
3929 self: dims={:?}, indices(order)={:?}\n\
3930 other_aligned: dims={:?}, indices(order)={:?}",
3931 self_expected_dims,
3932 fmt(&self.indices),
3933 other_expected_dims,
3934 fmt(&other_aligned.indices)
3935 )
3936 .into());
3937 }
3938
3939 self.axpby(
3940 AnyScalar::new_real(1.0),
3941 &other_aligned,
3942 AnyScalar::new_real(1.0),
3943 )
3944 }
3945
3946 pub fn axpby(
3971 &self,
3972 a: AnyScalar,
3973 other: &Self,
3974 b: AnyScalar,
3975 ) -> std::result::Result<Self, IdxTensorError> {
3976 if self.indices.len() != other.indices.len() {
3978 return Err(anyhow::anyhow!(
3979 "Index count mismatch: self has {} indices, other has {}",
3980 self.indices.len(),
3981 other.indices.len()
3982 )
3983 .into());
3984 }
3985
3986 let self_set: HashSet<_> = self.indices.iter().collect();
3988 let other_set: HashSet<_> = other.indices.iter().collect();
3989 if self_set != other_set {
3990 return Err(
3991 anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
3992 );
3993 }
3994
3995 let other_aligned = other.permute_indices(&self.indices)?;
3997
3998 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
4000 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
4001 if self_expected_dims != other_expected_dims {
4002 return Err(anyhow::anyhow!(
4003 "Dimension mismatch after alignment: self={:?}, other_aligned={:?}",
4004 self_expected_dims,
4005 other_expected_dims
4006 )
4007 .into());
4008 }
4009
4010 let axis_classes = if self.storage.axis_classes() == other_aligned.storage.axis_classes() {
4011 self.storage.axis_classes().to_vec()
4012 } else {
4013 Self::dense_axis_classes(self.indices.len())
4014 };
4015
4016 let same_compact_layout = self.storage.payload_dims()
4017 == other_aligned.storage.payload_dims()
4018 && self.storage.payload_strides_vec() == other_aligned.storage.payload_strides_vec()
4019 && self.storage.axis_classes() == other_aligned.storage.axis_classes();
4020 if same_compact_layout
4021 && matches!(&self.storage, IdxTensorStorage::Materialized(_))
4022 && matches!(&other_aligned.storage, IdxTensorStorage::Materialized(_))
4023 && !self.tracks_grad()
4024 && !other_aligned.tracks_grad()
4025 && !a.tracks_grad()
4026 && !b.tracks_grad()
4027 {
4028 let lhs_storage = self.storage.materialize(self.indices.len())?;
4029 let rhs_storage = other_aligned
4030 .storage
4031 .materialize(other_aligned.indices.len())?;
4032 let combined = lhs_storage
4033 .axpby(
4034 &a.to_backend_scalar(),
4035 rhs_storage.as_ref(),
4036 &b.to_backend_scalar(),
4037 )
4038 .map_err(|e| anyhow::anyhow!("storage axpby failed: {e}"))?;
4039 return Self::from_storage(self.indices.clone(), Arc::new(combined));
4040 }
4041
4042 let lhs = self.scale(a)?;
4043 let rhs = other_aligned.scale(b)?;
4044 let combined = lhs
4045 .try_materialized_inner()?
4046 .add(rhs.try_materialized_inner()?)
4047 .map_err(|e| anyhow::anyhow!("tensor addition failed: {e}"))?;
4048 Self::from_inner_with_axis_classes(self.indices.clone(), combined, axis_classes)
4049 .map_err(IdxTensorError::from)
4050 }
4051
4052 pub fn scale(&self, scalar: AnyScalar) -> std::result::Result<Self, IdxTensorError> {
4071 if matches!(
4072 &self.storage,
4073 IdxTensorStorage::Eager { .. }
4074 | IdxTensorStorage::Compact(_)
4075 | IdxTensorStorage::Materialized(_)
4076 ) {
4077 let storage = self.storage.scale_eager_payload(&scalar)?;
4082 return Ok(Self {
4083 indices: self.indices.clone(),
4084 storage,
4085 eager_cache: Self::empty_eager_cache(),
4086 });
4087 }
4088
4089 let self_dtype = self.try_materialized_inner()?.dtype();
4090 let scalar_dtype = scalar.as_tensor()?.try_materialized_inner()?.dtype();
4091 if self_dtype != scalar_dtype {
4092 let target_dtype = Self::scale_target_dtype(self_dtype, scalar_dtype)?;
4093 let self_inner = self.try_materialized_inner()?;
4094 let self_inner = if self_inner.dtype() == target_dtype {
4095 self_inner.clone()
4096 } else {
4097 self_inner.cast(target_dtype)?
4098 };
4099 let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
4100 let scalar_inner = if scalar_inner.dtype() == target_dtype {
4101 scalar_inner.clone()
4102 } else {
4103 scalar_inner.cast(target_dtype)?
4104 };
4105 let scaled = if self.indices.is_empty() {
4106 self_inner
4107 .mul(&scalar_inner)
4108 .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4109 } else {
4110 let subscripts = Self::scale_subscripts(self.indices.len())?;
4111 [&self_inner, &scalar_inner]
4112 .einsum_subscripts(&subscripts)
4113 .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4114 };
4115 return Self::from_inner_with_axis_classes(
4116 self.indices.clone(),
4117 scaled,
4118 self.storage.axis_classes().to_vec(),
4119 )
4120 .map_err(IdxTensorError::from);
4121 }
4122 let scaled = if self.indices.is_empty() {
4123 self.try_materialized_inner()?
4124 .mul(scalar.as_tensor()?.try_materialized_inner()?)
4125 .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4126 } else {
4127 let subscripts = Self::scale_subscripts(self.indices.len())?;
4128 [
4129 self.try_materialized_inner()?,
4130 scalar.as_tensor()?.try_materialized_inner()?,
4131 ]
4132 .einsum_subscripts(&subscripts)
4133 .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4134 };
4135 Self::from_inner_with_axis_classes(
4136 self.indices.clone(),
4137 scaled,
4138 self.storage.axis_classes().to_vec(),
4139 )
4140 .map_err(IdxTensorError::from)
4141 }
4142
4143 pub fn inner_product(&self, other: &Self) -> std::result::Result<AnyScalar, IdxTensorError> {
4164 if self.indices.len() == other.indices.len() {
4165 let self_set: HashSet<_> = self.indices.iter().collect();
4166 let other_set: HashSet<_> = other.indices.iter().collect();
4167 if self_set == other_set {
4168 let other_aligned = other.permute_indices(&self.indices)?;
4169 let result = super::contract::contract_pair_with_operand_options(
4170 self,
4171 &other_aligned,
4172 PairwiseContractionOptions::new().with_lhs_conj(true),
4173 )?;
4174 return result.sum();
4175 }
4176 }
4177
4178 let result = super::contract::contract_pair_with_operand_options(
4180 self,
4181 other,
4182 PairwiseContractionOptions::new().with_lhs_conj(true),
4183 )?;
4184 result.sum()
4186 }
4187}
4188
4189impl IdxTensor {
4194 pub fn replaceind(
4229 &self,
4230 old_index: &DynIndex,
4231 new_index: &DynIndex,
4232 ) -> std::result::Result<Self, IdxTensorError> {
4233 if old_index.dim() != new_index.dim() {
4235 return Err(IdxTensorError::ShapeMismatch {
4236 operation: "replaceind",
4237 expected: format!("dimension {}", old_index.dim()),
4238 actual: format!("dimension {}", new_index.dim()),
4239 });
4240 }
4241
4242 let new_indices: Vec<_> = self
4243 .indices
4244 .iter()
4245 .map(|idx| {
4246 if *idx == *old_index {
4247 new_index.clone()
4248 } else {
4249 idx.clone()
4250 }
4251 })
4252 .collect();
4253
4254 Ok(Self {
4255 indices: new_indices,
4256 storage: self.storage.clone(),
4257 eager_cache: Arc::clone(&self.eager_cache),
4258 })
4259 }
4260
4261 pub fn replace_indices(
4301 &self,
4302 old_indices: &[DynIndex],
4303 new_indices: &[DynIndex],
4304 ) -> std::result::Result<Self, IdxTensorError> {
4305 if old_indices.len() != new_indices.len() {
4306 return Err(IdxTensorError::ShapeMismatch {
4307 operation: "replace_indices",
4308 expected: format!("{} indices", old_indices.len()),
4309 actual: format!("{} indices", new_indices.len()),
4310 });
4311 }
4312
4313 for (old, new) in old_indices.iter().zip(new_indices.iter()) {
4315 if old.dim() != new.dim() {
4316 return Err(IdxTensorError::ShapeMismatch {
4317 operation: "replace_indices",
4318 expected: format!("dimension {}", old.dim()),
4319 actual: format!("dimension {}", new.dim()),
4320 });
4321 }
4322 }
4323
4324 let replacement_map: std::collections::HashMap<_, _> =
4326 old_indices.iter().zip(new_indices.iter()).collect();
4327
4328 let new_indices_vec: Vec<_> = self
4329 .indices
4330 .iter()
4331 .map(|idx| {
4332 if let Some(new_idx) = replacement_map.get(idx) {
4333 (*new_idx).clone()
4334 } else {
4335 idx.clone()
4336 }
4337 })
4338 .collect();
4339
4340 Ok(Self {
4341 indices: new_indices_vec,
4342 storage: self.storage.clone(),
4343 eager_cache: Arc::clone(&self.eager_cache),
4344 })
4345 }
4346}
4347
4348impl IdxTensor {
4353 pub fn conj(&self) -> Self {
4382 self.conj_with(&conjugate_eager)
4383 }
4384
4385 fn conj_with<F>(&self, conjugate: &F) -> Self
4386 where
4387 F: Fn(
4388 &EagerTensor,
4389 ) -> std::result::Result<
4390 EagerTensor,
4391 Arc<dyn std::error::Error + Send + Sync + 'static>,
4392 >,
4393 {
4394 let new_indices: Vec<DynIndex> = self.indices.iter().map(|idx| idx.conj()).collect();
4398 let mut storage = match self.storage.conjugate_with(conjugate) {
4399 Ok(storage) => storage,
4400 Err(error) => self.storage.clone().with_deferred_error(error),
4401 };
4402 let mut eager_cache = if storage.deferred_error().is_some() {
4403 Arc::clone(&self.eager_cache)
4404 } else {
4405 Self::empty_eager_cache()
4406 };
4407
4408 if storage.deferred_error().is_none() {
4409 if let Some(inner) = self.eager_cache.get() {
4410 match conjugate(inner.as_ref()) {
4411 Ok(conjugated) => eager_cache = Self::eager_cache_with(conjugated),
4412 Err(source) => {
4413 storage =
4414 storage.with_deferred_error(TensorStorageError::Conjugation { source });
4415 eager_cache = Arc::clone(&self.eager_cache);
4418 }
4419 }
4420 }
4421 }
4422
4423 Self {
4424 indices: new_indices,
4425 storage,
4426 eager_cache,
4427 }
4428 }
4429}
4430
4431#[derive(Debug, Default)]
4432struct Lassq {
4433 scale: f64,
4434 sumsq: f64,
4435 infinite: bool,
4436}
4437
4438impl Lassq {
4439 fn add_component(&mut self, value: f64) {
4440 let value = value.abs();
4441 if value == 0.0 {
4442 return;
4443 }
4444 if value.is_infinite() {
4445 self.infinite = true;
4446 return;
4447 }
4448 if self.scale < value {
4449 if self.scale == 0.0 {
4450 self.sumsq = 1.0;
4451 } else {
4452 let ratio = self.scale / value;
4453 self.sumsq = 1.0 + self.sumsq * ratio * ratio;
4454 }
4455 self.scale = value;
4456 } else {
4457 let ratio = value / self.scale;
4458 self.sumsq += ratio * ratio;
4459 }
4460 }
4461
4462 fn add_complex(&mut self, value: Complex64) {
4463 self.add_component(value.re);
4464 self.add_component(value.im);
4465 }
4466
4467 fn add_scaled(&mut self, scale: f64, coefficient: f64) {
4468 if scale == 0.0 || coefficient == 0.0 {
4469 return;
4470 }
4471 if self.scale < scale {
4472 if self.scale == 0.0 {
4473 self.sumsq = coefficient * coefficient;
4474 } else {
4475 let ratio = self.scale / scale;
4476 self.sumsq = coefficient * coefficient + self.sumsq * ratio * ratio;
4477 }
4478 self.scale = scale;
4479 } else {
4480 let ratio = scale / self.scale * coefficient;
4481 self.sumsq += ratio * ratio;
4482 }
4483 }
4484
4485 fn add_component_difference(&mut self, lhs: f64, rhs: f64) {
4486 if lhs == rhs {
4487 return;
4488 }
4489 let scale = lhs.abs().max(rhs.abs());
4490 if scale != 0.0 {
4491 self.add_scaled(scale, (lhs / scale - rhs / scale).abs());
4492 }
4493 }
4494
4495 fn add_complex_difference(&mut self, lhs: Complex64, rhs: Complex64) {
4496 self.add_component_difference(lhs.re, rhs.re);
4497 self.add_component_difference(lhs.im, rhs.im);
4498 }
4499
4500 fn is_zero(&self) -> bool {
4501 !self.infinite && self.scale == 0.0
4502 }
4503
4504 fn norm(&self) -> f64 {
4505 if self.infinite {
4506 f64::INFINITY
4507 } else if self.scale == 0.0 {
4508 0.0
4509 } else {
4510 self.scale * self.sumsq.sqrt()
4511 }
4512 }
4513
4514 fn norm_squared(&self) -> f64 {
4515 let norm = self.norm();
4516 norm * norm
4517 }
4518
4519 fn log_norm(&self) -> f64 {
4520 if self.infinite {
4521 f64::INFINITY
4522 } else if self.scale == 0.0 {
4523 f64::NEG_INFINITY
4524 } else {
4525 self.scale.ln() + 0.5 * self.sumsq.ln()
4526 }
4527 }
4528}
4529
4530impl IdxTensor {
4535 pub fn norm_squared(&self) -> std::result::Result<f64, IdxTensorError> {
4561 let dtype = self
4562 .scalar_dtype()
4563 .map_err(IdxTensorError::scalar_extraction)?;
4564 if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4565 return Err(IdxTensorError::ScalarTypeMismatch {
4566 expected: "f32, f64, c32, or c64",
4567 actual: Self::dtype_name(dtype).to_string(),
4568 });
4569 }
4570 let (has_nan, _) = self
4571 .compact_nonfinite_flags()
4572 .map_err(IdxTensorError::materialization)?;
4573 if has_nan {
4574 return Err(IdxTensorError::NaNInput {
4575 operation: "norm_squared",
4576 });
4577 }
4578
4579 let mut norm = Lassq::default();
4580 self.storage
4581 .for_each_payload_value(|value| norm.add_complex(value))
4582 .map_err(IdxTensorError::materialization)?;
4583 let value = norm.norm_squared();
4584 if value.is_nan() {
4585 return Err(IdxTensorError::NaNInput {
4586 operation: "norm_squared",
4587 });
4588 }
4589 Ok(value)
4590 }
4591
4592 pub fn norm(&self) -> std::result::Result<f64, IdxTensorError> {
4610 Ok(self.norm_squared()?.sqrt())
4611 }
4612
4613 pub fn maxabs(&self) -> std::result::Result<f64, IdxTensorError> {
4629 if let Some(error) = self.storage.deferred_error() {
4630 return Err(IdxTensorError::Storage {
4631 source: error.clone(),
4632 });
4633 }
4634 let dtype = self
4635 .storage
4636 .dtype()
4637 .ok_or_else(|| IdxTensorError::ScalarTypeMismatch {
4638 expected: "f32, f64, c32, or c64",
4639 actual: "unknown".to_string(),
4640 })?;
4641 if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4642 return Err(IdxTensorError::ScalarTypeMismatch {
4643 expected: "f32, f64, c32, or c64",
4644 actual: Self::dtype_name(dtype).to_string(),
4645 });
4646 }
4647 let (has_nan, _) = self
4648 .compact_nonfinite_flags()
4649 .map_err(IdxTensorError::materialization)?;
4650 if has_nan {
4651 return Err(IdxTensorError::NaNInput {
4652 operation: "maxabs",
4653 });
4654 }
4655 let mut value = 0.0_f64;
4656 self.storage
4657 .for_each_payload_value(|scalar| {
4658 let magnitude = scalar.re.hypot(scalar.im);
4659 value = value.max(magnitude);
4660 })
4661 .map_err(IdxTensorError::materialization)?;
4662 Ok(value)
4663 }
4664
4665 fn native_complex_payload_value_at(
4666 native: &EagerTensor,
4667 payload_coords: &[usize],
4668 ) -> Result<Complex64> {
4669 if !(payload_coords.len() == native.shape().len()) {
4670 return Err(anyhow::anyhow!(
4671 "payload coordinate rank {} does not match payload rank {}",
4672 payload_coords.len(),
4673 native.shape().len()
4674 ));
4675 };
4676 for (&coordinate, &dim) in payload_coords.iter().zip(native.shape().iter()) {
4677 if coordinate >= dim {
4678 return Err(anyhow::anyhow!(
4679 "payload coordinate {coordinate} is out of bounds for dim {dim}"
4680 ));
4681 }
4682 }
4683 let value = native.value()?;
4684 match value.as_tensor_view() {
4685 TensorView::F32(view) => view
4686 .get(payload_coords)
4687 .copied()
4688 .map(|value| Complex64::new(f64::from(value), 0.0))
4689 .ok_or_else(|| anyhow::anyhow!("failed to read f32 payload value")),
4690 TensorView::F64(view) => view
4691 .get(payload_coords)
4692 .copied()
4693 .map(|value| Complex64::new(value, 0.0))
4694 .ok_or_else(|| anyhow::anyhow!("failed to read f64 payload value")),
4695 TensorView::C32(view) => view
4696 .get(payload_coords)
4697 .copied()
4698 .map(|value| Complex64::new(f64::from(value.re), f64::from(value.im)))
4699 .ok_or_else(|| anyhow::anyhow!("failed to read c32 payload value")),
4700 TensorView::C64(view) => view
4701 .get(payload_coords)
4702 .copied()
4703 .ok_or_else(|| anyhow::anyhow!("failed to read c64 payload value")),
4704 view => Err(anyhow::anyhow!(
4705 "unsupported payload dtype {:?}",
4706 view.dtype()
4707 )),
4708 }
4709 }
4710
4711 fn native_sum_scalar(native: &EagerTensor) -> Result<AnyScalar> {
4712 let value = native.value()?;
4713 match native.dtype() {
4714 DType::F32 => Ok(AnyScalar::from_value(
4715 value.as_slice::<f32>()?.iter().copied().sum::<f32>(),
4716 )),
4717 DType::F64 => Ok(AnyScalar::from_value(
4718 value.as_slice::<f64>()?.iter().copied().sum::<f64>(),
4719 )),
4720 DType::C32 => Ok(AnyScalar::from_value(
4721 value
4722 .as_slice::<Complex32>()?
4723 .iter()
4724 .copied()
4725 .sum::<Complex32>(),
4726 )),
4727 DType::C64 => Ok(AnyScalar::from_value(
4728 value
4729 .as_slice::<Complex64>()?
4730 .iter()
4731 .copied()
4732 .sum::<Complex64>(),
4733 )),
4734 dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
4735 }
4736 }
4737
4738 fn native_nonfinite_flags_typed<T: TensorElement>(
4739 native: &EagerTensor,
4740 classify: impl Fn(T) -> (bool, bool),
4741 ) -> Result<(bool, bool)> {
4742 let value = native.value()?;
4743 if let Ok(values) = value.as_slice::<T>() {
4744 return Ok(values.iter().copied().map(&classify).fold(
4745 (false, false),
4746 |(has_nan, has_infinity), (is_nan, is_infinite)| {
4747 (has_nan | is_nan, has_infinity | is_infinite)
4748 },
4749 ));
4750 }
4751 drop(value);
4752 let value = IdxTensorStorage::materialize_eager_payload(native)?;
4753 Ok(value.as_slice::<T>()?.iter().copied().map(classify).fold(
4754 (false, false),
4755 |(has_nan, has_infinity), (is_nan, is_infinite)| {
4756 (has_nan | is_nan, has_infinity | is_infinite)
4757 },
4758 ))
4759 }
4760
4761 fn native_nonfinite_flags(native: &EagerTensor) -> Result<(bool, bool)> {
4762 match native.dtype() {
4763 DType::F32 => Self::native_nonfinite_flags_typed(native, |value: f32| {
4764 (value.is_nan(), value.is_infinite())
4765 }),
4766 DType::F64 => Self::native_nonfinite_flags_typed(native, |value: f64| {
4767 (value.is_nan(), value.is_infinite())
4768 }),
4769 DType::C32 => Self::native_nonfinite_flags_typed(native, |value: Complex32| {
4770 (
4771 value.re.is_nan() || value.im.is_nan(),
4772 value.re.is_infinite() || value.im.is_infinite(),
4773 )
4774 }),
4775 DType::C64 => Self::native_nonfinite_flags_typed(native, |value: Complex64| {
4776 (
4777 value.re.is_nan() || value.im.is_nan(),
4778 value.re.is_infinite() || value.im.is_infinite(),
4779 )
4780 }),
4781 dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
4782 }
4783 }
4784
4785 fn compact_nonfinite_flags(&self) -> Result<(bool, bool)> {
4786 self.storage.nonfinite_flags()
4787 }
4788
4789 pub fn sub(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
4799 self.axpby(AnyScalar::new_real(1.0), other, AnyScalar::new_real(-1.0))
4800 }
4801
4802 pub fn neg(&self) -> std::result::Result<Self, IdxTensorError> {
4809 self.scale(AnyScalar::new_real(-1.0))
4810 }
4811
4812 pub fn isapprox(
4824 &self,
4825 other: &Self,
4826 atol: f64,
4827 rtol: f64,
4828 ) -> std::result::Result<bool, IdxTensorError> {
4829 for (name, value) in [("atol", atol), ("rtol", rtol)] {
4830 if !value.is_finite() || value < 0.0 {
4831 return Err(IdxTensorError::InvalidTolerance { name, value });
4832 }
4833 }
4834 if self.indices.len() != other.indices.len() {
4835 return Err(IdxTensorError::ShapeMismatch {
4836 operation: "isapprox",
4837 expected: format!("indices {:?}", self.indices),
4838 actual: format!("indices {:?}", other.indices),
4839 });
4840 }
4841
4842 let other_axis_by_index = other
4843 .indices
4844 .iter()
4845 .cloned()
4846 .enumerate()
4847 .map(|(axis, index)| (index, axis))
4848 .collect::<HashMap<_, _>>();
4849 let other_positions = self
4850 .indices
4851 .iter()
4852 .map(|index| {
4853 other_axis_by_index.get(index).copied().ok_or_else(|| {
4854 IdxTensorError::ShapeMismatch {
4855 operation: "isapprox",
4856 expected: format!("indices {:?}", self.indices),
4857 actual: format!("indices {:?}", other.indices),
4858 }
4859 })
4860 })
4861 .collect::<std::result::Result<Vec<_>, _>>()?;
4862 let self_dims = self.dims();
4863 let other_dims = other.dims();
4864 for (axis, &other_axis) in other_positions.iter().enumerate() {
4865 if self_dims[axis] != other_dims[other_axis] {
4866 return Err(IdxTensorError::ShapeMismatch {
4867 operation: "isapprox",
4868 expected: format!("dims {:?}", self_dims),
4869 actual: format!("dims {:?}", other_dims),
4870 });
4871 }
4872 }
4873 for tensor in [self, other] {
4874 if tensor
4875 .compact_nonfinite_flags()
4876 .map_err(IdxTensorError::materialization)?
4877 .0
4878 {
4879 return Err(IdxTensorError::NaNInput {
4880 operation: "isapprox",
4881 });
4882 }
4883 }
4884
4885 let exact = atol == 0.0 && rtol == 0.0;
4886 let lhs_payload_dims = self.storage.payload_dims().to_vec();
4887 let rhs_payload_dims = other.storage.payload_dims().to_vec();
4888 let lhs_payload_len =
4889 checked_product(&lhs_payload_dims).map_err(IdxTensorError::materialization)?;
4890 let rhs_payload_len =
4891 checked_product(&rhs_payload_dims).map_err(IdxTensorError::materialization)?;
4892 let lhs_axis_classes = self.storage.axis_classes();
4893 let rhs_axis_classes = other.storage.axis_classes();
4894 let self_to_other = other_positions.clone();
4895 let mut other_to_self = vec![0usize; self_to_other.len()];
4896 for (self_axis, &other_axis) in self_to_other.iter().enumerate() {
4897 other_to_self[other_axis] = self_axis;
4898 }
4899
4900 let mut diff = Lassq::default();
4901 let mut lhs_norm = Lassq::default();
4902 let mut rhs_norm = Lassq::default();
4903 let mut compare = |lhs: Complex64, rhs: Complex64| -> bool {
4904 if exact {
4905 return lhs == rhs;
4906 }
4907 let lhs_infinite = lhs.re.is_infinite() || lhs.im.is_infinite();
4908 let rhs_infinite = rhs.re.is_infinite() || rhs.im.is_infinite();
4909 if lhs_infinite || rhs_infinite {
4910 return lhs == rhs;
4911 }
4912 lhs_norm.add_complex(lhs);
4913 rhs_norm.add_complex(rhs);
4914 diff.add_complex_difference(lhs, rhs);
4915 true
4916 };
4917
4918 let mut lhs_coords = vec![0usize; lhs_payload_dims.len()];
4922 let mut rhs_from_lhs = vec![0usize; rhs_payload_dims.len()];
4923 let mut rhs_seen = vec![false; rhs_payload_dims.len()];
4924 for _ in 0..lhs_payload_len {
4925 let lhs = self
4926 .storage
4927 .payload_value_at(&lhs_coords)
4928 .map_err(IdxTensorError::materialization)?;
4929 if map_payload_support_coordinate(
4930 lhs_axis_classes,
4931 rhs_axis_classes,
4932 &self_to_other,
4933 &lhs_coords,
4934 &mut rhs_from_lhs,
4935 &mut rhs_seen,
4936 ) {
4937 let rhs = other
4938 .storage
4939 .payload_value_at(&rhs_from_lhs)
4940 .map_err(IdxTensorError::materialization)?;
4941 if !compare(lhs, rhs) {
4942 return Ok(false);
4943 }
4944 } else if !compare(lhs, Complex64::new(0.0, 0.0)) {
4945 return Ok(false);
4946 }
4947
4948 increment_col_major_coordinate(&mut lhs_coords, &lhs_payload_dims);
4949 }
4950
4951 let mut rhs_coords = vec![0usize; rhs_payload_dims.len()];
4955 let mut lhs_from_rhs = vec![0usize; lhs_payload_dims.len()];
4956 let mut lhs_seen = vec![false; lhs_payload_dims.len()];
4957 for _ in 0..rhs_payload_len {
4958 let rhs = other
4959 .storage
4960 .payload_value_at(&rhs_coords)
4961 .map_err(IdxTensorError::materialization)?;
4962 if !map_payload_support_coordinate(
4963 rhs_axis_classes,
4964 lhs_axis_classes,
4965 &other_to_self,
4966 &rhs_coords,
4967 &mut lhs_from_rhs,
4968 &mut lhs_seen,
4969 ) && !compare(Complex64::new(0.0, 0.0), rhs)
4970 {
4971 return Ok(false);
4972 }
4973 increment_col_major_coordinate(&mut rhs_coords, &rhs_payload_dims);
4974 }
4975
4976 if exact {
4977 return Ok(true);
4978 }
4979 let absolute_ok = if diff.is_zero() {
4980 true
4981 } else if atol == 0.0 || diff.infinite {
4982 false
4983 } else {
4984 diff.log_norm() <= atol.ln()
4985 };
4986 let relative_ok = if rtol == 0.0 || diff.is_zero() {
4987 diff.is_zero()
4988 } else if diff.infinite {
4989 lhs_norm.infinite || rhs_norm.infinite
4990 } else if lhs_norm.infinite || rhs_norm.infinite {
4991 true
4992 } else {
4993 let reference_log = lhs_norm.log_norm().max(rhs_norm.log_norm());
4994 diff.log_norm() <= rtol.ln() + reference_log
4995 };
4996 Ok(absolute_ok || relative_ok)
4997 }
4998
4999 pub fn diagonal(
5006 input_index: &DynIndex,
5007 output_index: &DynIndex,
5008 ) -> std::result::Result<Self, IdxTensorError> {
5009 <Self as TensorConstructionLike>::diagonal(input_index, output_index)
5010 }
5011
5012 pub fn delta(
5018 input_indices: &[DynIndex],
5019 output_indices: &[DynIndex],
5020 ) -> std::result::Result<Self, IdxTensorError> {
5021 <Self as TensorConstructionLike>::delta(input_indices, output_indices)
5022 }
5023
5024 pub fn scalar_one() -> std::result::Result<Self, IdxTensorError> {
5031 <Self as TensorConstructionLike>::scalar_one()
5032 }
5033
5034 pub fn ones(indices: &[DynIndex]) -> std::result::Result<Self, IdxTensorError> {
5041 <Self as TensorConstructionLike>::ones(indices)
5042 }
5043
5044 pub fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, IdxTensorError> {
5051 <Self as TensorConstructionLike>::onehot(index_vals)
5052 }
5053
5054 pub fn mask_index(
5093 &self,
5094 index: &DynIndex,
5095 position: usize,
5096 ) -> std::result::Result<Self, IdxTensorError> {
5097 if !(self.indices.iter().any(|candidate| candidate == index)) {
5098 return Err(anyhow::anyhow!("mask_index: index is not present in tensor").into());
5099 };
5100 if !(position < index.dim()) {
5101 return Err(anyhow::anyhow!(
5102 "mask_index: position {position} is out of range for dimension {}",
5103 index.dim()
5104 )
5105 .into());
5106 };
5107
5108 let mask = match self.scalar_dtype()? {
5113 DType::F32 => Self::from_dense(
5114 vec![index.clone()],
5115 (0..index.dim())
5116 .map(|value| if value == position { 1.0_f32 } else { 0.0 })
5117 .collect(),
5118 ),
5119 DType::F64 => Self::from_dense(
5120 vec![index.clone()],
5121 (0..index.dim())
5122 .map(|value| if value == position { 1.0_f64 } else { 0.0 })
5123 .collect(),
5124 ),
5125 DType::C32 => Self::from_dense(
5126 vec![index.clone()],
5127 (0..index.dim())
5128 .map(|value| {
5129 if value == position {
5130 num_complex::Complex32::new(1.0, 0.0)
5131 } else {
5132 num_complex::Complex32::new(0.0, 0.0)
5133 }
5134 })
5135 .collect(),
5136 ),
5137 DType::C64 => Self::from_dense(
5138 vec![index.clone()],
5139 (0..index.dim())
5140 .map(|value| {
5141 if value == position {
5142 Complex64::new(1.0, 0.0)
5143 } else {
5144 Complex64::new(0.0, 0.0)
5145 }
5146 })
5147 .collect(),
5148 ),
5149 dtype => {
5150 return Err(anyhow::anyhow!("mask_index does not support dtype {dtype:?}").into())
5151 }
5152 }?;
5153 super::contract::contract_pair_with_options(
5154 self,
5155 &mask,
5156 super::contract::ContractionOptions::new()
5157 .with_retain_indices(std::slice::from_ref(index)),
5158 )
5159 }
5160
5161 pub fn distance(&self, other: &Self) -> std::result::Result<f64, IdxTensorError> {
5196 let norm_self = self.norm()?;
5197
5198 let neg_other = other.scale(AnyScalar::new_real(-1.0))?;
5200 let diff = self.add(&neg_other)?;
5201 let norm_diff = diff.norm()?;
5202
5203 if norm_self > 0.0 {
5204 Ok(norm_diff / norm_self)
5205 } else {
5206 Ok(norm_diff)
5207 }
5208 }
5209}
5210
5211impl std::fmt::Debug for IdxTensor {
5212 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5213 f.debug_struct("IdxTensor")
5214 .field("indices", &self.indices)
5215 .field("dims", &self.dims())
5216 .field("is_diag", &self.is_diag())
5217 .finish()
5218 }
5219}
5220
5221pub fn diag_idx_tensor(
5245 indices: Vec<DynIndex>,
5246 diag_data: Vec<f64>,
5247) -> std::result::Result<IdxTensor, IdxTensorError> {
5248 IdxTensor::from_diag(indices, diag_data)
5249}
5250
5251#[allow(clippy::type_complexity)]
5252pub(crate) type UnfoldSplitInnerResult = (
5253 EagerTensor,
5254 usize,
5255 usize,
5256 usize,
5257 Vec<DynIndex>,
5258 Vec<DynIndex>,
5259);
5260
5261#[allow(clippy::type_complexity)]
5301pub fn unfold_split(
5302 t: &IdxTensor,
5303 left_inds: &[DynIndex],
5304) -> std::result::Result<
5305 (
5306 NativeTensor,
5307 usize,
5308 usize,
5309 usize,
5310 Vec<DynIndex>,
5311 Vec<DynIndex>,
5312 ),
5313 IdxTensorError,
5314> {
5315 let (matrix_inner, left_len, m, n, left_indices, right_indices) =
5316 unfold_split_inner(t, left_inds)?;
5317
5318 Ok((
5319 matrix_inner.duplicate_value()?,
5320 left_len,
5321 m,
5322 n,
5323 left_indices,
5324 right_indices,
5325 ))
5326}
5327
5328pub(crate) fn unfold_split_inner(
5329 t: &IdxTensor,
5330 left_inds: &[DynIndex],
5331) -> Result<UnfoldSplitInnerResult> {
5332 let rank = t.indices.len();
5333
5334 if !(rank >= 2) {
5336 return Err(anyhow::anyhow!(
5337 "Tensor must have rank >= 2, got rank {}",
5338 rank
5339 ));
5340 };
5341
5342 let left_len = left_inds.len();
5343
5344 if !(left_len > 0 && left_len < rank) {
5346 return Err(anyhow::anyhow!("Left indices must be a non-empty proper subset of tensor indices (0 < left_len < rank), got left_len={}, rank={}",
5347 left_len,
5348 rank));
5349 };
5350
5351 let tensor_set: HashSet<_> = t.indices.iter().collect();
5353 let mut left_set = HashSet::new();
5354
5355 for left_idx in left_inds {
5356 if !(tensor_set.contains(left_idx)) {
5357 return Err(anyhow::anyhow!("Index in left_inds not found in tensor"));
5358 };
5359 if !(left_set.insert(left_idx)) {
5360 return Err(anyhow::anyhow!("Duplicate index in left_inds"));
5361 };
5362 }
5363
5364 let mut right_inds = Vec::new();
5366 for idx in &t.indices {
5367 if !left_set.contains(idx) {
5368 right_inds.push(idx.clone());
5369 }
5370 }
5371
5372 let mut new_indices = Vec::with_capacity(rank);
5374 new_indices.extend_from_slice(left_inds);
5375 new_indices.extend_from_slice(&right_inds);
5376
5377 let unfolded = t.permute_indices(&new_indices)?;
5379
5380 let unfolded_dims = unfolded.dims();
5382 let m = checked_product(&unfolded_dims[..left_len])?;
5383 let n = checked_product(&unfolded_dims[left_len..])?;
5384
5385 let matrix_tensor = unfolded.try_materialized_inner()?.reshape(&[m, n])?;
5386
5387 Ok((
5388 matrix_tensor,
5389 left_len,
5390 m,
5391 n,
5392 left_inds.to_vec(),
5393 right_inds,
5394 ))
5395}
5396
5397use crate::tensor_index::TensorIndex;
5402
5403impl TensorIndex for IdxTensor {
5404 type Index = DynIndex;
5405 type Error = IdxTensorError;
5406
5407 fn external_indices(&self) -> Vec<DynIndex> {
5408 self.indices.clone()
5410 }
5411
5412 fn num_external_indices(&self) -> usize {
5413 self.indices.len()
5414 }
5415
5416 fn replaceind(
5417 &self,
5418 old_index: &DynIndex,
5419 new_index: &DynIndex,
5420 ) -> std::result::Result<Self, Self::Error> {
5421 IdxTensor::replaceind(self, old_index, new_index)
5423 }
5424
5425 fn replace_indices(
5426 &self,
5427 old_indices: &[DynIndex],
5428 new_indices: &[DynIndex],
5429 ) -> std::result::Result<Self, Self::Error> {
5430 IdxTensor::replace_indices(self, old_indices, new_indices)
5432 }
5433}
5434
5435use crate::tensor_like::{
5440 FactorizeError, FactorizeOptions, FactorizeResult, TensorConstructionLike,
5441 TensorContractionLike, TensorFactorizationLike, TensorVectorSpace,
5442};
5443
5444impl TensorVectorSpace for IdxTensor {
5445 fn norm_squared(&self) -> std::result::Result<f64, Self::Error> {
5446 IdxTensor::norm_squared(self)
5447 }
5448
5449 fn maxabs(&self) -> std::result::Result<f64, Self::Error> {
5450 IdxTensor::maxabs(self)
5451 }
5452
5453 fn isapprox(
5454 &self,
5455 other: &Self,
5456 atol: f64,
5457 rtol: f64,
5458 ) -> std::result::Result<bool, Self::Error> {
5459 IdxTensor::isapprox(self, other, atol, rtol)
5460 }
5461
5462 fn axpby(
5463 &self,
5464 a: crate::AnyScalar,
5465 other: &Self,
5466 b: crate::AnyScalar,
5467 ) -> std::result::Result<Self, Self::Error> {
5468 IdxTensor::axpby(self, a, other, b)
5469 }
5470
5471 fn scale(&self, scalar: crate::AnyScalar) -> std::result::Result<Self, Self::Error> {
5472 IdxTensor::scale(self, scalar)
5473 }
5474
5475 fn inner_product(&self, other: &Self) -> std::result::Result<crate::AnyScalar, Self::Error> {
5476 IdxTensor::inner_product(self, other)
5477 }
5478}
5479
5480impl TensorFactorizationLike for IdxTensor {
5481 fn factorize(
5482 &self,
5483 left_inds: &[DynIndex],
5484 options: &FactorizeOptions,
5485 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5486 crate::factorize::factorize(self, left_inds, options)
5487 }
5488
5489 fn factorize_auto(
5490 &self,
5491 left_inds: &[DynIndex],
5492 options: &FactorizeOptions,
5493 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5494 crate::factorize::factorize_auto(self, left_inds, options)
5495 }
5496
5497 fn factorize_full_rank(
5498 &self,
5499 left_inds: &[DynIndex],
5500 alg: crate::FactorizeAlg,
5501 canonical: crate::Canonical,
5502 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5503 crate::factorize::factorize_full_rank(self, left_inds, alg, canonical)
5504 }
5505}
5506
5507impl TensorContractionLike for IdxTensor {
5508 fn conj(&self) -> Self {
5509 IdxTensor::conj(self)
5511 }
5512
5513 fn direct_sum(
5514 &self,
5515 other: &Self,
5516 pairs: &[(DynIndex, DynIndex)],
5517 ) -> std::result::Result<crate::tensor_like::DirectSumResult<Self>, Self::Error> {
5518 let (tensor, new_indices) = crate::direct_sum::direct_sum(self, other, pairs)?;
5519 Ok(crate::tensor_like::DirectSumResult {
5520 tensor,
5521 new_indices,
5522 })
5523 }
5524
5525 fn outer_product(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
5526 super::contract::outer_product(self, other)
5527 }
5528
5529 fn permuteinds(&self, new_order: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
5530 IdxTensor::permute_indices(self, new_order)
5532 }
5533
5534 fn fuse_indices(
5535 &self,
5536 old_indices: &[DynIndex],
5537 new_index: DynIndex,
5538 order: LinearizationOrder,
5539 ) -> std::result::Result<Self, Self::Error> {
5540 IdxTensor::fuse_indices(self, old_indices, new_index, order)
5541 }
5542
5543 fn contract(tensors: &[&Self]) -> std::result::Result<Self, Self::Error> {
5544 super::contract::contract(tensors)
5545 }
5546
5547 fn contract_pair(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
5548 super::contract::contract_pair(self, other)
5549 }
5550}
5551
5552impl TensorConstructionLike for IdxTensor {
5553 fn select_indices(
5554 &self,
5555 selected_indices: &[DynIndex],
5556 positions: &[usize],
5557 ) -> std::result::Result<Self, Self::Error> {
5558 IdxTensor::select_indices(self, selected_indices, positions)
5559 }
5560
5561 fn diagonal(
5562 input_index: &DynIndex,
5563 output_index: &DynIndex,
5564 ) -> std::result::Result<Self, Self::Error> {
5565 let dim = input_index.dim();
5566 if dim != output_index.dim() {
5567 return Err(anyhow::anyhow!(
5568 "Dimension mismatch: input index has dim {}, output has dim {}",
5569 dim,
5570 output_index.dim(),
5571 )
5572 .into());
5573 }
5574
5575 IdxTensor::from_diag(
5576 vec![input_index.clone(), output_index.clone()],
5577 vec![1.0_f64; dim],
5578 )
5579 }
5580
5581 fn scalar_one() -> std::result::Result<Self, Self::Error> {
5582 IdxTensor::from_dense(vec![], vec![1.0_f64])
5583 }
5584
5585 fn ones(indices: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
5586 if indices.is_empty() {
5587 return <Self as TensorConstructionLike>::scalar_one();
5588 }
5589 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
5590 let total_size = checked_total_size(&dims)?;
5591 IdxTensor::from_dense(indices.to_vec(), vec![1.0_f64; total_size])
5592 }
5593
5594 fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, Self::Error> {
5595 if index_vals.is_empty() {
5596 return <Self as TensorConstructionLike>::scalar_one();
5597 }
5598 let indices: Vec<DynIndex> = index_vals.iter().map(|(idx, _)| idx.clone()).collect();
5599 let vals: Vec<usize> = index_vals.iter().map(|(_, v)| *v).collect();
5600 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
5601
5602 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
5603 if v >= d {
5604 return Err(anyhow::anyhow!(
5605 "onehot: value {} at position {} is >= dimension {}",
5606 v,
5607 k,
5608 d
5609 )
5610 .into());
5611 }
5612 }
5613
5614 let total_size = checked_total_size(&dims).map_err(Self::Error::from)?;
5615 let mut data = vec![0.0_f64; total_size];
5616
5617 let offset = column_major_offset(&dims, &vals).map_err(Self::Error::from)?;
5618 data[offset] = 1.0;
5619
5620 Self::from_dense(indices, data)
5621 }
5622
5623 }
5625
5626fn checked_total_size(dims: &[usize]) -> Result<usize> {
5627 dims.iter().try_fold(1_usize, |acc, &d| {
5628 if d == 0 {
5629 return Err(anyhow::anyhow!("invalid dimension 0"));
5630 }
5631 acc.checked_mul(d)
5632 .ok_or_else(|| anyhow::anyhow!("tensor size overflow"))
5633 })
5634}
5635
5636fn column_major_offset(dims: &[usize], vals: &[usize]) -> Result<usize> {
5637 if dims.len() != vals.len() {
5638 return Err(anyhow::anyhow!(
5639 "column_major_offset: dims.len() != vals.len()"
5640 ));
5641 }
5642 checked_total_size(dims)?;
5643
5644 let mut offset = 0usize;
5645 let mut stride = 1usize;
5646 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
5647 if d == 0 {
5648 return Err(anyhow::anyhow!("invalid dimension 0 at position {}", k));
5649 }
5650 if v >= d {
5651 return Err(anyhow::anyhow!(
5652 "column_major_offset: value {} at position {} is >= dimension {}",
5653 v,
5654 k,
5655 d
5656 ));
5657 }
5658 let term = v
5659 .checked_mul(stride)
5660 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
5661 offset = offset
5662 .checked_add(term)
5663 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
5664 stride = stride
5665 .checked_mul(d)
5666 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
5667 }
5668 Ok(offset)
5669}
5670
5671impl IdxTensor {
5676 fn any_scalar_payload_to_complex(data: Vec<AnyScalar>) -> Vec<Complex64> {
5677 data.into_iter()
5678 .map(|value| {
5679 value
5680 .as_c64()
5681 .unwrap_or_else(|| Complex64::new(value.real(), 0.0))
5682 })
5683 .collect()
5684 }
5685
5686 fn any_scalar_payload_to_real(data: Vec<AnyScalar>) -> Vec<f64> {
5687 data.into_iter().map(|value| value.real()).collect()
5688 }
5689
5690 fn validate_dense_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
5691 let expected_len = checked_total_size(dims)?;
5692 if !(data_len == expected_len) {
5693 return Err(anyhow::anyhow!(
5694 "dense payload length {} does not match dims {:?} (expected {})",
5695 data_len,
5696 dims,
5697 expected_len
5698 ));
5699 };
5700 Ok(())
5701 }
5702
5703 fn validate_diag_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
5704 if !(!dims.is_empty()) {
5705 return Err(anyhow::anyhow!(
5706 "diagonal tensor construction requires at least one index"
5707 ));
5708 };
5709 Self::validate_diag_dims(dims)?;
5710 if !(data_len == dims[0]) {
5711 return Err(anyhow::anyhow!(
5712 "diagonal payload length {} does not match diagonal dimension {}",
5713 data_len,
5714 dims[0]
5715 ));
5716 };
5717 Ok(())
5718 }
5719
5720 pub fn from_dense<T: TensorElement>(
5747 indices: Vec<DynIndex>,
5748 data: Vec<T>,
5749 ) -> std::result::Result<Self, IdxTensorError> {
5750 let dims = Self::expected_dims_from_indices(&indices);
5751 Self::validate_indices(&indices)?;
5752 Self::validate_dense_payload_len(data.len(), &dims)?;
5753 let native = dense_native_tensor_from_col_major(&data, &dims)?;
5754 Self::from_native(indices, native).map_err(IdxTensorError::from)
5755 }
5756
5757 pub fn from_dense_any(
5786 indices: Vec<DynIndex>,
5787 data: Vec<AnyScalar>,
5788 ) -> std::result::Result<Self, IdxTensorError> {
5789 if data.iter().any(AnyScalar::is_complex) {
5790 Self::from_dense(indices, Self::any_scalar_payload_to_complex(data))
5791 } else {
5792 Self::from_dense(indices, Self::any_scalar_payload_to_real(data))
5793 }
5794 }
5795
5796 pub fn from_diag<T: TensorElement>(
5828 indices: Vec<DynIndex>,
5829 data: Vec<T>,
5830 ) -> std::result::Result<Self, IdxTensorError> {
5831 let dims = Self::expected_dims_from_indices(&indices);
5832 Self::validate_indices(&indices)?;
5833 Self::validate_diag_payload_len(data.len(), &dims)?;
5834 let native = diag_native_tensor_from_col_major(&data, dims.len())?;
5835 Self::from_native_with_axis_classes(indices, native, Self::diag_axis_classes(dims.len()))
5836 .map_err(IdxTensorError::from)
5837 }
5838
5839 pub fn from_diag_any(
5865 indices: Vec<DynIndex>,
5866 data: Vec<AnyScalar>,
5867 ) -> std::result::Result<Self, IdxTensorError> {
5868 if data.iter().any(AnyScalar::is_complex) {
5869 Self::from_diag(indices, Self::any_scalar_payload_to_complex(data))
5870 } else {
5871 Self::from_diag(indices, Self::any_scalar_payload_to_real(data))
5872 }
5873 }
5874
5875 pub fn copy_tensor(
5899 indices: Vec<DynIndex>,
5900 value: AnyScalar,
5901 ) -> std::result::Result<Self, IdxTensorError> {
5902 if indices.is_empty() {
5903 return Self::from_dense_any(vec![], vec![value]);
5904 }
5905 let dim = indices[0].dim();
5906 let data = vec![value; dim];
5907 Self::from_diag_any(indices, data)
5908 }
5909
5910 pub fn fuse_indices(
5967 &self,
5968 old_indices: &[DynIndex],
5969 new_index: DynIndex,
5970 order: LinearizationOrder,
5971 ) -> std::result::Result<Self, IdxTensorError> {
5972 if !(!old_indices.is_empty()) {
5973 return Err(anyhow::anyhow!("fuse_indices requires at least one index to fuse").into());
5974 };
5975
5976 let old_dims = self.dims();
5977 let mut seen_indices = HashSet::new();
5978 let mut old_axes = Vec::with_capacity(old_indices.len());
5979 for old_index in old_indices {
5980 if !(seen_indices.insert(old_index)) {
5981 return Err(anyhow::anyhow!("duplicate index in old_indices").into());
5982 };
5983 let axis = self
5984 .indices
5985 .iter()
5986 .position(|idx| idx == old_index)
5987 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
5988 if !(old_index.dim() == old_dims[axis]) {
5989 return Err(anyhow::anyhow!(
5990 "old index dimension does not match tensor axis dimension"
5991 )
5992 .into());
5993 };
5994 old_axes.push(axis);
5995 }
5996
5997 let fused_dims: Vec<usize> = old_axes.iter().map(|&axis| old_dims[axis]).collect();
5998 let fused_product = checked_product(&fused_dims)?;
5999 if !(fused_product == new_index.dim()) {
6000 return Err(anyhow::anyhow!(
6001 "product of old index dimensions must match the replacement index dimension"
6002 )
6003 .into());
6004 };
6005
6006 let insertion_axis =
6007 old_axes.iter().copied().min().ok_or_else(|| {
6008 anyhow::anyhow!("fuse_indices requires at least one index to fuse")
6009 })?;
6010 let old_axis_set: HashSet<usize> = old_axes.iter().copied().collect();
6011
6012 let mut result_indices =
6013 Vec::with_capacity(self.indices.len() - old_indices.len() + 1usize);
6014 for (axis, index) in self.indices.iter().enumerate() {
6015 if axis == insertion_axis {
6016 result_indices.push(new_index.clone());
6017 }
6018 if !old_axis_set.contains(&axis) {
6019 result_indices.push(index.clone());
6020 }
6021 }
6022 let mut result_seen = HashSet::new();
6023 for index in &result_indices {
6024 if !(result_seen.insert(index)) {
6025 return Err(
6026 anyhow::anyhow!("fuse_indices result would contain duplicate index").into(),
6027 );
6028 };
6029 }
6030 Self::validate_indices(&result_indices)?;
6031
6032 let mut new_dims = Vec::with_capacity(old_dims.len() - old_indices.len() + 1usize);
6033 for (axis, dim) in old_dims.iter().copied().enumerate() {
6034 if axis == insertion_axis {
6035 new_dims.push(new_index.dim());
6036 }
6037 if !old_axis_set.contains(&axis) {
6038 new_dims.push(dim);
6039 }
6040 }
6041
6042 self.ensure_shape_packing_preserves_ad("fuse_indices")?;
6043
6044 let mut grouped_axes = old_axes.clone();
6045 if matches!(order, LinearizationOrder::RowMajor) {
6046 grouped_axes.reverse();
6047 }
6048 let mut perm = Vec::with_capacity(self.indices.len());
6049 perm.extend((0..insertion_axis).filter(|axis| !old_axis_set.contains(axis)));
6050 perm.extend(grouped_axes);
6051 perm.extend(
6052 ((insertion_axis + 1)..self.indices.len()).filter(|axis| !old_axis_set.contains(axis)),
6053 );
6054 debug_assert_eq!(perm.len(), self.indices.len());
6055
6056 let packed = self.permute(&perm)?;
6057 let reshaped = packed.try_materialized_inner()?.reshape(&new_dims)?;
6058 Self::from_inner(result_indices, reshaped).map_err(IdxTensorError::from)
6059 }
6060
6061 pub fn unfuse_index(
6086 &self,
6087 old_index: &DynIndex,
6088 new_indices: &[DynIndex],
6089 order: LinearizationOrder,
6090 ) -> std::result::Result<Self, IdxTensorError> {
6091 if !(!new_indices.is_empty()) {
6092 return Err(
6093 anyhow::anyhow!("unfuse_index requires at least one replacement index").into(),
6094 );
6095 };
6096
6097 let axis = self
6098 .indices
6099 .iter()
6100 .position(|idx| idx == old_index)
6101 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
6102
6103 let replacement_dims: Vec<usize> = new_indices.iter().map(DynIndex::dim).collect();
6104 let replacement_product = checked_product(&replacement_dims)?;
6105 if !(replacement_product == old_index.dim()) {
6106 return Err(anyhow::anyhow!(
6107 "product of new index dimensions must match the replaced index dimension"
6108 )
6109 .into());
6110 };
6111
6112 let mut result_indices =
6113 Vec::with_capacity(self.indices.len() - 1usize + new_indices.len());
6114 result_indices.extend_from_slice(&self.indices[..axis]);
6115 result_indices.extend(new_indices.iter().cloned());
6116 result_indices.extend_from_slice(&self.indices[axis + 1..]);
6117 Self::validate_indices(&result_indices)?;
6118
6119 let old_dims = self.dims();
6120 let mut new_dims = Vec::with_capacity(old_dims.len() - 1usize + replacement_dims.len());
6121 new_dims.extend_from_slice(&old_dims[..axis]);
6122 new_dims.extend_from_slice(&replacement_dims);
6123 new_dims.extend_from_slice(&old_dims[axis + 1..]);
6124
6125 self.ensure_shape_packing_preserves_ad("unfuse_index")?;
6126
6127 let mut grouped_indices = new_indices.to_vec();
6128 let mut grouped_dims = replacement_dims.clone();
6129 if matches!(order, LinearizationOrder::RowMajor) {
6130 grouped_indices.reverse();
6131 grouped_dims.reverse();
6132 }
6133 let mut packed_indices =
6134 Vec::with_capacity(self.indices.len() - 1usize + grouped_indices.len());
6135 packed_indices.extend_from_slice(&self.indices[..axis]);
6136 packed_indices.extend(grouped_indices);
6137 packed_indices.extend_from_slice(&self.indices[axis + 1..]);
6138
6139 let mut packed_dims = Vec::with_capacity(old_dims.len() - 1usize + grouped_dims.len());
6140 packed_dims.extend_from_slice(&old_dims[..axis]);
6141 packed_dims.extend_from_slice(&grouped_dims);
6142 packed_dims.extend_from_slice(&old_dims[axis + 1..]);
6143
6144 let reshaped = self.try_materialized_inner()?.reshape(&packed_dims)?;
6145 let packed = Self::from_inner(packed_indices, reshaped)?;
6146 if matches!(order, LinearizationOrder::ColumnMajor) {
6147 Ok(packed)
6148 } else {
6149 packed.permute_indices(&result_indices)
6150 }
6151 }
6152
6153 pub fn scalar<T: TensorElement>(value: T) -> std::result::Result<Self, IdxTensorError> {
6167 Self::from_dense(vec![], vec![value])
6168 }
6169
6170 pub fn zeros<T: TensorElement + Zero + Clone>(
6186 indices: Vec<DynIndex>,
6187 ) -> std::result::Result<Self, IdxTensorError> {
6188 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
6189 let size = checked_product(&dims)?;
6190 Self::from_dense(indices, vec![T::zero(); size])
6191 }
6192}
6193
6194impl IdxTensor {
6199 pub fn to_vec<T: TensorElement>(&self) -> std::result::Result<Vec<T>, IdxTensorError> {
6223 self.as_inner()?
6224 .duplicate_value()?
6225 .as_slice::<T>()
6226 .map(|values| values.to_vec())
6227 .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))
6228 }
6229
6230 pub fn with_dense_slice<T: TensorElement, R>(
6269 &self,
6270 read: impl FnOnce(&[T]) -> R,
6271 ) -> std::result::Result<R, IdxTensorError> {
6272 let inner = self.as_inner()?;
6273 let value = inner.value()?;
6274 if let Ok(values) = value.as_slice::<T>() {
6275 return Ok(read(values));
6276 }
6277
6278 drop(value);
6283 let values = inner.duplicate_value()?;
6284 let values = values
6285 .as_slice::<T>()
6286 .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))?;
6287 Ok(read(values))
6288 }
6289
6290 pub fn into_dense_col_major_parts<T: TensorElement>(
6325 self,
6326 ) -> std::result::Result<(Vec<DynIndex>, Vec<T>), IdxTensorError> {
6327 if !(!self.tracks_grad()) {
6328 return Err(anyhow::anyhow!("IdxTensor::into_dense_col_major_parts cannot consume tensors with tracked autodiff state").into());
6329 };
6330 let data = self.to_vec::<T>()?;
6331 Ok((self.indices, data))
6332 }
6333
6334 pub fn is_f64(&self) -> bool {
6347 self.storage.dtype() == Some(DType::F64)
6348 }
6349
6350 pub fn is_f32(&self) -> bool {
6365 self.storage.dtype() == Some(DType::F32)
6366 }
6367
6368 pub fn is_c32(&self) -> bool {
6384 self.storage.dtype() == Some(DType::C32)
6385 }
6386
6387 pub fn is_c64(&self) -> bool {
6403 self.storage.is_c64()
6404 }
6405
6406 pub fn is_diag(&self) -> bool {
6433 self.storage.is_diag()
6434 }
6435
6436 pub fn is_complex(&self) -> bool {
6455 self.storage.is_complex()
6456 }
6457}
6458
6459fn checked_product(dims: &[usize]) -> Result<usize> {
6460 dims.iter().try_fold(1usize, |acc, &dim| {
6461 acc.checked_mul(dim)
6462 .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
6463 })
6464}
6465
6466fn increment_col_major_coordinate(coords: &mut [usize], dims: &[usize]) {
6467 let mut carry = true;
6468 for (coordinate, &dim) in coords.iter_mut().zip(dims.iter()) {
6469 if !carry {
6470 break;
6471 }
6472 *coordinate += 1;
6473 if *coordinate == dim {
6474 *coordinate = 0;
6475 } else {
6476 carry = false;
6477 }
6478 }
6479}
6480
6481fn map_payload_support_coordinate(
6482 source_axis_classes: &[usize],
6483 target_axis_classes: &[usize],
6484 source_to_target_axes: &[usize],
6485 source_coords: &[usize],
6486 target_coords: &mut [usize],
6487 target_seen: &mut [bool],
6488) -> bool {
6489 if source_axis_classes.len() != source_to_target_axes.len()
6490 || target_coords.len() != target_seen.len()
6491 {
6492 return false;
6493 }
6494 target_seen.fill(false);
6495 for (source_axis, &target_axis) in source_to_target_axes.iter().enumerate() {
6496 let Some(&source_class) = source_axis_classes.get(source_axis) else {
6497 return false;
6498 };
6499 let Some(&target_class) = target_axis_classes.get(target_axis) else {
6500 return false;
6501 };
6502 let Some(&source_value) = source_coords.get(source_class) else {
6503 return false;
6504 };
6505 let Some(target_value) = target_coords.get_mut(target_class) else {
6506 return false;
6507 };
6508 if target_seen[target_class] {
6509 if *target_value != source_value {
6510 return false;
6511 }
6512 } else {
6513 *target_value = source_value;
6514 target_seen[target_class] = true;
6515 }
6516 }
6517 target_seen.iter().all(|&seen| seen)
6518}
6519
6520fn decode_col_major_linear(linear: usize, dims: &[usize]) -> Result<Vec<usize>> {
6521 let total = checked_product(dims)?;
6522 if !(linear < total) {
6523 return Err(anyhow::anyhow!(
6524 "linear offset {} out of bounds for dims {:?}",
6525 linear,
6526 dims
6527 ));
6528 };
6529 let mut remaining = linear;
6530 let mut out = Vec::with_capacity(dims.len());
6531 for &dim in dims {
6532 out.push(remaining % dim);
6533 remaining /= dim;
6534 }
6535 Ok(out)
6536}
6537
6538fn encode_col_major_linear(indices: &[usize], dims: &[usize]) -> Result<usize> {
6539 if !(indices.len() == dims.len()) {
6540 return Err(anyhow::anyhow!(
6541 "index rank {} does not match dims {:?}",
6542 indices.len(),
6543 dims
6544 ));
6545 };
6546 let mut linear = 0usize;
6547 let mut stride = 1usize;
6548 for (&index, &dim) in indices.iter().zip(dims.iter()) {
6549 if !(index < dim) {
6550 return Err(anyhow::anyhow!(
6551 "index {} out of bounds for dimension {}",
6552 index,
6553 dim
6554 ));
6555 };
6556 let term = index
6557 .checked_mul(stride)
6558 .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
6559 linear = linear
6560 .checked_add(term)
6561 .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
6562 stride = stride
6563 .checked_mul(dim)
6564 .ok_or_else(|| anyhow::anyhow!("stride overflow"))?;
6565 }
6566 Ok(linear)
6567}
6568
6569#[cfg(test)]
6570mod tests {
6571 use super::*;
6572 use num_complex::{Complex32, Complex64};
6573 use std::cell::Cell;
6574 use tensor4all_tensorbackend::StorageError;
6575
6576 #[test]
6577 fn structured_contraction_does_not_install_logical_dense_cache() {
6578 let n = 8;
6579 let left = DynIndex::new_dyn(n);
6580 let site = DynIndex::new_dyn(3);
6581 let right = DynIndex::new_dyn(n);
6582 let far = DynIndex::new_dyn(n);
6583 let end = DynIndex::new_dyn(n);
6584 let a =
6585 IdxTensor::from_copy_selector(left, site.clone(), right.clone(), 1, 1.0_f64).unwrap();
6586 let b =
6587 IdxTensor::from_copy_selector(right, site.clone(), far.clone(), 1, 2.0_f64).unwrap();
6588 let c = IdxTensor::from_copy_selector(far, site.clone(), end, 1, 3.0_f64).unwrap();
6589 let result = crate::defaults::contract::contract_with_options(
6590 &[&a, &b, &c],
6591 crate::defaults::contract::ContractionOptions::new()
6592 .with_retain_indices(std::slice::from_ref(&site)),
6593 )
6594 .unwrap();
6595
6596 assert_eq!(result.storage_kind(), StorageKind::Structured);
6597 assert!(result.eager_cache.get().is_none());
6598 assert_eq!(result.storage().unwrap().payload_len(), n * 3);
6599 }
6600
6601 #[test]
6602 fn structured_metrics_use_authoritative_compact_payload_for_all_dtypes() {
6603 fn check(tensor: IdxTensor, expected_sum: f64, expected_norm_squared: f64) {
6604 assert!(matches!(tensor.storage, IdxTensorStorage::Compact(_)));
6605 assert!(tensor.eager_cache.get().is_none());
6606 assert!((tensor.sum().unwrap().real() - expected_sum).abs() < 1.0e-6);
6607 assert!((tensor.norm_squared().unwrap() - expected_norm_squared).abs() < 1.0e-6);
6608 assert!((tensor.maxabs().unwrap() - 2.0).abs() < 1.0e-6);
6609 assert!(tensor.isapprox(&tensor, 0.0, 0.0).unwrap());
6610 assert!(tensor.eager_cache.get().is_none());
6611 }
6612
6613 let indices = || vec![DynIndex::new_dyn(2), DynIndex::new_dyn(2)];
6614 check(
6615 IdxTensor::from_diag(indices(), vec![1.0_f32, 2.0]).unwrap(),
6616 3.0,
6617 5.0,
6618 );
6619 check(
6620 IdxTensor::from_diag(indices(), vec![1.0_f64, 2.0]).unwrap(),
6621 3.0,
6622 5.0,
6623 );
6624 check(
6625 IdxTensor::from_diag(
6626 indices(),
6627 vec![Complex32::new(1.0, 0.0), Complex32::new(2.0, 0.0)],
6628 )
6629 .unwrap(),
6630 3.0,
6631 5.0,
6632 );
6633 check(
6634 IdxTensor::from_diag(
6635 indices(),
6636 vec![Complex64::new(1.0, 0.0), Complex64::new(2.0, 0.0)],
6637 )
6638 .unwrap(),
6639 3.0,
6640 5.0,
6641 );
6642 }
6643
6644 #[test]
6645 fn payload_storage_error_retains_typed_source() {
6646 let storage = Storage::from_dense_col_major(vec![1.0_f64, 2.0], &[2]).unwrap();
6647 let error = storage.scalar_at(&[2]).unwrap_err();
6648 assert!(matches!(error, StorageError::InvalidStructuredStorage(_)));
6649 }
6650
6651 #[test]
6652 fn materialization_error_retains_backend_source() {
6653 let native = NativeTensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
6654 let inner = EagerTensor::from_tensor_in(native, default_eager_ctx().unwrap()).unwrap();
6655 let storage = IdxTensorStorage::Eager {
6656 inner: Arc::new(inner),
6657 axis_classes: vec![0, 0],
6658 };
6659
6660 let error = storage.materialize(2).unwrap_err();
6661 assert!(matches!(error, TensorStorageError::Materialization { .. }));
6662 assert!(std::error::Error::source(&error).is_some());
6663 }
6664
6665 fn conjugate_with_injected_failure(
6666 tensor: &IdxTensor,
6667 target: *const EagerTensor,
6668 message: &'static str,
6669 ) -> IdxTensor {
6670 let calls = Cell::new(0usize);
6671 let conjugated = tensor.conj_with(&|inner| {
6672 calls.set(calls.get() + 1);
6673 if std::ptr::eq(inner, target) {
6674 Err(Arc::new(std::io::Error::other(message)) as _)
6675 } else {
6676 conjugate_eager(inner)
6677 }
6678 });
6679 assert!(calls.get() > 0, "injected closure was not reached");
6680 conjugated
6681 }
6682
6683 fn assert_unwrapped_conjugation_error(
6684 tensor: IdxTensor,
6685 target: *const EagerTensor,
6686 message: &'static str,
6687 ) -> IdxTensor {
6688 let conjugated = conjugate_with_injected_failure(&tensor, target, message);
6689 let error = conjugated.to_storage().unwrap_err();
6690 assert!(matches!(error, TensorStorageError::Conjugation { .. }));
6691 let source = std::error::Error::source(&error).unwrap();
6692 assert_eq!(source.to_string(), message);
6693 assert!(
6694 source.source().is_none(),
6695 "source was wrapped more than once"
6696 );
6697 conjugated
6698 }
6699
6700 #[test]
6701 fn authoritative_storage_conjugation_failure_is_deferred_without_detaching() {
6702 let i = DynIndex::new_dyn(2);
6703 let native = NativeTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
6704 let inner = EagerTensor::requires_grad_in(native, default_eager_ctx().unwrap()).unwrap();
6705 let tensor = IdxTensor::from_inner(vec![i], inner).unwrap();
6706 let source = match &tensor.storage {
6707 IdxTensorStorage::Eager { inner, .. } => Arc::clone(inner),
6708 IdxTensorStorage::Compact(payload) => Arc::clone(&payload.payload),
6709 IdxTensorStorage::Materialized(_) | IdxTensorStorage::Deferred { .. } => {
6710 panic!("tracked eager source expected")
6711 }
6712 };
6713
6714 let conjugated = conjugate_with_injected_failure(
6715 &tensor,
6716 Arc::as_ptr(&source),
6717 "forced authoritative eager conjugation failure",
6718 );
6719 assert!(conjugated.tracks_grad());
6720 assert!(conjugated.is_f64());
6721 assert!(!conjugated.is_complex());
6722 assert!(!conjugated.is_diag());
6723 assert_eq!(conjugated.dims(), vec![2]);
6724
6725 let error = conjugated.to_storage().unwrap_err();
6726 let source = std::error::Error::source(&error).unwrap();
6727 assert_eq!(
6728 source.to_string(),
6729 "forced authoritative eager conjugation failure"
6730 );
6731 assert!(conjugated.detach().is_err());
6732 }
6733
6734 #[test]
6735 fn structured_payload_conjugation_failure_retains_graph_and_blocks_detached_primal() {
6736 let i = DynIndex::new_dyn(2);
6737 let j = DynIndex::new_dyn(2);
6738 let tensor = IdxTensor::from_diag(
6739 vec![i, j],
6740 vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, -4.0)],
6741 )
6742 .unwrap()
6743 .enable_grad()
6744 .unwrap();
6745
6746 let target = Arc::as_ptr(
6747 &tensor
6748 .storage
6749 .compact_payload()
6750 .expect("tracked compact payload")
6751 .payload,
6752 );
6753 let conjugated = assert_unwrapped_conjugation_error(
6754 tensor,
6755 target,
6756 "forced structured AD conjugation failure",
6757 );
6758 assert!(conjugated.tracks_grad());
6759 assert!(conjugated.detach().is_err());
6760 assert!(conjugated.clone().enable_grad().is_err());
6761 assert!(conjugated.sum().is_err());
6762 assert!(conjugated.grad().is_err());
6763 assert!(conjugated.clear_grad().is_err());
6764 assert!(conjugated.maxabs().is_err());
6765 assert!(conjugated.norm_squared().is_err());
6766
6767 let twice_conjugated = conjugated.conj();
6768 let error = twice_conjugated.to_storage().unwrap_err();
6769 assert_eq!(
6770 std::error::Error::source(&error).unwrap().to_string(),
6771 "forced structured AD conjugation failure"
6772 );
6773 }
6774
6775 #[test]
6776 fn eager_cache_conjugation_failure_is_deferred_with_original_diagnostic() {
6777 let i = DynIndex::new_dyn(2);
6778 let j = DynIndex::new_dyn(2);
6779 let tensor = IdxTensor::from_diag(vec![i, j], vec![1.0_f64, 2.0]).unwrap();
6780 tensor.as_inner().unwrap();
6781 let target = Arc::as_ptr(tensor.eager_cache.get().unwrap());
6782
6783 let conjugated = assert_unwrapped_conjugation_error(
6784 tensor,
6785 target,
6786 "forced eager cache conjugation failure",
6787 );
6788 assert!(!conjugated.tracks_grad());
6789 assert!(conjugated.detach().is_err());
6790 }
6791
6792 #[test]
6793 fn encode_col_major_linear_rejects_offset_overflow() {
6794 let error =
6795 encode_col_major_linear(&[usize::MAX - 1, usize::MAX - 1], &[usize::MAX, usize::MAX])
6796 .unwrap_err();
6797 assert!(error.to_string().contains("linear offset overflow"));
6798 }
6799}