Skip to main content

tensor4all_core/defaults/
idx_tensor.rs

1use crate::defaults::DynIndex;
2use crate::index_like::IndexLike;
3use crate::index_ops::{common_ind_positions, prepare_contraction, prepare_contraction_pairs};
4use crate::tensor_like::LinearizationOrder;
5use crate::AnyScalar;
6use anyhow::{Context, Result};
7use num_complex::{Complex32, Complex64};
8use num_traits::Zero;
9use rand::Rng;
10use rand_distr::{Distribution, StandardNormal};
11use std::cell::RefCell;
12use std::cmp::Reverse;
13use std::collections::{HashMap, HashSet};
14use std::env;
15use std::sync::{Arc, OnceLock};
16use std::time::{Duration, Instant};
17use tenferro::{
18    DType, DotGeneralConfig, Tensor as NativeTensor, TensorRead, TensorValue, TensorView,
19};
20use tenferro_ad::{extension::adopt_untracked_eager_value, EagerRuntime, EagerTensor};
21use tenferro_einsum::{EagerEinsumExt, EinsumSubscripts};
22use tenferro_linalg::{EagerTensorLinalgExt, RankRevealingQrOptions};
23use tensor4all_tensorbackend::{
24    contract_native_tensor, default_eager_ctx, dense_native_tensor_from_col_major,
25    dense_native_tensor_from_col_major_owned, diag_native_tensor_from_col_major,
26    native_tensor_primal_to_diag, storage_payload_native_read_input, storage_to_native_tensor,
27    ExecutionContext, NativeTensorReadInput, TensorElement,
28};
29use tensor4all_tensorbackend::{
30    src_error_estimate as backend_src_error_estimate,
31    src_error_estimate_general as backend_src_error_estimate_general, Matrix,
32};
33use tensor4all_tensorbackend::{IncrementalQr, IncrementalQrScalar};
34use tensor4all_tensorbackend::{Storage, StorageKind};
35
36use super::contract::{ContractionOptions, PairwiseContractionOptions};
37use super::structured_contraction::{
38    normalize_payload_read_for_roots, storage_from_payload_native, storage_payload_native,
39    OperandLayout, StructuredContractionPlan, StructuredContractionSpec,
40};
41
42fn conjugate_eager(
43    inner: &EagerTensor,
44) -> std::result::Result<EagerTensor, Arc<dyn std::error::Error + Send + Sync + 'static>> {
45    inner.conj().map_err(|source| Arc::new(source) as _)
46}
47
48#[derive(Debug, Default, Clone)]
49struct PairwiseContractProfileEntry {
50    calls: usize,
51    total_time: Duration,
52    total_bytes: usize,
53}
54
55/// Hermitian eigendecomposition of a rank-2 [`IdxTensor`].
56/// Eigenvectors are returned as a rank-2 tensor whose first index is the input
57/// matrix row index and whose second index labels eigenvector columns. The
58/// eigenvalues are detached primal values intended for nonsmooth selection
59/// logic such as truncation cutoffs.
60/// # Examples
61/// ```
62/// use tensor4all_core::{DynIndex, IdxTensor};
63/// let row = DynIndex::new_dyn(2);
64/// let col = DynIndex::new_dyn(2);
65/// let matrix = IdxTensor::from_dense(
66///     vec![row.clone(), col],
67///     vec![1.0_f64, 0.0, 0.0, 2.0],
68/// ).unwrap();
69/// let decomp = matrix.hermitian_eigendecomposition(1.0e-12).unwrap();
70/// assert_eq!(decomp.eigenvalues, vec![1.0, 2.0]);
71/// assert_eq!(
72///     decomp.eigenvectors.indices(),
73///     &[row, decomp.eigenvector_index.clone()]
74/// );
75/// ```
76#[derive(Debug, Clone)]
77pub struct TensorHermitianEigendecomposition {
78    /// Real eigenvalues in backend Hermitian eigensolver order.
79    pub eigenvalues: Vec<f64>,
80    /// Eigenvector matrix with one eigenvector in each column.
81    pub eigenvectors: IdxTensor,
82    /// Index labeling the eigenvector columns.
83    pub eigenvector_index: DynIndex,
84}
85
86thread_local! {
87    static PAIRWISE_CONTRACT_PROFILE_STATE: RefCell<HashMap<&'static str, PairwiseContractProfileEntry>> =
88        RefCell::new(HashMap::new());
89}
90
91fn pairwise_contract_profile_enabled() -> bool {
92    static ENABLED: OnceLock<bool> = OnceLock::new();
93    *ENABLED.get_or_init(|| env::var("T4A_PROFILE_PAIRWISE_CONTRACT").is_ok())
94}
95
96fn record_pairwise_contract_profile(section: &'static str, elapsed: Duration) {
97    if !pairwise_contract_profile_enabled() {
98        return;
99    }
100    PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
101        let mut state = state.borrow_mut();
102        let entry = state.entry(section).or_default();
103        entry.calls += 1;
104        entry.total_time += elapsed;
105    });
106}
107
108fn record_pairwise_contract_profile_bytes(section: &'static str, bytes: usize) {
109    if !pairwise_contract_profile_enabled() {
110        return;
111    }
112    PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
113        let mut state = state.borrow_mut();
114        let entry = state.entry(section).or_default();
115        entry.total_bytes += bytes;
116    });
117}
118
119fn profile_pairwise_contract_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
120    if !pairwise_contract_profile_enabled() {
121        return f();
122    }
123    let started = Instant::now();
124    let result = f();
125    record_pairwise_contract_profile(section, started.elapsed());
126    result
127}
128
129/// Reset the aggregated pairwise `IdxTensor` contraction profile.
130pub fn reset_pairwise_contract_profile() {
131    PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| state.borrow_mut().clear());
132}
133
134/// Print and clear the aggregated pairwise `IdxTensor` contraction profile.
135pub fn print_and_reset_pairwise_contract_profile() {
136    if !pairwise_contract_profile_enabled() {
137        return;
138    }
139    PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
140        let mut entries: Vec<_> = state
141            .borrow()
142            .iter()
143            .map(|(section, entry)| (*section, entry.clone()))
144            .collect();
145        state.borrow_mut().clear();
146        entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
147
148        eprintln!("=== IdxTensor pairwise contract profile ===");
149        for (section, entry) in entries {
150            let per_call_us = if entry.calls == 0 {
151                0.0
152            } else {
153                entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64
154            };
155            eprintln!(
156                "{section}: calls={} total={:.6}ms per_call={:.3}us bytes={}",
157                entry.calls,
158                entry.total_time.as_secs_f64() * 1.0e3,
159                per_call_us,
160                entry.total_bytes,
161            );
162        }
163    });
164}
165
166fn tensor_profile_bytes(dtype: DType, shape: &[usize]) -> usize {
167    let element_size = match dtype {
168        DType::F32 => 4,
169        DType::F64 => 8,
170        DType::C32 => 8,
171        DType::C64 => 16,
172        DType::I32 => 4,
173        DType::I64 => 8,
174        DType::Bool => 1,
175    };
176    shape
177        .iter()
178        .try_fold(1usize, |bytes, &dim| bytes.checked_mul(dim))
179        .and_then(|elements| elements.checked_mul(element_size))
180        .unwrap_or(usize::MAX)
181}
182
183/// Trait for scalar types that can generate random values from a standard
184/// normal distribution.
185/// This enables the generic [`IdxTensor::random`] constructor.
186pub trait RandomScalar: TensorElement {
187    /// Generate a random value from the standard normal distribution.
188    fn random_value<R: Rng>(rng: &mut R) -> Self;
189}
190
191impl RandomScalar for f64 {
192    fn random_value<R: Rng>(rng: &mut R) -> Self {
193        StandardNormal.sample(rng)
194    }
195}
196
197impl RandomScalar for Complex64 {
198    fn random_value<R: Rng>(rng: &mut R) -> Self {
199        Complex64::new(StandardNormal.sample(rng), StandardNormal.sample(rng))
200    }
201}
202
203/// Compute the permutation array from original indices to new indices.
204/// This function finds the mapping from new indices to original indices by
205/// matching index IDs. The result is a permutation array `perm` such that
206/// `new_indices[i]` corresponds to `original_indices[perm[i]]`.
207/// # Arguments
208/// * `original_indices` - The original indices in their current order
209/// * `new_indices` - The desired new indices order (must be a permutation of original_indices)
210/// # Returns
211/// A `Vec<usize>` representing the permutation: `perm[i]` is the position in
212/// `original_indices` of the index that should be at position `i` in `new_indices`.
213/// # Errors
214/// Returns an error when `new_order` contains indices not present in
215/// `original` (a missing-index failure) or the two lists differ in length
216/// (a length mismatch).
217/// # Example
218/// ```
219/// use tensor4all_core::tensor::compute_permutation_from_indices;
220/// use tensor4all_core::DynIndex;
221/// let i = DynIndex::new_dyn(2);
222/// let j = DynIndex::new_dyn(3);
223/// let original = vec![i.clone(), j.clone()];
224/// let new_order = vec![j.clone(), i.clone()];
225/// let perm = compute_permutation_from_indices(&original, &new_order).unwrap();
226/// assert_eq!(perm, vec![1, 0]);  // j is at position 1, i is at position 0
227/// ```
228pub fn compute_permutation_from_indices(
229    original_indices: &[DynIndex],
230    new_indices: &[DynIndex],
231) -> std::result::Result<Vec<usize>, IdxTensorError> {
232    if !(new_indices.len() == original_indices.len()) {
233        return Err(
234            anyhow::anyhow!("new_indices length must match original_indices length").into(),
235        );
236    };
237
238    let mut perm = Vec::with_capacity(new_indices.len());
239    let mut used = std::collections::HashSet::new();
240
241    for new_idx in new_indices {
242        // Find the position of this index in the original indices
243        // DynIndex implements Eq, so we can compare directly
244        let pos = original_indices
245            .iter()
246            .position(|old_idx| old_idx == new_idx)
247            .ok_or_else(|| {
248                anyhow::anyhow!("new_indices must be a permutation of original_indices")
249            })?;
250
251        if !(used.insert(pos)) {
252            return Err(anyhow::anyhow!("duplicate index in new_indices").into());
253        };
254        perm.push(pos);
255    }
256
257    Ok(perm)
258}
259
260/// Compact structured payload kept in the authoritative eager representation.
261/// The payload may use any supported eager dtype (`f32`, `f64`, `c32`, or
262/// `c64`) and may be either tracked or untracked. Tracking is a property of
263/// `payload`, never of the presence of this metadata container.
264#[derive(Clone)]
265pub(crate) struct StructuredPayload {
266    payload: Arc<EagerTensor>,
267    payload_dims: Vec<usize>,
268    axis_classes: Vec<usize>,
269}
270
271/// Error returned when [`IdxTensor::storage`] or
272/// [`IdxTensor::to_storage`] cannot produce a compact `f64`/`Complex64`
273/// storage snapshot from the authoritative payload.
274/// Backend diagnostics remain available through [`std::error::Error::source`]
275/// instead of being erased into a display string. The error is cloneable so a
276/// deferred failure can be retained by cloned tensors without rebuilding a
277/// detached primal value.
278/// # Examples
279/// ```
280/// use std::error::Error;
281/// use std::sync::Arc;
282/// use tensor4all_core::TensorStorageError;
283/// let error = TensorStorageError::Materialization {
284///     source: Arc::new(std::io::Error::other("backend unavailable")),
285/// };
286/// assert!(error.source().is_some());
287/// assert!(error.to_string().contains("backend unavailable"));
288/// ```
289#[derive(Debug, Clone, thiserror::Error)]
290pub enum TensorStorageError {
291    /// The eager or structured payload could not be converted to compact storage.
292    #[error("failed to materialize IdxTensor storage: {source}")]
293    Materialization {
294        /// Original diagnostic returned by the backend or storage conversion seam.
295        #[source]
296        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
297    },
298    /// An eager payload uses a scalar dtype that compact [`Storage`] cannot hold.
299    #[error(
300        "compact IdxTensor storage does not support dtype {dtype}; the eager payload remains authoritative"
301    )]
302    UnsupportedDtype {
303        /// Native scalar dtype retained by the eager representation.
304        dtype: &'static str,
305    },
306    /// An eager conjugation operation failed and was deferred by the infallible
307    /// [`IdxTensor::conj`] API.
308    #[error("failed to conjugate IdxTensor storage: {source}")]
309    Conjugation {
310        /// Original diagnostic returned by the eager AD backend.
311        #[source]
312        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
313    },
314}
315
316/// Errors returned by the fallible numerical and comparison methods on
317/// [`IdxTensor`].
318/// The enum is intentionally owned by `tensor4all-core`: callers can match
319/// storage, shape, scalar, subtraction, and invalid-value failures without
320/// depending on the internal `anyhow` plumbing. Wrapped backend diagnostics
321/// retain their complete [`std::error::Error::source`] chain.
322///
323/// # Remedies
324/// - Storage failures: check the operation against the storage kind
325///   (structured vs dense) and dtype before calling; the eager payload may
326///   remain authoritative for unsupported dtypes.
327/// - Shape/index failures: validate indices and dimensions at the call site
328///   (`dims`, `indices`, external index sets) before the operation.
329/// - NaN/invalid-value failures: inspect the payload for non-finite entries
330///   before numerical comparisons.
331/// - Backend failures: the wrapped source chain identifies the backend stage;
332///   re-run with the backend diagnostic visible (see `source`).
333/// # Examples
334/// ```
335/// use tensor4all_core::IdxTensorError;
336/// let error = IdxTensorError::NaNInput {
337///     operation: "norm_squared",
338/// };
339/// assert!(error.to_string().contains("NaN"));
340/// ```
341#[derive(Debug, Clone, thiserror::Error)]
342pub enum IdxTensorError {
343    /// Compact storage or deferred storage materialization failed.
344    #[error("IdxTensor storage operation failed: {source}")]
345    Storage {
346        /// Original storage diagnostic, including its backend source chain.
347        #[source]
348        source: TensorStorageError,
349    },
350    /// A native eager payload could not be materialized for a numerical operation.
351    #[error("IdxTensor materialization failed: {source}")]
352    Materialization {
353        /// Original backend or eager-runtime diagnostic.
354        #[source]
355        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
356    },
357    /// A rank-zero scalar could not be extracted from a reduction result.
358    #[error("IdxTensor scalar extraction failed: {source}")]
359    ScalarExtraction {
360        /// Original scalar-wrapper or backend diagnostic.
361        #[source]
362        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
363    },
364    /// The reduction result has a scalar dtype that this real-valued operation
365    /// cannot interpret.
366    #[error("IdxTensor scalar type mismatch: expected {expected}, got {actual}")]
367    ScalarTypeMismatch {
368        /// Scalar dtype required by the operation.
369        expected: &'static str,
370        /// Scalar dtype returned by the reduction.
371        actual: String,
372    },
373    /// Tensor shapes, index spaces, or dimension metadata cannot be aligned
374    /// for an operation (comparison, index replacement, or other shape-sensitive
375    /// transformations).
376    #[error("IdxTensor shape mismatch during {operation}: expected {expected}, got {actual}")]
377    ShapeMismatch {
378        /// Operation that attempted the alignment.
379        operation: &'static str,
380        /// Expected index/dimension description.
381        expected: String,
382        /// Actual index/dimension description.
383        actual: String,
384    },
385    /// Tensor subtraction failed while evaluating a comparison.
386    #[error("IdxTensor subtraction failed: {source}")]
387    Subtraction {
388        /// Original arithmetic or backend diagnostic.
389        #[source]
390        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
391    },
392    /// An input contained a NaN and the operation rejected it rather than
393    /// silently converting it to zero.
394    #[error("IdxTensor {operation} received NaN input")]
395    NaNInput {
396        /// Numerical operation that observed the NaN.
397        operation: &'static str,
398    },
399    /// A comparison tolerance was NaN, infinite, or negative.
400    #[error("IdxTensor tolerance {name} is invalid: {value}")]
401    InvalidTolerance {
402        /// Name of the invalid tolerance.
403        name: &'static str,
404        /// Supplied tolerance value.
405        value: f64,
406    },
407    /// Another eager tensor operation failed while preparing a comparison.
408    #[error("IdxTensor {operation} failed: {source}")]
409    Operation {
410        /// Name of the eager operation that failed.
411        operation: &'static str,
412        /// Original backend or tensor diagnostic.
413        #[source]
414        source: Arc<dyn std::error::Error + Send + Sync + 'static>,
415    },
416}
417
418impl From<anyhow::Error> for IdxTensorError {
419    fn from(source: anyhow::Error) -> Self {
420        Self::operation("IdxTensor", source)
421    }
422}
423
424impl From<TensorStorageError> for IdxTensorError {
425    fn from(source: TensorStorageError) -> Self {
426        Self::Storage { source }
427    }
428}
429
430impl From<tenferro_ad::Error> for IdxTensorError {
431    fn from(source: tenferro_ad::Error) -> Self {
432        Self::Materialization {
433            source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
434        }
435    }
436}
437
438impl From<tensor4all_tensorbackend::EagerContextError> for IdxTensorError {
439    fn from(source: tensor4all_tensorbackend::EagerContextError) -> Self {
440        Self::Materialization {
441            source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
442        }
443    }
444}
445
446impl From<tensor4all_tensorbackend::BridgeError> for IdxTensorError {
447    fn from(source: tensor4all_tensorbackend::BridgeError) -> Self {
448        Self::Materialization {
449            source: Arc::new(source),
450        }
451    }
452}
453
454impl IdxTensorError {
455    fn boxed(error: anyhow::Error) -> Arc<dyn std::error::Error + Send + Sync + 'static> {
456        Arc::from(error.into_boxed_dyn_error())
457    }
458
459    fn materialization(error: anyhow::Error) -> Self {
460        Self::Materialization {
461            source: Self::boxed(error),
462        }
463    }
464
465    fn scalar_extraction(error: anyhow::Error) -> Self {
466        Self::ScalarExtraction {
467            source: Self::boxed(error),
468        }
469    }
470
471    fn operation(operation: &'static str, error: anyhow::Error) -> Self {
472        Self::Operation {
473            operation,
474            source: Self::boxed(error),
475        }
476    }
477}
478
479#[derive(Clone)]
480pub(crate) enum IdxTensorStorage {
481    Materialized(Arc<Storage>),
482    Eager {
483        inner: Arc<EagerTensor>,
484        axis_classes: Vec<usize>,
485    },
486    /// One authoritative compact eager payload and its logical layout.
487    Compact(Arc<StructuredPayload>),
488    /// A storage representation whose eager operation failed before the
489    /// infallible tensor API could return an error.
490    Deferred {
491        source: Box<Self>,
492        error: Arc<TensorStorageError>,
493    },
494}
495
496impl IdxTensorStorage {
497    fn from_storage(storage: Arc<Storage>) -> Self {
498        Self::Materialized(storage)
499    }
500
501    fn from_eager_dense(inner: EagerTensor, rank: usize) -> Self {
502        Self::Eager {
503            inner: Arc::new(inner),
504            axis_classes: IdxTensor::dense_axis_classes(rank),
505        }
506    }
507
508    fn eager(&self) -> Option<&EagerTensor> {
509        match self {
510            Self::Materialized(_) => None,
511            Self::Eager { inner, .. } => Some(inner.as_ref()),
512            Self::Compact(payload) => Some(payload.payload.as_ref()),
513            Self::Deferred { source, .. } => source.eager(),
514        }
515    }
516
517    fn deferred_error(&self) -> Option<&TensorStorageError> {
518        match self {
519            Self::Deferred { error, .. } => Some(error.as_ref()),
520            _ => None,
521        }
522    }
523
524    fn with_deferred_error(self, error: TensorStorageError) -> Self {
525        if self.deferred_error().is_some() {
526            self
527        } else {
528            Self::Deferred {
529                source: Box::new(self),
530                error: Arc::new(error),
531            }
532        }
533    }
534
535    fn axis_classes(&self) -> &[usize] {
536        match self {
537            Self::Materialized(storage) => storage.axis_classes(),
538            Self::Eager { axis_classes, .. } => axis_classes,
539            Self::Compact(payload) => &payload.axis_classes,
540            Self::Deferred { source, .. } => source.axis_classes(),
541        }
542    }
543
544    fn payload_dims(&self) -> &[usize] {
545        match self {
546            Self::Materialized(storage) => storage.payload_dims(),
547            Self::Eager { inner, .. } => inner.shape(),
548            Self::Compact(payload) => &payload.payload_dims,
549            Self::Deferred { source, .. } => source.payload_dims(),
550        }
551    }
552
553    fn payload_strides_vec(&self) -> Vec<isize> {
554        match self {
555            Self::Materialized(storage) => storage.payload_strides().to_vec(),
556            Self::Eager { inner, .. } => {
557                IdxTensor::col_major_strides(inner.shape()).unwrap_or_default()
558            }
559            Self::Compact(payload) => {
560                IdxTensor::col_major_strides(&payload.payload_dims).unwrap_or_default()
561            }
562            Self::Deferred { source, .. } => source.payload_strides_vec(),
563        }
564    }
565
566    fn is_f64(&self) -> bool {
567        match self {
568            Self::Materialized(storage) => storage.is_f64(),
569            Self::Eager { inner, .. } => inner.dtype() == DType::F64,
570            Self::Compact(payload) => payload.payload.dtype() == DType::F64,
571            Self::Deferred { source, .. } => source.is_f64(),
572        }
573    }
574
575    fn is_c64(&self) -> bool {
576        match self {
577            Self::Materialized(storage) => storage.is_c64(),
578            Self::Eager { inner, .. } => inner.dtype() == DType::C64,
579            Self::Compact(payload) => payload.payload.dtype() == DType::C64,
580            Self::Deferred { source, .. } => source.is_c64(),
581        }
582    }
583
584    fn dtype(&self) -> Option<DType> {
585        match self {
586            Self::Materialized(storage) => Some(if storage.is_c64() {
587                DType::C64
588            } else {
589                DType::F64
590            }),
591            Self::Eager { inner, .. } => Some(inner.dtype()),
592            Self::Compact(payload) => Some(payload.payload.dtype()),
593            Self::Deferred { source, .. } => source.dtype(),
594        }
595    }
596
597    fn is_complex(&self) -> bool {
598        match self {
599            Self::Materialized(storage) => storage.is_complex(),
600            Self::Eager { inner, .. } => matches!(inner.dtype(), DType::C32 | DType::C64),
601            Self::Compact(payload) => {
602                matches!(payload.payload.dtype(), DType::C32 | DType::C64)
603            }
604            Self::Deferred { source, .. } => source.is_complex(),
605        }
606    }
607
608    fn is_diag(&self) -> bool {
609        match self {
610            Self::Materialized(storage) => storage.is_diag(),
611            Self::Eager { axis_classes, .. } => IdxTensor::is_diag_axis_classes(axis_classes),
612            Self::Compact(payload) => IdxTensor::is_diag_axis_classes(&payload.axis_classes),
613            Self::Deferred { source, .. } => source.is_diag(),
614        }
615    }
616
617    fn storage_kind(&self) -> StorageKind {
618        match self {
619            Self::Materialized(storage) => storage.storage_kind(),
620            Self::Eager { axis_classes, .. } => {
621                if axis_classes.iter().copied().eq(0..axis_classes.len()) {
622                    StorageKind::Dense
623                } else if IdxTensor::is_diag_axis_classes(axis_classes) {
624                    StorageKind::Diagonal
625                } else {
626                    StorageKind::Structured
627                }
628            }
629            Self::Compact(payload) => {
630                if payload
631                    .axis_classes
632                    .iter()
633                    .copied()
634                    .eq(0..payload.axis_classes.len())
635                {
636                    StorageKind::Dense
637                } else if IdxTensor::is_diag_axis_classes(&payload.axis_classes) {
638                    StorageKind::Diagonal
639                } else {
640                    StorageKind::Structured
641                }
642            }
643            Self::Deferred { source, .. } => source.storage_kind(),
644        }
645    }
646
647    fn materialize_eager_payload(inner: &EagerTensor) -> Result<NativeTensor> {
648        let native = inner.duplicate_value()?;
649        if native.is_col_major_contiguous()? {
650            return Ok(native);
651        }
652        let read = inner.tensor_read();
653        Ok(inner.runtime().with_execution_session(|session| {
654            session.to_contiguous_read(TensorRead::from_view(read.tensor_view()))
655        })??)
656    }
657
658    fn materialize(
659        &self,
660        logical_rank: usize,
661    ) -> std::result::Result<Arc<Storage>, TensorStorageError> {
662        match self {
663            Self::Materialized(storage) => Ok(Arc::clone(storage)),
664            Self::Eager {
665                inner,
666                axis_classes,
667            } => {
668                let native = Self::materialize_eager_payload(inner).map_err(|source| {
669                    TensorStorageError::Materialization {
670                        source: Arc::from(source.into_boxed_dyn_error()),
671                    }
672                })?;
673                let dtype = native.dtype();
674                if matches!(dtype, DType::F32 | DType::C32) {
675                    return Err(TensorStorageError::UnsupportedDtype {
676                        dtype: IdxTensor::dtype_name(dtype),
677                    });
678                }
679                IdxTensor::storage_from_native_with_axis_classes(
680                    &native,
681                    axis_classes,
682                    logical_rank,
683                )
684                .map(Arc::new)
685                .map_err(|source| TensorStorageError::Materialization {
686                    source: Arc::from(source.into_boxed_dyn_error()),
687                })
688            }
689            Self::Compact(payload) => {
690                let native =
691                    Self::materialize_eager_payload(&payload.payload).map_err(|source| {
692                        TensorStorageError::Materialization {
693                            source: Arc::from(source.into_boxed_dyn_error()),
694                        }
695                    })?;
696                let dtype = native.dtype();
697                if matches!(dtype, DType::F32 | DType::C32) {
698                    return Err(TensorStorageError::UnsupportedDtype {
699                        dtype: IdxTensor::dtype_name(dtype),
700                    });
701                }
702                IdxTensor::storage_from_native_with_axis_classes(
703                    &native,
704                    &payload.axis_classes,
705                    logical_rank,
706                )
707                .map(Arc::new)
708                .map_err(|source| TensorStorageError::Materialization {
709                    source: Arc::from(source.into_boxed_dyn_error()),
710                })
711            }
712            Self::Deferred { error, .. } => Err((**error).clone()),
713        }
714    }
715
716    fn scale_eager_payload(&self, scalar: &AnyScalar) -> Result<Self> {
717        let (payload, payload_dims, axis_classes) = match self {
718            Self::Materialized(storage) => {
719                let native = if storage.is_f64() {
720                    let values = storage
721                        .payload_f64_col_major_vec()
722                        .map_err(anyhow::Error::new)?;
723                    dense_native_tensor_from_col_major(&values, storage.payload_dims())?
724                } else {
725                    let values = storage
726                        .payload_c64_col_major_vec()
727                        .map_err(anyhow::Error::new)?;
728                    dense_native_tensor_from_col_major(&values, storage.payload_dims())?
729                };
730                (
731                    EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
732                    storage.payload_dims().to_vec(),
733                    storage.axis_classes().to_vec(),
734                )
735            }
736            Self::Eager {
737                inner,
738                axis_classes,
739            } => (
740                (**inner).clone(),
741                inner.shape().to_vec(),
742                axis_classes.clone(),
743            ),
744            Self::Compact(compact) => (
745                (*compact.payload).clone(),
746                compact.payload_dims.clone(),
747                compact.axis_classes.clone(),
748            ),
749            Self::Deferred { error, .. } => return Err(anyhow::Error::new((**error).clone())),
750        };
751        let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
752        let target_dtype = IdxTensor::scale_target_dtype(payload.dtype(), scalar_inner.dtype())?;
753        let payload = if payload.dtype() == target_dtype {
754            payload
755        } else {
756            payload.cast(target_dtype)?
757        };
758        let scalar_inner = if scalar_inner.dtype() == target_dtype {
759            scalar_inner.clone()
760        } else {
761            scalar_inner.cast(target_dtype)?
762        };
763        let scaled = if payload.shape().is_empty() {
764            payload.mul(&scalar_inner)?
765        } else {
766            let subscripts = IdxTensor::scale_subscripts(payload.shape().len())?;
767            [&payload, &scalar_inner].einsum_subscripts(&subscripts)?
768        };
769        match self {
770            Self::Eager { .. } => Ok(Self::Eager {
771                inner: Arc::new(scaled),
772                axis_classes,
773            }),
774            Self::Compact(_) | Self::Materialized(_) => {
775                Ok(Self::Compact(Arc::new(StructuredPayload {
776                    payload: Arc::new(scaled),
777                    payload_dims,
778                    axis_classes,
779                })))
780            }
781            Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
782        }
783    }
784
785    fn conjugate_with<F>(&self, conjugate: &F) -> std::result::Result<Self, TensorStorageError>
786    where
787        F: Fn(
788            &EagerTensor,
789        ) -> std::result::Result<
790            EagerTensor,
791            Arc<dyn std::error::Error + Send + Sync + 'static>,
792        >,
793    {
794        match self {
795            Self::Materialized(storage) => Ok(Self::Materialized(Arc::new(storage.conj()))),
796            Self::Eager {
797                inner,
798                axis_classes,
799            } => conjugate(inner)
800                .map(|conjugated| Self::Eager {
801                    inner: Arc::new(conjugated),
802                    axis_classes: axis_classes.clone(),
803                })
804                .map_err(|source| TensorStorageError::Conjugation { source }),
805            Self::Compact(payload) => conjugate(payload.payload.as_ref())
806                .map(|conjugated| {
807                    Self::Compact(Arc::new(StructuredPayload {
808                        payload: Arc::new(conjugated),
809                        payload_dims: payload.payload_dims.clone(),
810                        axis_classes: payload.axis_classes.clone(),
811                    }))
812                })
813                .map_err(|source| TensorStorageError::Conjugation { source }),
814            Self::Deferred { error, .. } => Err((**error).clone()),
815        }
816    }
817
818    fn sum_scalar(&self) -> Result<AnyScalar> {
819        match self {
820            Self::Materialized(storage) => {
821                if storage.is_f64() {
822                    Ok(AnyScalar::new_real(storage.sum::<f64>()))
823                } else {
824                    let value = storage.sum::<Complex64>();
825                    Ok(AnyScalar::new_complex(value.re, value.im))
826                }
827            }
828            Self::Eager { inner, .. } => IdxTensor::native_sum_scalar(inner),
829            Self::Compact(payload) => IdxTensor::native_sum_scalar(&payload.payload),
830            Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
831        }
832    }
833
834    fn nonfinite_flags(&self) -> Result<(bool, bool)> {
835        match self {
836            Self::Materialized(storage) => Ok(storage.payload_nonfinite_flags()),
837            Self::Eager { inner, .. } => IdxTensor::native_nonfinite_flags(inner),
838            Self::Compact(payload) => IdxTensor::native_nonfinite_flags(&payload.payload),
839            Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
840        }
841    }
842
843    fn payload_value_at(&self, payload_coords: &[usize]) -> Result<Complex64> {
844        match self {
845            Self::Materialized(storage) => storage
846                .scalar_at(payload_coords)
847                .map(Complex64::from)
848                .map_err(anyhow::Error::new),
849            Self::Eager { inner, .. } => {
850                IdxTensor::native_complex_payload_value_at(inner, payload_coords)
851            }
852            Self::Compact(payload) => {
853                IdxTensor::native_complex_payload_value_at(&payload.payload, payload_coords)
854            }
855            Self::Deferred { error, .. } => Err(anyhow::Error::new((**error).clone())),
856        }
857    }
858
859    fn for_each_payload_value(&self, mut f: impl FnMut(Complex64)) -> Result<()> {
860        let payload_dims = self.payload_dims();
861        let payload_len = checked_product(payload_dims)?;
862        let mut payload_coords = vec![0usize; payload_dims.len()];
863        for _ in 0..payload_len {
864            f(self.payload_value_at(&payload_coords)?);
865            let mut carry = true;
866            for (coordinate, &dim) in payload_coords.iter_mut().zip(payload_dims.iter()) {
867                if !carry {
868                    break;
869                }
870                *coordinate += 1;
871                if *coordinate == dim {
872                    *coordinate = 0;
873                } else {
874                    carry = false;
875                }
876            }
877        }
878        Ok(())
879    }
880
881    fn compact_payload(&self) -> Option<&StructuredPayload> {
882        match self {
883            Self::Compact(payload) => Some(payload.as_ref()),
884            Self::Deferred { source, .. } => source.compact_payload(),
885            _ => None,
886        }
887    }
888}
889
890/// Errors returned when constructing a compact copy-selector tensor.
891/// A copy-selector has logical values
892/// `scale * delta(left, right) * delta(site, selected_value)` and is used to
893/// carry a bond through a fixed physical site without dense bond-squared storage.
894/// # Examples
895/// ```
896/// use tensor4all_core::{DynIndex, StructuredSelectorError, IdxTensor};
897/// let left = DynIndex::new_dyn(2);
898/// let site = DynIndex::new_dyn(3);
899/// let right = DynIndex::new_dyn(4);
900/// let error = IdxTensor::from_copy_selector(left, site, right, 1, 1.0_f64)
901///     .unwrap_err();
902/// assert!(matches!(error, StructuredSelectorError::BondDimensionMismatch { .. }));
903/// ```
904#[derive(Debug, thiserror::Error)]
905pub enum StructuredSelectorError {
906    /// The two logical copy axes have different dimensions.
907    #[error("copy-selector bond dimensions differ: left={left}, right={right}")]
908    BondDimensionMismatch {
909        /// Dimension of the left copy axis.
910        left: usize,
911        /// Dimension of the right copy axis.
912        right: usize,
913    },
914    /// One of the logical axes has dimension zero.
915    #[error("copy-selector {axis} dimension must be positive")]
916    ZeroDimension {
917        /// Name of the zero-dimensional axis.
918        axis: &'static str,
919    },
920    /// The selected physical coordinate is outside the site dimension.
921    #[error("selected site value {value} is outside 0..{site_dim}")]
922    SelectedValueOutOfBounds {
923        /// Requested zero-based physical coordinate.
924        value: usize,
925        /// Dimension of the physical site.
926        site_dim: usize,
927    },
928    /// The compact payload element count cannot be represented by `usize`.
929    #[error("copy-selector payload size overflows usize for dimensions {bond_dim} x {site_dim}")]
930    PayloadSizeOverflow {
931        /// Dimension shared by the copy axes.
932        bond_dim: usize,
933        /// Dimension of the physical site.
934        site_dim: usize,
935    },
936    /// A compact payload stride cannot be represented by `isize`.
937    #[error("copy-selector bond stride {bond_dim} exceeds isize::MAX")]
938    StrideOverflow {
939        /// Bond dimension that could not be converted to a stride.
940        bond_dim: usize,
941    },
942    /// Reserving the compact payload failed.
943    #[error("could not allocate copy-selector payload with {elements} elements")]
944    AllocationFailed {
945        /// Number of compact payload elements requested.
946        elements: usize,
947    },
948    /// Backend structured-storage validation failed.
949    #[error("invalid copy-selector storage: {message}")]
950    InvalidStorage {
951        /// Diagnostic returned by structured-storage validation.
952        message: String,
953    },
954}
955
956/// Dynamic-rank tensor with structured payload storage -- the central data type
957/// of tensor4all.
958/// `IdxTensor` stores a logical multi-dimensional tensor of supported scalar
959/// values (`f32`, `f64`, `Complex32`, or `Complex64`) together with a list of
960/// [`DynIndex`] labels. `f64`/`Complex64` tensors may use compact [`Storage`]
961/// snapshots; `f32`/`Complex32` tensors retain an eager payload as the
962/// authoritative representation because compact storage supports only the
963/// 64-bit dtypes. The logical layout may be dense, diagonal, or explicitly
964/// structured. The indices carry unique
965/// identities (UUIDs) so that contraction, addition, and other binary
966/// operations can automatically match legs by identity rather than position.
967/// # Key Operations
968/// | Operation | Method |
969/// |-----------|--------|
970/// | Create from data | [`from_dense`](Self::from_dense), [`from_diag`](Self::from_diag), [`zeros`](Self::zeros) |
971/// | Extract data | [`to_vec`](Self::to_vec), [`into_dense_col_major_parts`](Self::into_dense_col_major_parts), [`sum`](Self::sum), [`only`](Self::only) |
972/// | Contraction | [`contract`](Self::contract) |
973/// | Arithmetic | [`add`](Self::add), [`scale`](Self::scale), [`axpby`](Self::axpby) |
974/// | Factorization | via [`TensorFactorizationLike::factorize`](crate::TensorFactorizationLike::factorize) |
975/// | Norms | [`norm`](Self::norm), [`norm_squared`](Self::norm_squared), [`maxabs`](Self::maxabs) |
976/// | Index ops | [`replaceind`](Self::replaceind), [`permute_indices`](Self::permute_indices) |
977/// # Data Layout
978/// Logical dense extraction uses **column-major** order (first index varies
979/// fastest), matching Fortran, Julia, and ITensors.jl conventions. Compact
980/// structured payloads additionally carry explicit payload dimensions, strides,
981/// and logical-axis classes.
982/// # Examples
983/// ```
984/// use tensor4all_core::{IdxTensor, DynIndex};
985/// // Create a 2x3 real tensor
986/// let i = DynIndex::new_dyn(2);
987/// let j = DynIndex::new_dyn(3);
988/// let data = vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0];
989/// let t = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
990/// assert_eq!(t.dims(), vec![2, 3]);
991/// assert!(t.is_f64());
992/// // Sum all elements: 1+2+3+4+5+6 = 21
993/// let s = t.sum().unwrap();
994/// assert!((s.real() - 21.0).abs() < 1e-12);
995/// // Extract data back out
996/// let data_out = t.to_vec::<f64>().unwrap();
997/// assert_eq!(data_out, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
998/// ```
999#[derive(Clone)]
1000pub struct IdxTensor {
1001    /// Full index information (includes tags and other metadata).
1002    pub indices: Vec<DynIndex>,
1003    /// Authoritative payload representation. Compact storage is used when the
1004    /// dtype is supported by [`Storage`]; otherwise this retains an eager
1005    /// payload without promotion.
1006    pub(crate) storage: IdxTensorStorage,
1007    /// Lazily materialized logical-dense eager payload for native execution and AD.
1008    pub(crate) eager_cache: Arc<OnceLock<Arc<EagerTensor>>>,
1009}
1010
1011impl IdxTensor {
1012    fn dense_axis_classes(rank: usize) -> Vec<usize> {
1013        (0..rank).collect()
1014    }
1015
1016    fn dtype_name(dtype: DType) -> &'static str {
1017        match dtype {
1018            DType::F32 => "f32",
1019            DType::F64 => "f64",
1020            DType::C32 => "c32",
1021            DType::C64 => "c64",
1022            DType::I32 => "i32",
1023            DType::I64 => "i64",
1024            DType::Bool => "bool",
1025        }
1026    }
1027
1028    fn scalar_dtype(&self) -> Result<DType> {
1029        if let Some(inner) = self.storage.eager() {
1030            return Ok(inner.dtype());
1031        }
1032        if self.storage.is_f64() {
1033            Ok(DType::F64)
1034        } else if self.storage.is_c64() {
1035            Ok(DType::C64)
1036        } else {
1037            Err(anyhow::anyhow!(
1038                "unable to determine IdxTensor scalar dtype"
1039            ))
1040        }
1041    }
1042
1043    fn diag_axis_classes(rank: usize) -> Vec<usize> {
1044        if rank == 0 {
1045            vec![]
1046        } else {
1047            vec![0; rank]
1048        }
1049    }
1050
1051    fn canonicalize_axis_classes(axis_classes: &[usize]) -> Vec<usize> {
1052        let mut map = std::collections::HashMap::new();
1053        let mut next = 0usize;
1054        axis_classes
1055            .iter()
1056            .map(|&class_id| {
1057                *map.entry(class_id).or_insert_with(|| {
1058                    let canonical = next;
1059                    next += 1;
1060                    canonical
1061                })
1062            })
1063            .collect()
1064    }
1065
1066    fn permute_axis_classes(&self, perm: &[usize]) -> Vec<usize> {
1067        let axis_classes = self.storage.axis_classes();
1068        let permuted: Vec<usize> = perm.iter().map(|&index| axis_classes[index]).collect();
1069        Self::canonicalize_axis_classes(&permuted)
1070    }
1071
1072    fn normalize_insert_axis(op: &str, axis: isize, rank: usize) -> Result<usize> {
1073        let normalized = if axis < 0 {
1074            rank as isize + 1 + axis
1075        } else {
1076            axis
1077        };
1078        if !(normalized >= 0 && normalized <= rank as isize) {
1079            return Err(anyhow::anyhow!(
1080                "{op}: axis {axis} is out of bounds for inserting into rank {rank}"
1081            ));
1082        };
1083        Ok(normalized as usize)
1084    }
1085
1086    fn is_diag_axis_classes(axis_classes: &[usize]) -> bool {
1087        axis_classes.len() >= 2 && axis_classes.iter().all(|&class_id| class_id == 0)
1088    }
1089
1090    fn validate_axis_classes(axis_classes: &[usize], rank: usize) -> Result<()> {
1091        if axis_classes.len() != rank {
1092            return Err(anyhow::anyhow!(
1093                "axis-class rank {} does not match tensor rank {rank}",
1094                axis_classes.len()
1095            ));
1096        }
1097        if Self::canonicalize_axis_classes(axis_classes) != axis_classes {
1098            return Err(anyhow::anyhow!(
1099                "axis classes must be canonical first-occurrence labels: {axis_classes:?}"
1100            ));
1101        }
1102        Ok(())
1103    }
1104
1105    fn einsum_subscripts_from_usize_ids(
1106        inputs: &[Vec<usize>],
1107        output: &[usize],
1108    ) -> Result<EinsumSubscripts> {
1109        let input_labels = inputs
1110            .iter()
1111            .map(|ids| {
1112                ids.iter()
1113                    .map(|&id| {
1114                        u32::try_from(id)
1115                            .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1116                    })
1117                    .collect::<Result<Vec<_>>>()
1118            })
1119            .collect::<Result<Vec<_>>>()?;
1120        let output_labels = output
1121            .iter()
1122            .map(|&id| {
1123                u32::try_from(id)
1124                    .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
1125            })
1126            .collect::<Result<Vec<_>>>()?;
1127        let input_refs = input_labels.iter().map(Vec::as_slice).collect::<Vec<_>>();
1128        Ok(EinsumSubscripts::new(&input_refs, &output_labels))
1129    }
1130
1131    fn build_binary_einsum_subscripts(
1132        lhs_rank: usize,
1133        axes_a: &[usize],
1134        rhs_rank: usize,
1135        axes_b: &[usize],
1136    ) -> Result<EinsumSubscripts> {
1137        if !(axes_a.len() == axes_b.len()) {
1138            return Err(anyhow::anyhow!(
1139                "contract axis length mismatch: lhs {:?}, rhs {:?}",
1140                axes_a,
1141                axes_b
1142            ));
1143        };
1144
1145        let mut lhs_ids = vec![usize::MAX; lhs_rank];
1146        let mut rhs_ids = vec![usize::MAX; rhs_rank];
1147        let mut next_id = 0usize;
1148
1149        let mut seen_lhs = vec![false; lhs_rank];
1150        let mut seen_rhs = vec![false; rhs_rank];
1151
1152        for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1153            if !(lhs_axis < lhs_rank) {
1154                return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1155            };
1156            if !(rhs_axis < rhs_rank) {
1157                return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1158            };
1159            if !(!seen_lhs[lhs_axis]) {
1160                return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1161            };
1162            if !(!seen_rhs[rhs_axis]) {
1163                return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1164            };
1165            seen_lhs[lhs_axis] = true;
1166            seen_rhs[rhs_axis] = true;
1167            lhs_ids[lhs_axis] = next_id;
1168            rhs_ids[rhs_axis] = next_id;
1169            next_id += 1;
1170        }
1171
1172        let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1173        for id in &mut lhs_ids {
1174            if *id == usize::MAX {
1175                *id = next_id;
1176                output_ids.push(next_id);
1177                next_id += 1;
1178            }
1179        }
1180        for id in &mut rhs_ids {
1181            if *id == usize::MAX {
1182                *id = next_id;
1183                output_ids.push(next_id);
1184                next_id += 1;
1185            }
1186        }
1187
1188        Self::einsum_subscripts_from_usize_ids(&[lhs_ids, rhs_ids], &output_ids)
1189    }
1190
1191    fn binary_dot_general_config(axes_a: &[usize], axes_b: &[usize]) -> Result<DotGeneralConfig> {
1192        if !(axes_a.len() == axes_b.len()) {
1193            return Err(anyhow::anyhow!(
1194                "contract axis length mismatch: lhs {:?}, rhs {:?}",
1195                axes_a,
1196                axes_b
1197            ));
1198        };
1199        Ok(DotGeneralConfig {
1200            lhs_contracting_dims: axes_a.to_vec(),
1201            rhs_contracting_dims: axes_b.to_vec(),
1202            lhs_batch_dims: vec![],
1203            rhs_batch_dims: vec![],
1204        })
1205    }
1206
1207    fn binary_contraction_axis_classes(
1208        lhs_axis_classes: &[usize],
1209        axes_a: &[usize],
1210        rhs_axis_classes: &[usize],
1211        axes_b: &[usize],
1212    ) -> Result<Vec<usize>> {
1213        debug_assert_eq!(axes_a.len(), axes_b.len());
1214
1215        fn find(parent: &mut [usize], value: usize) -> usize {
1216            if parent[value] != value {
1217                parent[value] = find(parent, parent[value]);
1218            }
1219            parent[value]
1220        }
1221
1222        fn union(parent: &mut [usize], lhs: usize, rhs: usize) {
1223            let lhs_root = find(parent, lhs);
1224            let rhs_root = find(parent, rhs);
1225            if lhs_root != rhs_root {
1226                parent[rhs_root] = lhs_root;
1227            }
1228        }
1229
1230        let lhs_payload_rank = match lhs_axis_classes.iter().copied().max() {
1231            Some(value) => value
1232                .checked_add(1)
1233                .ok_or_else(|| anyhow::anyhow!("left payload rank overflows usize"))?,
1234            None => 0,
1235        };
1236        let rhs_payload_rank = match rhs_axis_classes.iter().copied().max() {
1237            Some(value) => value
1238                .checked_add(1)
1239                .ok_or_else(|| anyhow::anyhow!("right payload rank overflows usize"))?,
1240            None => 0,
1241        };
1242        let rhs_offset = lhs_payload_rank;
1243        let parent_len = lhs_payload_rank
1244            .checked_add(rhs_payload_rank)
1245            .ok_or_else(|| anyhow::anyhow!("payload rank sum overflows usize"))?;
1246        let mut parent: Vec<usize> = (0..parent_len).collect();
1247
1248        for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1249            union(
1250                &mut parent,
1251                lhs_axis_classes[lhs_axis],
1252                rhs_offset
1253                    .checked_add(rhs_axis_classes[rhs_axis])
1254                    .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?,
1255            );
1256        }
1257
1258        let mut lhs_contracted = vec![false; lhs_axis_classes.len()];
1259        for &axis in axes_a {
1260            lhs_contracted[axis] = true;
1261        }
1262        let mut rhs_contracted = vec![false; rhs_axis_classes.len()];
1263        for &axis in axes_b {
1264            rhs_contracted[axis] = true;
1265        }
1266
1267        let mut root_to_class = std::collections::HashMap::new();
1268        let mut next_class = 0usize;
1269        let mut axis_classes = Vec::with_capacity(
1270            lhs_axis_classes
1271                .len()
1272                .saturating_add(rhs_axis_classes.len()),
1273        );
1274
1275        for (axis, &class_id) in lhs_axis_classes.iter().enumerate() {
1276            if !lhs_contracted[axis] {
1277                let root = find(&mut parent, class_id);
1278                let class = *root_to_class.entry(root).or_insert_with(|| {
1279                    let value = next_class;
1280                    next_class += 1;
1281                    value
1282                });
1283                axis_classes.push(class);
1284            }
1285        }
1286        for (axis, &class_id) in rhs_axis_classes.iter().enumerate() {
1287            if !rhs_contracted[axis] {
1288                let rhs_class = rhs_offset
1289                    .checked_add(class_id)
1290                    .ok_or_else(|| anyhow::anyhow!("rhs axis-class offset overflows usize"))?;
1291                let root = find(&mut parent, rhs_class);
1292                let class = *root_to_class.entry(root).or_insert_with(|| {
1293                    let value = next_class;
1294                    next_class += 1;
1295                    value
1296                });
1297                axis_classes.push(class);
1298            }
1299        }
1300
1301        Ok(axis_classes)
1302    }
1303
1304    fn scale_subscripts(rank: usize) -> Result<EinsumSubscripts> {
1305        let ids: Vec<usize> = (0..rank).collect();
1306        Self::einsum_subscripts_from_usize_ids(&[ids.clone(), Vec::new()], &ids)
1307    }
1308
1309    fn scale_target_dtype(payload: DType, scalar: DType) -> Result<DType> {
1310        let target = match payload {
1311            DType::F32 => match scalar {
1312                DType::C32 | DType::C64 => DType::C32,
1313                DType::F32 | DType::F64 => DType::F32,
1314                dtype => {
1315                    return Err(anyhow::anyhow!(
1316                        "unsupported scalar dtype {dtype:?} for f32 scaling"
1317                    ));
1318                }
1319            },
1320            DType::C32 => match scalar {
1321                DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C32,
1322                dtype => {
1323                    return Err(anyhow::anyhow!(
1324                        "unsupported scalar dtype {dtype:?} for c32 scaling"
1325                    ));
1326                }
1327            },
1328            DType::F64 => match scalar {
1329                DType::C32 | DType::C64 => DType::C64,
1330                DType::F32 | DType::F64 => DType::F64,
1331                dtype => {
1332                    return Err(anyhow::anyhow!(
1333                        "unsupported scalar dtype {dtype:?} for f64 scaling"
1334                    ));
1335                }
1336            },
1337            DType::C64 => match scalar {
1338                DType::F32 | DType::F64 | DType::C32 | DType::C64 => DType::C64,
1339                dtype => {
1340                    return Err(anyhow::anyhow!(
1341                        "unsupported scalar dtype {dtype:?} for c64 scaling"
1342                    ));
1343                }
1344            },
1345            dtype => {
1346                return Err(anyhow::anyhow!(
1347                    "unsupported tensor dtype {dtype:?} for scaling"
1348                ));
1349            }
1350        };
1351        Ok(target)
1352    }
1353
1354    fn validate_indices(indices: &[DynIndex]) -> Result<()> {
1355        let mut seen = HashSet::new();
1356        for idx in indices {
1357            if !(seen.insert(idx.clone())) {
1358                return Err(anyhow::anyhow!("Tensor indices must all be unique"));
1359            };
1360        }
1361        Ok(())
1362    }
1363
1364    fn validate_diag_dims(dims: &[usize]) -> Result<()> {
1365        if !dims.is_empty() {
1366            let first_dim = dims[0];
1367            for (i, &dim) in dims.iter().enumerate() {
1368                if !(dim == first_dim) {
1369                    return Err(anyhow::anyhow!("DiagTensor requires all indices to have the same dimension, but dims[{i}] = {dim} != dims[0] = {first_dim}"));
1370                };
1371            }
1372        }
1373        Ok(())
1374    }
1375
1376    fn seed_native_payload(storage: &Storage, dims: &[usize]) -> Result<NativeTensor> {
1377        Ok(storage_to_native_tensor(storage, dims)?)
1378    }
1379
1380    fn empty_eager_cache() -> Arc<OnceLock<Arc<EagerTensor>>> {
1381        Arc::new(OnceLock::new())
1382    }
1383
1384    fn eager_cache_with(inner: EagerTensor) -> Arc<OnceLock<Arc<EagerTensor>>> {
1385        let cache = Arc::new(OnceLock::new());
1386        let _ = cache.set(Arc::new(inner));
1387        cache
1388    }
1389
1390    fn compact_payload_inner(&self) -> Result<EagerTensor> {
1391        self.ensure_storage_ready()?;
1392        if let Some(inner) = self.storage.eager() {
1393            return Ok(inner.clone());
1394        }
1395        Ok(EagerTensor::from_tensor_in(
1396            storage_payload_native(self.storage.materialize(self.indices.len())?.as_ref())?,
1397            default_eager_ctx()?,
1398        )?)
1399    }
1400
1401    fn dense_inner_from_payload(
1402        payload: &EagerTensor,
1403        axis_classes: &[usize],
1404        logical_dims: &[usize],
1405    ) -> Result<EagerTensor> {
1406        let payload_rank = match axis_classes.iter().copied().max() {
1407            Some(class_id) => class_id
1408                .checked_add(1)
1409                .ok_or_else(|| anyhow::anyhow!("structured payload class rank overflows usize"))?,
1410            None => 0,
1411        };
1412        if !(payload.shape().len() == payload_rank) {
1413            return Err(anyhow::anyhow!(
1414                "structured payload rank {} does not match axis classes {:?}",
1415                payload.shape().len(),
1416                axis_classes
1417            ));
1418        };
1419        if !(logical_dims.len() == axis_classes.len()) {
1420            return Err(anyhow::anyhow!(
1421                "logical rank {} does not match axis class rank {}",
1422                logical_dims.len(),
1423                axis_classes.len()
1424            ));
1425        };
1426
1427        if axis_classes == Self::dense_axis_classes(logical_dims.len()) {
1428            if !(payload.shape() == logical_dims) {
1429                return Err(anyhow::anyhow!(
1430                    "dense payload dims {:?} do not match logical dims {:?}",
1431                    payload.shape(),
1432                    logical_dims
1433                ));
1434            };
1435            return Ok(payload.clone());
1436        }
1437
1438        let mut first_axis_by_class = vec![None; payload_rank];
1439        let mut dense = payload.clone();
1440        for (logical_axis, &class_id) in axis_classes.iter().enumerate() {
1441            let first_axis = match first_axis_by_class[class_id] {
1442                Some(first_axis) => first_axis,
1443                None => {
1444                    first_axis_by_class[class_id] = Some(logical_axis);
1445                    continue;
1446                }
1447            };
1448            dense = dense.embed_diag(first_axis, logical_axis)?;
1449        }
1450        if !(dense.shape() == logical_dims) {
1451            return Err(anyhow::anyhow!(
1452                "expanded structured payload dims {:?} do not match logical dims {:?}",
1453                dense.shape(),
1454                logical_dims
1455            ));
1456        };
1457        Ok(dense)
1458    }
1459
1460    fn tracked_compact_payload_value(&self) -> Option<&StructuredPayload> {
1461        self.storage
1462            .deferred_error()
1463            .is_none()
1464            .then_some(self.storage.compact_payload())
1465            .flatten()
1466            .filter(|value| value.payload.tracks_grad())
1467    }
1468
1469    fn ensure_storage_ready(&self) -> Result<()> {
1470        if let Some(error) = self.storage.deferred_error() {
1471            return Err(anyhow::Error::new(error.clone()));
1472        }
1473        Ok(())
1474    }
1475
1476    fn compact_payload_is_logical_dense(&self, payload_dims: &[usize]) -> bool {
1477        self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len())
1478            && payload_dims == self.dims()
1479    }
1480
1481    fn uses_tracked_compact_storage(&self) -> bool {
1482        self.tracked_compact_payload_value()
1483            .is_some_and(|value| !self.compact_payload_is_logical_dense(&value.payload_dims))
1484    }
1485
1486    fn ensure_shape_packing_preserves_ad(&self, op_name: &str) -> Result<()> {
1487        self.ensure_storage_ready()?;
1488        if !(!self.uses_tracked_compact_storage()) {
1489            return Err(anyhow::anyhow!("{op_name}: structured AD tensors with compact storage are not supported because materializing compact storage would detach gradients"));
1490        };
1491        Ok(())
1492    }
1493
1494    fn operand_indices_for_contraction(&self, conjugate: bool) -> Vec<DynIndex> {
1495        if conjugate {
1496            self.indices.iter().map(|index| index.conj()).collect()
1497        } else {
1498            self.indices.clone()
1499        }
1500    }
1501
1502    fn build_binary_contraction_labels(
1503        lhs_rank: usize,
1504        axes_a: &[usize],
1505        rhs_rank: usize,
1506        axes_b: &[usize],
1507    ) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
1508        if !(axes_a.len() == axes_b.len()) {
1509            return Err(anyhow::anyhow!(
1510                "contract axis length mismatch: lhs {:?}, rhs {:?}",
1511                axes_a,
1512                axes_b
1513            ));
1514        };
1515
1516        let mut lhs_ids = vec![usize::MAX; lhs_rank];
1517        let mut rhs_ids = vec![usize::MAX; rhs_rank];
1518        let mut next_id = 0usize;
1519
1520        let mut seen_lhs = vec![false; lhs_rank];
1521        let mut seen_rhs = vec![false; rhs_rank];
1522
1523        for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
1524            if !(lhs_axis < lhs_rank) {
1525                return Err(anyhow::anyhow!("lhs contract axis {lhs_axis} out of range"));
1526            };
1527            if !(rhs_axis < rhs_rank) {
1528                return Err(anyhow::anyhow!("rhs contract axis {rhs_axis} out of range"));
1529            };
1530            if !(!seen_lhs[lhs_axis]) {
1531                return Err(anyhow::anyhow!("duplicate lhs contract axis {lhs_axis}"));
1532            };
1533            if !(!seen_rhs[rhs_axis]) {
1534                return Err(anyhow::anyhow!("duplicate rhs contract axis {rhs_axis}"));
1535            };
1536            seen_lhs[lhs_axis] = true;
1537            seen_rhs[rhs_axis] = true;
1538            lhs_ids[lhs_axis] = next_id;
1539            rhs_ids[rhs_axis] = next_id;
1540            next_id += 1;
1541        }
1542
1543        let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
1544        for id in &mut lhs_ids {
1545            if *id == usize::MAX {
1546                *id = next_id;
1547                output_ids.push(next_id);
1548                next_id += 1;
1549            }
1550        }
1551        for id in &mut rhs_ids {
1552            if *id == usize::MAX {
1553                *id = next_id;
1554                output_ids.push(next_id);
1555                next_id += 1;
1556            }
1557        }
1558
1559        Ok((lhs_ids, rhs_ids, output_ids))
1560    }
1561
1562    fn build_payload_einsum_subscripts(
1563        input_roots: &[Vec<usize>],
1564        output_roots: &[usize],
1565    ) -> Result<EinsumSubscripts> {
1566        Self::einsum_subscripts_from_usize_ids(input_roots, output_roots)
1567    }
1568
1569    fn normalize_eager_payload_for_roots(
1570        payload: &EagerTensor,
1571        roots: &[usize],
1572    ) -> Result<(Option<EagerTensor>, Vec<usize>)> {
1573        if !(payload.shape().len() == roots.len()) {
1574            return Err(anyhow::anyhow!(
1575                "payload rank {} does not match root label count {}",
1576                payload.shape().len(),
1577                roots.len()
1578            ));
1579        };
1580
1581        let mut current_payload = None;
1582        let mut current_roots = roots.to_vec();
1583        while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(&current_roots) {
1584            let source = current_payload.as_ref().unwrap_or(payload);
1585            current_payload = Some(source.extract_diag(axis_a, axis_b)?);
1586            current_roots.remove(axis_b);
1587        }
1588
1589        Ok((current_payload, current_roots))
1590    }
1591
1592    fn first_duplicate_pair(values: &[usize]) -> Option<(usize, usize)> {
1593        let mut first_axis_by_value = std::collections::HashMap::new();
1594        for (axis, &value) in values.iter().enumerate() {
1595            if let Some(&first_axis) = first_axis_by_value.get(&value) {
1596                return Some((first_axis, axis));
1597            }
1598            first_axis_by_value.insert(value, axis);
1599        }
1600        None
1601    }
1602
1603    fn from_structured_payload_inner(
1604        indices: Vec<DynIndex>,
1605        payload_inner: EagerTensor,
1606        payload_dims: Vec<usize>,
1607        axis_classes: Vec<usize>,
1608    ) -> Result<Self> {
1609        Self::validate_indices(&indices)?;
1610        if payload_inner.shape() != payload_dims {
1611            return Err(anyhow::anyhow!(
1612                "structured payload dims {:?} do not match planned payload dims {:?}",
1613                payload_inner.shape(),
1614                payload_dims
1615            ));
1616        }
1617        if axis_classes == Self::dense_axis_classes(indices.len()) {
1618            return Self::from_inner_with_axis_classes(indices, payload_inner, axis_classes);
1619        }
1620        let structured_payload = Arc::new(StructuredPayload {
1621            payload: Arc::new(payload_inner),
1622            payload_dims,
1623            axis_classes,
1624        });
1625        Ok(Self {
1626            indices,
1627            storage: IdxTensorStorage::Compact(structured_payload),
1628            eager_cache: Self::empty_eager_cache(),
1629        })
1630    }
1631
1632    fn contract_structured_payloads(
1633        &self,
1634        other: &Self,
1635        result_indices: Vec<DynIndex>,
1636        axes_a: &[usize],
1637        axes_b: &[usize],
1638    ) -> Result<Self> {
1639        let (lhs_labels, rhs_labels, output_labels) = Self::build_binary_contraction_labels(
1640            self.indices.len(),
1641            axes_a,
1642            other.indices.len(),
1643            axes_b,
1644        )?;
1645        Self::contract_structured_payloads_nary(
1646            &[self, other],
1647            result_indices,
1648            vec![lhs_labels, rhs_labels],
1649            output_labels,
1650        )
1651    }
1652
1653    pub(crate) fn contract_structured_payloads_nary(
1654        operands: &[&Self],
1655        result_indices: Vec<DynIndex>,
1656        input_labels: Vec<Vec<usize>>,
1657        output_labels: Vec<usize>,
1658    ) -> Result<Self> {
1659        if !(!operands.is_empty()) {
1660            return Err(anyhow::anyhow!("structured contraction needs operands"));
1661        };
1662        for operand in operands {
1663            operand.ensure_storage_ready()?;
1664        }
1665        let layouts = operands
1666            .iter()
1667            .map(|operand| {
1668                OperandLayout::new(operand.dims(), operand.storage.axis_classes().to_vec())
1669            })
1670            .collect::<Result<Vec<_>>>()?;
1671        let spec = StructuredContractionSpec {
1672            input_labels,
1673            output_labels,
1674            retained_labels: Default::default(),
1675        };
1676        let plan = StructuredContractionPlan::new(&layouts, &spec)?;
1677        let any_grad = operands.iter().any(|operand| operand.tracks_grad());
1678
1679        if any_grad {
1680            let dtypes = operands
1681                .iter()
1682                .map(|operand| {
1683                    operand.storage.dtype().ok_or_else(|| {
1684                        anyhow::anyhow!("structured contraction operand has no scalar dtype")
1685                    })
1686                })
1687                .collect::<Result<Vec<_>>>()?;
1688            let target = Self::common_eager_dtype(&dtypes)?;
1689            let mut payloads = Vec::with_capacity(operands.len());
1690            for operand in operands {
1691                let payload = operand.compact_payload_inner()?;
1692                payloads.push(if payload.dtype() == target {
1693                    payload
1694                } else {
1695                    payload.cast(target)?
1696                });
1697            }
1698
1699            let mut normalized = Vec::with_capacity(payloads.len());
1700            let mut labels = Vec::with_capacity(payloads.len());
1701            for (operand_idx, (payload, operand_plan)) in
1702                payloads.iter().zip(plan.operand_plans.iter()).enumerate()
1703            {
1704                let (payload, roots) =
1705                    Self::normalize_eager_payload_for_roots(payload, &operand_plan.class_roots)?;
1706                normalized.push(payload.unwrap_or_else(|| payloads[operand_idx].clone()));
1707                labels.push(roots);
1708            }
1709            let refs = normalized.iter().collect::<Vec<_>>();
1710            let subscripts =
1711                Self::build_payload_einsum_subscripts(&labels, &plan.output_payload_roots)?;
1712            let payload = refs.as_slice().einsum_subscripts(&subscripts)?;
1713            return Self::from_structured_payload_inner(
1714                result_indices,
1715                payload,
1716                plan.output_payload_dims,
1717                plan.output_axis_classes,
1718            );
1719        }
1720
1721        // The native backend borrows contiguous compact payloads and promotes
1722        // only operands whose compact dtype differs. No logical dense tensor is
1723        // constructed on this path.
1724        let storage_owners = operands
1725            .iter()
1726            .map(|operand| match &operand.storage {
1727                IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1728                IdxTensorStorage::Deferred { source, .. } => match source.as_ref() {
1729                    IdxTensorStorage::Materialized(storage) => Some(Arc::clone(storage)),
1730                    _ => None,
1731                },
1732                _ => None,
1733            })
1734            .collect::<Vec<_>>();
1735        let mut inputs = Vec::with_capacity(operands.len());
1736        for (operand_idx, operand) in operands.iter().enumerate() {
1737            if let Some(storage) = storage_owners[operand_idx].as_ref() {
1738                inputs.push(storage_payload_native_read_input(storage.as_ref())?);
1739            } else {
1740                let inner = operand
1741                    .storage
1742                    .eager()
1743                    .ok_or_else(|| anyhow::anyhow!("structured operand has no compact payload"))?;
1744                inputs.push(NativeTensorReadInput::Borrowed(inner.tensor_read()));
1745            }
1746        }
1747        let mut normalized = Vec::with_capacity(inputs.len());
1748        let mut labels = Vec::with_capacity(inputs.len());
1749        for (input, operand_plan) in inputs.into_iter().zip(plan.operand_plans.iter()) {
1750            let (input, roots) =
1751                normalize_payload_read_for_roots(input, &operand_plan.class_roots)?;
1752            normalized.push(input);
1753            labels.push(roots);
1754        }
1755        let refs = normalized
1756            .iter()
1757            .zip(labels.iter())
1758            .map(|(input, labels)| (input, labels.as_slice()))
1759            .collect::<Vec<_>>();
1760        let payload = tensor4all_tensorbackend::einsum_native_tensor_reads(
1761            &refs,
1762            &plan.output_payload_roots,
1763        )?;
1764        // Wrap the result in the operands' common eager runtime when they share
1765        // one, so explicit-context tensors (e.g. SVD factors) stay in their
1766        // context instead of falling back to the process-global default.
1767        // Mixed runtimes keep the legacy default wrap; rejecting them is #623
1768        // scope, not this seam's.
1769        let mut common: Option<Arc<EagerRuntime>> = None;
1770        let mut mixed = false;
1771        for operand in operands {
1772            if let Some(inner) = operand.storage.eager() {
1773                match &common {
1774                    None => common = Some(Arc::clone(inner.runtime())),
1775                    Some(runtime) => {
1776                        if runtime.id() != inner.ctx_id() {
1777                            mixed = true;
1778                        }
1779                    }
1780                }
1781            }
1782        }
1783        let payload_inner = match (common, mixed) {
1784            (Some(runtime), false) => EagerTensor::from_tensor_in(payload, runtime)?,
1785            _ => EagerTensor::from_tensor_in(payload, default_eager_ctx()?)?,
1786        };
1787        Self::from_structured_payload_inner(
1788            result_indices,
1789            payload_inner,
1790            plan.output_payload_dims,
1791            plan.output_axis_classes,
1792        )
1793    }
1794
1795    fn common_eager_dtype(dtypes: &[DType]) -> Result<DType> {
1796        let target = if dtypes.contains(&DType::C64)
1797            || (dtypes.contains(&DType::C32) && dtypes.contains(&DType::F64))
1798        {
1799            DType::C64
1800        } else if dtypes.contains(&DType::F64) || dtypes.contains(&DType::C32) {
1801            if dtypes.contains(&DType::C32) {
1802                DType::C32
1803            } else {
1804                DType::F64
1805            }
1806        } else {
1807            DType::F32
1808        };
1809        if !(dtypes
1810            .iter()
1811            .all(|dtype| matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64)))
1812        {
1813            return Err(anyhow::anyhow!(
1814                "structured contraction supports only f32, f64, c32, and c64 operands"
1815            ));
1816        };
1817        Ok(target)
1818    }
1819
1820    fn should_use_structured_payload_contract(&self, other: &Self) -> bool {
1821        self.tracks_grad()
1822            || other.tracks_grad()
1823            || self.storage.axis_classes() != Self::dense_axis_classes(self.indices.len())
1824            || other.storage.axis_classes() != Self::dense_axis_classes(other.indices.len())
1825    }
1826
1827    fn storage_from_native_with_axis_classes(
1828        native: &NativeTensor,
1829        axis_classes: &[usize],
1830        logical_rank: usize,
1831    ) -> Result<Storage> {
1832        if matches!(native.dtype(), DType::F32 | DType::C32) {
1833            return Err(anyhow::anyhow!(
1834                "compact IdxTensor storage does not support dtype {:?}; retain the eager payload",
1835                native.dtype()
1836            ));
1837        }
1838        if Self::is_diag_axis_classes(axis_classes) {
1839            match native.dtype() {
1840                DType::F64 | DType::I32 | DType::I64 | DType::Bool => Storage::from_diag_col_major(
1841                    native_tensor_primal_to_diag::<f64>(native)?,
1842                    logical_rank,
1843                ),
1844                DType::C64 => Storage::from_diag_col_major(
1845                    native_tensor_primal_to_diag::<Complex64>(native)?,
1846                    logical_rank,
1847                ),
1848                DType::F32 | DType::C32 => Err(anyhow::anyhow!(
1849                    "compact IdxTensor storage does not support dtype {:?}",
1850                    native.dtype()
1851                )),
1852            }
1853        } else {
1854            storage_from_payload_native(native.duplicate()?, native.shape(), axis_classes.to_vec())
1855        }
1856    }
1857
1858    fn dense_selected_diag_payload<T: TensorElement + Copy + Zero>(
1859        payload: Vec<T>,
1860        kept_dims: &[usize],
1861        selected_positions: &[usize],
1862    ) -> Result<Vec<T>> {
1863        let output_len = checked_product(kept_dims)?;
1864        let mut data = vec![T::zero(); output_len];
1865        if output_len == 0 {
1866            return Ok(data);
1867        }
1868
1869        let Some((&first_position, rest)) = selected_positions.split_first() else {
1870            return Ok(data);
1871        };
1872        if rest.iter().any(|&position| position != first_position) {
1873            return Ok(data);
1874        }
1875
1876        let value = payload[first_position];
1877        if kept_dims.is_empty() {
1878            data[0] = value;
1879            return Ok(data);
1880        }
1881
1882        let mut offset = 0usize;
1883        let mut stride = 1usize;
1884        for &dim in kept_dims {
1885            let term = first_position
1886                .checked_mul(stride)
1887                .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1888            offset = offset
1889                .checked_add(term)
1890                .ok_or_else(|| anyhow::anyhow!("diagonal selection offset overflow"))?;
1891            stride = stride
1892                .checked_mul(dim)
1893                .ok_or_else(|| anyhow::anyhow!("diagonal selection stride overflow"))?;
1894        }
1895        data[offset] = value;
1896        Ok(data)
1897    }
1898
1899    fn select_diag_indices(
1900        &self,
1901        kept_indices: Vec<DynIndex>,
1902        kept_dims: Vec<usize>,
1903        positions: &[usize],
1904    ) -> Result<Self> {
1905        if self.storage.is_f64() {
1906            let storage = self.storage.materialize(self.indices.len())?;
1907            let payload = storage
1908                .payload_f64_col_major_vec()
1909                .map_err(anyhow::Error::new)?;
1910            let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1911            Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1912        } else if self.storage.is_c64() {
1913            let storage = self.storage.materialize(self.indices.len())?;
1914            let payload = storage
1915                .payload_c64_col_major_vec()
1916                .map_err(anyhow::Error::new)?;
1917            let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1918            Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1919        } else if self.storage.dtype() == Some(DType::F32) {
1920            let inner = self
1921                .storage
1922                .eager()
1923                .ok_or_else(|| anyhow::anyhow!("failed to read f32 diagonal payload"))?;
1924            let payload = inner.value()?.as_slice::<f32>()?.to_vec();
1925            let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1926            Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1927        } else if self.storage.dtype() == Some(DType::C32) {
1928            let inner = self
1929                .storage
1930                .eager()
1931                .ok_or_else(|| anyhow::anyhow!("failed to read c32 diagonal payload"))?;
1932            let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
1933            let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions)?;
1934            Self::from_dense(kept_indices, data).map_err(anyhow::Error::from)
1935        } else {
1936            Err(anyhow::anyhow!("unsupported diagonal storage scalar type"))
1937        }
1938    }
1939
1940    fn col_major_strides(dims: &[usize]) -> Result<Vec<isize>> {
1941        let mut strides = Vec::with_capacity(dims.len());
1942        let mut stride = 1isize;
1943        for &dim in dims {
1944            strides.push(stride);
1945            let dim = isize::try_from(dim)
1946                .map_err(|_| anyhow::anyhow!("dimension does not fit in isize"))?;
1947            stride = stride
1948                .checked_mul(dim)
1949                .ok_or_else(|| anyhow::anyhow!("column-major stride overflow"))?;
1950        }
1951        Ok(strides)
1952    }
1953
1954    fn zero_structured_selection<T>(
1955        kept_indices: Vec<DynIndex>,
1956        kept_dims: &[usize],
1957    ) -> Result<Self>
1958    where
1959        T: TensorElement + Zero,
1960    {
1961        let output_len = checked_product(kept_dims)?;
1962        Self::from_dense(kept_indices, vec![T::zero(); output_len]).map_err(anyhow::Error::from)
1963    }
1964
1965    fn selected_structured_class_positions(
1966        axis_classes: &[usize],
1967        payload_rank: usize,
1968        selected_axes: &[usize],
1969        positions: &[usize],
1970    ) -> Option<Vec<Option<usize>>> {
1971        let mut selected_class_positions = vec![None; payload_rank];
1972        for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
1973            let class_id = axis_classes[axis];
1974            if let Some(existing) = selected_class_positions[class_id] {
1975                if existing != position {
1976                    return None;
1977                }
1978            } else {
1979                selected_class_positions[class_id] = Some(position);
1980            }
1981        }
1982        Some(selected_class_positions)
1983    }
1984
1985    fn select_structured_indices_typed<T, F>(
1986        &self,
1987        payload: Vec<T>,
1988        kept_axes: &[usize],
1989        kept_indices: Vec<DynIndex>,
1990        kept_dims: Vec<usize>,
1991        selected: (&[usize], &[usize]),
1992        make_output: F,
1993    ) -> Result<Self>
1994    where
1995        T: TensorElement + Zero,
1996        F: FnOnce(Vec<T>, Vec<usize>, Vec<isize>, Vec<usize>) -> Result<Self>,
1997    {
1998        let (selected_axes, positions) = selected;
1999        let payload_dims = self.storage.payload_dims();
2000        let axis_classes = self.storage.axis_classes();
2001        let payload_rank = payload_dims.len();
2002        let Some(selected_class_positions) = Self::selected_structured_class_positions(
2003            axis_classes,
2004            payload_rank,
2005            selected_axes,
2006            positions,
2007        ) else {
2008            return Self::zero_structured_selection::<T>(kept_indices, &kept_dims);
2009        };
2010
2011        let selected_class_kept = kept_axes
2012            .iter()
2013            .any(|&axis| selected_class_positions[axis_classes[axis]].is_some());
2014        if selected_class_kept {
2015            return self.select_structured_indices_dense(
2016                payload,
2017                kept_axes,
2018                kept_indices,
2019                kept_dims,
2020                &selected_class_positions,
2021            );
2022        }
2023
2024        let mut old_to_new_class = vec![None; payload_rank];
2025        let mut output_payload_dims = Vec::new();
2026        let mut output_axis_classes = Vec::with_capacity(kept_axes.len());
2027        for &axis in kept_axes {
2028            let class_id = axis_classes[axis];
2029            let new_class = match old_to_new_class[class_id] {
2030                Some(new_class) => new_class,
2031                None => {
2032                    let new_class = output_payload_dims.len();
2033                    old_to_new_class[class_id] = Some(new_class);
2034                    output_payload_dims.push(payload_dims[class_id]);
2035                    new_class
2036                }
2037            };
2038            output_axis_classes.push(new_class);
2039        }
2040
2041        let output_len = checked_product(&output_payload_dims)?;
2042        let mut output_payload = Vec::with_capacity(output_len);
2043        for linear in 0..output_len {
2044            let output_payload_index = decode_col_major_linear(linear, &output_payload_dims)?;
2045            let mut input_payload_index = vec![0usize; payload_rank];
2046            for class_id in 0..payload_rank {
2047                input_payload_index[class_id] =
2048                    if let Some(position) = selected_class_positions[class_id] {
2049                        position
2050                    } else if let Some(new_class) = old_to_new_class[class_id] {
2051                        output_payload_index[new_class]
2052                    } else {
2053                        return Err(anyhow::anyhow!(
2054                            "structured payload class {class_id} is neither selected nor kept"
2055                        ));
2056                    };
2057            }
2058            let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2059            output_payload.push(payload[input_linear]);
2060        }
2061
2062        let output_strides = Self::col_major_strides(&output_payload_dims)?;
2063        make_output(
2064            output_payload,
2065            output_payload_dims,
2066            output_strides,
2067            output_axis_classes,
2068        )
2069    }
2070
2071    fn select_structured_indices_dense<T>(
2072        &self,
2073        payload: Vec<T>,
2074        kept_axes: &[usize],
2075        kept_indices: Vec<DynIndex>,
2076        kept_dims: Vec<usize>,
2077        selected_class_positions: &[Option<usize>],
2078    ) -> Result<Self>
2079    where
2080        T: TensorElement + Zero,
2081    {
2082        let payload_dims = self.storage.payload_dims();
2083        let axis_classes = self.storage.axis_classes();
2084        let output_len = checked_product(&kept_dims)?;
2085        let mut output = Vec::with_capacity(output_len);
2086
2087        for linear in 0..output_len {
2088            let kept_position = decode_col_major_linear(linear, &kept_dims)?;
2089            let mut input_payload_index = selected_class_positions.to_vec();
2090            let mut is_structural_zero = false;
2091
2092            for (&axis, &position) in kept_axes.iter().zip(kept_position.iter()) {
2093                let class_id = axis_classes[axis];
2094                match input_payload_index[class_id] {
2095                    Some(existing) if existing != position => {
2096                        is_structural_zero = true;
2097                        break;
2098                    }
2099                    Some(_) => {}
2100                    None => input_payload_index[class_id] = Some(position),
2101                }
2102            }
2103
2104            if is_structural_zero {
2105                output.push(T::zero());
2106                continue;
2107            }
2108
2109            let input_payload_index = input_payload_index
2110                .into_iter()
2111                .enumerate()
2112                .map(|(class_id, position)| {
2113                    position.ok_or_else(|| {
2114                        anyhow::anyhow!(
2115                            "structured payload class {class_id} is neither selected nor kept"
2116                        )
2117                    })
2118                })
2119                .collect::<Result<Vec<_>>>()?;
2120            let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
2121            output.push(payload[input_linear]);
2122        }
2123
2124        Self::from_dense(kept_indices, output).map_err(anyhow::Error::from)
2125    }
2126
2127    fn select_structured_indices(
2128        &self,
2129        kept_axes: &[usize],
2130        kept_indices: Vec<DynIndex>,
2131        kept_dims: Vec<usize>,
2132        selected_axes: &[usize],
2133        positions: &[usize],
2134    ) -> Result<Self> {
2135        if self.storage.is_f64() {
2136            let storage = self.storage.materialize(self.indices.len())?;
2137            let payload = storage
2138                .payload_f64_col_major_vec()
2139                .map_err(anyhow::Error::new)?;
2140            let output_indices = kept_indices.clone();
2141            self.select_structured_indices_typed(
2142                payload,
2143                kept_axes,
2144                kept_indices,
2145                kept_dims,
2146                (selected_axes, positions),
2147                move |payload, dims, strides, classes| {
2148                    let storage = Storage::new_structured(payload, dims, strides, classes)?;
2149                    Self::from_storage(output_indices, Arc::new(storage))
2150                        .map_err(anyhow::Error::from)
2151                },
2152            )
2153        } else if self.storage.is_c64() {
2154            let storage = self.storage.materialize(self.indices.len())?;
2155            let payload = storage
2156                .payload_c64_col_major_vec()
2157                .map_err(anyhow::Error::new)?;
2158            let output_indices = kept_indices.clone();
2159            self.select_structured_indices_typed(
2160                payload,
2161                kept_axes,
2162                kept_indices,
2163                kept_dims,
2164                (selected_axes, positions),
2165                move |payload, dims, strides, classes| {
2166                    let storage = Storage::new_structured(payload, dims, strides, classes)?;
2167                    Self::from_storage(output_indices, Arc::new(storage))
2168                        .map_err(anyhow::Error::from)
2169                },
2170            )
2171        } else if self.storage.dtype() == Some(DType::F32) {
2172            let inner = self
2173                .storage
2174                .eager()
2175                .ok_or_else(|| anyhow::anyhow!("failed to read f32 structured payload"))?;
2176            let payload = inner.value()?.as_slice::<f32>()?.to_vec();
2177            let output_indices = kept_indices.clone();
2178            self.select_structured_indices_typed(
2179                payload,
2180                kept_axes,
2181                kept_indices,
2182                kept_dims,
2183                (selected_axes, positions),
2184                move |payload, dims, _strides, classes| {
2185                    let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2186                    let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2187                    Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2188                },
2189            )
2190        } else if self.storage.dtype() == Some(DType::C32) {
2191            let inner = self
2192                .storage
2193                .eager()
2194                .ok_or_else(|| anyhow::anyhow!("failed to read c32 structured payload"))?;
2195            let payload = inner.value()?.as_slice::<Complex32>()?.to_vec();
2196            let output_indices = kept_indices.clone();
2197            self.select_structured_indices_typed(
2198                payload,
2199                kept_axes,
2200                kept_indices,
2201                kept_dims,
2202                (selected_axes, positions),
2203                move |payload, dims, _strides, classes| {
2204                    let native = dense_native_tensor_from_col_major(&payload, &dims)?;
2205                    let inner = EagerTensor::from_tensor_in(native, default_eager_ctx()?)?;
2206                    Self::from_structured_payload_inner(output_indices, inner, dims, classes)
2207                },
2208            )
2209        } else {
2210            Err(anyhow::anyhow!(
2211                "unsupported structured storage scalar type"
2212            ))
2213        }
2214    }
2215
2216    fn validate_storage_matches_indices(indices: &[DynIndex], storage: &Storage) -> Result<()> {
2217        let dims = Self::expected_dims_from_indices(indices);
2218        let storage_dims = storage.logical_dims();
2219        if storage_dims != dims {
2220            return Err(anyhow::anyhow!(
2221                "storage logical dims {:?} do not match indices dims {:?}",
2222                storage_dims,
2223                dims
2224            ));
2225        }
2226        if storage.is_diag() {
2227            Self::validate_diag_dims(&dims)?;
2228        }
2229        Ok(())
2230    }
2231
2232    fn try_materialized_inner(&self) -> Result<&EagerTensor> {
2233        self.ensure_storage_ready()?;
2234        let logical_dims = self.dims();
2235        if let Some(value) = self.tracked_compact_payload_value() {
2236            if self.compact_payload_is_logical_dense(&value.payload_dims) {
2237                return Ok(value.payload.as_ref());
2238            }
2239            if self.eager_cache.get().is_none() {
2240                let dense = Self::dense_inner_from_payload(
2241                    value.payload.as_ref(),
2242                    &value.axis_classes,
2243                    &logical_dims,
2244                )?;
2245                let _ = self.eager_cache.set(Arc::new(dense));
2246            }
2247            return self
2248                .eager_cache
2249                .get()
2250                .map(|inner| inner.as_ref())
2251                .ok_or_else(|| {
2252                    anyhow::anyhow!("IdxTensor structured AD cache was not initialized")
2253                });
2254        }
2255        if let Some(inner) = self.storage.eager() {
2256            if self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len()) {
2257                return Ok(inner);
2258            }
2259            if self.eager_cache.get().is_none() {
2260                let dense = Self::dense_inner_from_payload(
2261                    inner,
2262                    self.storage.axis_classes(),
2263                    &logical_dims,
2264                )?;
2265                let _ = self.eager_cache.set(Arc::new(dense));
2266            }
2267            return self
2268                .eager_cache
2269                .get()
2270                .map(|inner| inner.as_ref())
2271                .ok_or_else(|| {
2272                    anyhow::anyhow!("IdxTensor structured eager cache was not initialized")
2273                });
2274        }
2275        if self.eager_cache.get().is_none() {
2276            let native = profile_pairwise_contract_section("materialize_storage_to_native", || {
2277                let storage = self.storage.materialize(self.indices.len())?;
2278                Self::seed_native_payload(storage.as_ref(), &logical_dims)
2279            })
2280            .context("IdxTensor materialization failed")?;
2281            record_pairwise_contract_profile_bytes(
2282                "materialize_storage_to_native",
2283                tensor_profile_bytes(native.dtype(), native.shape()),
2284            );
2285            let _ = self.eager_cache.set(Arc::new(EagerTensor::from_tensor_in(
2286                native,
2287                default_eager_ctx()?,
2288            )?));
2289        }
2290        self.eager_cache
2291            .get()
2292            .map(|inner| inner.as_ref())
2293            .ok_or_else(|| anyhow::anyhow!("IdxTensor materialization cache was not initialized"))
2294    }
2295
2296    pub(crate) fn as_inner(&self) -> Result<&EagerTensor> {
2297        self.try_materialized_inner()
2298    }
2299
2300    /// Compute dims from `indices` order.
2301    #[inline]
2302    fn expected_dims_from_indices(indices: &[DynIndex]) -> Vec<usize> {
2303        indices.iter().map(|idx| idx.dim()).collect()
2304    }
2305
2306    /// Get dims in the current `indices` order.
2307    ///
2308    /// This is computed on-demand from `indices` (single source of truth).
2309    ///
2310    /// # Examples
2311    ///
2312    /// ```
2313    /// use tensor4all_core::{DynIndex, IdxTensor};
2314    ///
2315    /// let i = DynIndex::new_dyn(2);
2316    /// let j = DynIndex::new_dyn(3);
2317    /// let k = DynIndex::new_dyn(4);
2318    /// let t = IdxTensor::from_dense(
2319    ///     vec![i, j, k],
2320    ///     vec![0.0; 24],
2321    /// ).unwrap();
2322    /// assert_eq!(t.dims(), vec![2, 3, 4]);
2323    /// ```
2324    pub fn dims(&self) -> Vec<usize> {
2325        Self::expected_dims_from_indices(&self.indices)
2326    }
2327
2328    /// Select fixed coordinates for tensor indices and drop those axes.
2329    ///
2330    /// The `selected_indices` slice identifies tensor axes by index identity,
2331    /// and `positions` gives the zero-based coordinate to take on each
2332    /// selected axis. Unselected indices are preserved in their original order.
2333    ///
2334    /// # Arguments
2335    ///
2336    /// * `selected_indices` - Indices to fix and remove from the result. Each
2337    ///
2338    ///   index must appear exactly once in this tensor.
2339    /// * `positions` - Coordinates for `selected_indices`. Each coordinate must
2340    ///
2341    ///   be less than the corresponding index dimension.
2342    ///
2343    /// # Returns
2344    ///
2345    /// A tensor over the unselected indices. Selecting no indices returns a
2346    /// clone of the original tensor. Selecting all indices returns a rank-0
2347    /// scalar tensor. Diagonal and structured tensors are sliced from their
2348    /// compact payload without materializing the original full tensor; the
2349    /// result keeps structured storage when the remaining logical axes can
2350    /// still be represented by axis classes.
2351    ///
2352    /// # Errors
2353    /// Returns an error when a selected coordinate is out of range for its index
2354    /// (an out of bounds failure) or when `selected_indices` and `positions`
2355    /// differ in length (a length mismatch).
2356    /// # Examples
2357    ///
2358    /// ```
2359    /// use tensor4all_core::{DynIndex, IdxTensor};
2360    ///
2361    /// let i = DynIndex::new_dyn(2);
2362    /// let j = DynIndex::new_dyn(3);
2363    /// let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
2364    /// let tensor = IdxTensor::from_dense(vec![i.clone(), j.clone()], data).unwrap();
2365    ///
2366    /// let selected = tensor.select_indices(&[j], &[1]).unwrap();
2367    /// assert_eq!(selected.dims(), vec![2]);
2368    /// assert_eq!(selected.to_vec::<f64>().unwrap(), vec![3.0, 4.0]);
2369    /// ```
2370    pub fn select_indices(
2371        &self,
2372        selected_indices: &[DynIndex],
2373        positions: &[usize],
2374    ) -> std::result::Result<Self, IdxTensorError> {
2375        if selected_indices.len() != positions.len() {
2376            return Err(anyhow::anyhow!(
2377                "selected_indices length {} does not match positions length {}",
2378                selected_indices.len(),
2379                positions.len()
2380            )
2381            .into());
2382        }
2383        if selected_indices.is_empty() {
2384            return Ok(self.clone());
2385        }
2386
2387        let mut selected_axes = Vec::with_capacity(selected_indices.len());
2388        let mut seen_axes = HashSet::with_capacity(selected_indices.len());
2389        for (selected, &position) in selected_indices.iter().zip(positions.iter()) {
2390            let axis = self
2391                .indices
2392                .iter()
2393                .position(|index| index == selected)
2394                .ok_or_else(|| anyhow::anyhow!("selected index is not present in tensor"))?;
2395            if !seen_axes.insert(axis) {
2396                return Err(anyhow::anyhow!("selected index appears more than once").into());
2397            }
2398            let dim = self.indices[axis].dim();
2399            if position >= dim {
2400                return Err(anyhow::anyhow!(
2401                    "selected coordinate {position} is out of range for axis {axis} with dim {dim}"
2402                )
2403                .into());
2404            }
2405            selected_axes.push(axis);
2406        }
2407
2408        let kept_axes = self
2409            .indices
2410            .iter()
2411            .enumerate()
2412            .filter(|(axis, _)| !seen_axes.contains(axis))
2413            .map(|(axis, _)| axis)
2414            .collect::<Vec<_>>();
2415        let kept_indices = kept_axes
2416            .iter()
2417            .map(|&axis| self.indices[axis].clone())
2418            .collect::<Vec<_>>();
2419        let kept_dims = kept_axes
2420            .iter()
2421            .map(|&axis| self.indices[axis].dim())
2422            .collect::<Vec<_>>();
2423
2424        if matches!(
2425            self.storage.storage_kind(),
2426            StorageKind::Diagonal | StorageKind::Structured
2427        ) {
2428            self.ensure_shape_packing_preserves_ad("select_indices")?;
2429        }
2430        if self.storage.storage_kind() == StorageKind::Diagonal {
2431            return self
2432                .select_diag_indices(kept_indices, kept_dims, positions)
2433                .map_err(IdxTensorError::from);
2434        }
2435        if self.storage.storage_kind() == StorageKind::Structured {
2436            return self
2437                .select_structured_indices(
2438                    &kept_axes,
2439                    kept_indices,
2440                    kept_dims,
2441                    &selected_axes,
2442                    positions,
2443                )
2444                .map_err(IdxTensorError::from);
2445        }
2446        if self.storage.storage_kind() != StorageKind::Dense {
2447            return Err(anyhow::anyhow!(
2448                "select_indices got unsupported storage kind {:?}",
2449                self.storage.storage_kind()
2450            )
2451            .into());
2452        }
2453
2454        // Slice positionally axis by axis in the operand's own runtime.
2455        // (A starts-tensor + dynamic_slice build would need an explicit
2456        // host-to-device upload for resident operands, which this
2457        // context-free helper cannot name.)
2458        let mut sliced = self.try_materialized_inner()?.clone();
2459        for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
2460            sliced = sliced
2461                .slice_axis(axis, position..position + 1)
2462                .map_err(|error| anyhow::anyhow!("select_indices slicing failed: {error}"))?;
2463        }
2464        Self::from_inner(kept_indices, sliced.reshape(&kept_dims)?).map_err(IdxTensorError::from)
2465    }
2466
2467    /// Stack tensors along a newly inserted index.
2468    ///
2469    /// Each input must have exactly the same index order and dimensions. The
2470    /// `new_index` dimension must match the number of input tensors. The
2471    /// `axis` argument follows tenferro/PyTorch-style insertion semantics:
2472    /// `0` inserts before the first existing axis and `-1` appends a trailing
2473    /// axis. Use `axis = -1` for batched contractions because tenferro uses
2474    /// trailing batch dimensions as the canonical batched-GEMM layout.
2475    ///
2476    /// # Errors
2477    ///
2478    /// Returns an error if no tensors are provided, the new index dimension
2479    /// does not match the number of tensors, an input has a different index
2480    /// order, `axis` is outside the valid insertion range, or a tracked
2481    /// structured-AD tensor uses compact storage that would need dense
2482    /// materialization.
2483    ///
2484    /// # Examples
2485    ///
2486    /// ```
2487    /// use tensor4all_core::{DynIndex, IdxTensor};
2488    ///
2489    /// let i = DynIndex::new_dyn(2);
2490    /// let batch = DynIndex::new_dyn(2);
2491    /// let a = IdxTensor::from_dense(vec![i.clone()], vec![1.0_f64, 2.0]).unwrap();
2492    /// let b = IdxTensor::from_dense(vec![i.clone()], vec![3.0_f64, 4.0]).unwrap();
2493    ///
2494    /// let stacked = IdxTensor::stack_along_new_index(&[&a, &b], batch.clone(), -1).unwrap();
2495    ///
2496    /// assert_eq!(stacked.indices(), &[i, batch]);
2497    /// assert_eq!(stacked.to_vec::<f64>().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
2498    /// ```
2499    pub fn stack_along_new_index(
2500        tensors: &[&Self],
2501        new_index: DynIndex,
2502        axis: isize,
2503    ) -> std::result::Result<Self, IdxTensorError> {
2504        let first = tensors
2505            .first()
2506            .copied()
2507            .ok_or_else(|| anyhow::anyhow!("stack_along_new_index requires at least one tensor"))?;
2508        if !(new_index.dim() == tensors.len()) {
2509            return Err(anyhow::anyhow!(
2510                "stack_along_new_index: new index dim {} does not match tensor count {}",
2511                new_index.dim(),
2512                tensors.len()
2513            )
2514            .into());
2515        };
2516
2517        let base_indices = first.indices.clone();
2518        for tensor in tensors.iter().copied().skip(1) {
2519            if !(tensor.indices == base_indices) {
2520                return Err(anyhow::anyhow!(
2521                    "stack_along_new_index: input tensors must have identical index order"
2522                )
2523                .into());
2524            };
2525        }
2526        for &tensor in tensors {
2527            tensor.ensure_shape_packing_preserves_ad("stack_along_new_index")?;
2528        }
2529
2530        let insert_axis =
2531            Self::normalize_insert_axis("stack_along_new_index", axis, base_indices.len())?;
2532        let mut result_indices = base_indices;
2533        result_indices.insert(insert_axis, new_index);
2534
2535        let inner_refs = tensors
2536            .iter()
2537            .map(|tensor| tensor.try_materialized_inner())
2538            .collect::<Result<Vec<_>>>()?;
2539        let stacked = EagerTensor::stack(&inner_refs, axis)?;
2540        Self::from_inner(result_indices, stacked).map_err(IdxTensorError::from)
2541    }
2542
2543    /// Select positions along one index and replace it with a new index.
2544    ///
2545    /// This is the retained-axis counterpart to [`Self::select_indices`]:
2546    /// instead of fixing one coordinate and removing the index, it gathers a
2547    /// list of positions and keeps the gathered axis under `target_index`.
2548    /// Repeated positions are allowed; reverse-mode AD accumulates repeated
2549    /// cotangents through tenferro's scatter-add gather transpose.
2550    ///
2551    /// # Errors
2552    /// Returns an error when a selected position is out of range for the source
2553    /// index (an out of bounds failure) or when the source and target index
2554    /// dimensions are incompatible (a shape mismatch).
2555    /// # Examples
2556    ///
2557    /// ```
2558    /// use tensor4all_core::{DynIndex, IdxTensor};
2559    ///
2560    /// let source = DynIndex::new_dyn(3);
2561    /// let target = DynIndex::new_dyn(2);
2562    /// let tensor = IdxTensor::from_dense(
2563    ///     vec![source.clone()],
2564    ///     vec![10.0_f64, 20.0, 30.0],
2565    /// ).unwrap();
2566    ///
2567    /// let selected = tensor.index_select(&source, target.clone(), &[2, 0]).unwrap();
2568    ///
2569    /// assert_eq!(selected.indices(), &[target]);
2570    /// assert_eq!(selected.to_vec::<f64>().unwrap(), vec![30.0, 10.0]);
2571    /// ```
2572    pub fn index_select(
2573        &self,
2574        source_index: &DynIndex,
2575        target_index: DynIndex,
2576        positions: &[usize],
2577    ) -> std::result::Result<Self, IdxTensorError> {
2578        if !(target_index.dim() == positions.len()) {
2579            return Err(anyhow::anyhow!(
2580                "index_select: target index dim {} does not match position count {}",
2581                target_index.dim(),
2582                positions.len()
2583            )
2584            .into());
2585        };
2586        let axis = self
2587            .indices
2588            .iter()
2589            .position(|index| index == source_index)
2590            .ok_or_else(|| anyhow::anyhow!("index_select: source index is not present"))?;
2591        let source_dim = self.indices[axis].dim();
2592        for &position in positions {
2593            if !(position < source_dim) {
2594                return Err(anyhow::anyhow!(
2595                    "index_select: position {position} is out of range for source dim {source_dim}"
2596                )
2597                .into());
2598            };
2599        }
2600        self.ensure_shape_packing_preserves_ad("index_select")?;
2601
2602        let axis = isize::try_from(axis)
2603            .map_err(|_| anyhow::anyhow!("index_select: axis does not fit in isize"))?;
2604        let selected = self
2605            .try_materialized_inner()?
2606            .index_select(axis, positions)?;
2607        let mut result_indices = self.indices.clone();
2608        result_indices[axis as usize] = target_index;
2609        Self::from_inner(result_indices, selected).map_err(IdxTensorError::from)
2610    }
2611
2612    /// Create a new tensor with dynamic rank.
2613    ///
2614    /// # Errors
2615    /// Returns an error when the storage logical dimension does not match the
2616    /// index dimension product (a shape mismatch) or when duplicate
2617    /// indices are provided.
2618    /// # Examples
2619    ///
2620    /// ```
2621    /// use tensor4all_core::{DynIndex, IdxTensor};
2622    /// use tensor4all_tensorbackend::Storage;
2623    /// use std::sync::Arc;
2624    ///
2625    /// let i = DynIndex::new_dyn(3);
2626    /// let storage = Arc::new(Storage::new_dense::<f64>(3).unwrap());
2627    /// let t = IdxTensor::new(vec![i], storage).unwrap();
2628    /// assert_eq!(t.dims(), vec![3]);
2629    /// ```
2630    pub fn new(
2631        indices: Vec<DynIndex>,
2632        storage: Arc<Storage>,
2633    ) -> std::result::Result<Self, IdxTensorError> {
2634        Self::from_storage(indices, storage)
2635    }
2636
2637    /// Create a new tensor with dynamic rank, automatically computing dimensions from indices.
2638    ///
2639    /// This is a convenience constructor that extracts dimensions from indices using `IndexLike::dim()`.
2640    ///
2641    /// # Errors
2642    /// Returns an error when the storage logical dimension does not match the
2643    /// index dimension product (a shape mismatch) or when duplicate
2644    /// indices are provided.
2645    /// # Examples
2646    ///
2647    /// ```
2648    /// use tensor4all_core::{DynIndex, IdxTensor};
2649    /// use tensor4all_tensorbackend::Storage;
2650    /// use std::sync::Arc;
2651    ///
2652    /// let i = DynIndex::new_dyn(4);
2653    /// let storage = Arc::new(Storage::new_dense::<f64>(4).unwrap());
2654    /// let t = IdxTensor::from_indices(vec![i], storage).unwrap();
2655    /// assert_eq!(t.dims(), vec![4]);
2656    /// ```
2657    pub fn from_indices(
2658        indices: Vec<DynIndex>,
2659        storage: Arc<Storage>,
2660    ) -> std::result::Result<Self, IdxTensorError> {
2661        Self::new(indices, storage)
2662    }
2663
2664    /// Create a tensor from explicit compact storage.
2665    ///
2666    /// # Errors
2667    /// Returns an error when the storage scalar kind is incompatible with the
2668    /// requested operations (a scalar-kind mismatch) or the storage cannot
2669    /// represent the given index space (a shape mismatch).
2670    /// # Examples
2671    ///
2672    /// ```
2673    /// use tensor4all_core::{DynIndex, IdxTensor};
2674    /// use tensor4all_tensorbackend::Storage;
2675    /// use std::sync::Arc;
2676    ///
2677    /// let i = DynIndex::new_dyn(2);
2678    /// let j = DynIndex::new_dyn(2);
2679    /// let storage = Arc::new(Storage::new_diag(vec![1.0_f64, 2.0]).unwrap());
2680    /// let t = IdxTensor::from_storage(vec![i, j], storage).unwrap();
2681    /// assert_eq!(t.dims(), vec![2, 2]);
2682    /// ```
2683    pub fn from_storage(
2684        indices: Vec<DynIndex>,
2685        storage: Arc<Storage>,
2686    ) -> std::result::Result<Self, IdxTensorError> {
2687        Self::validate_indices(&indices)?;
2688        Self::validate_storage_matches_indices(&indices, storage.as_ref())?;
2689        Ok(Self {
2690            indices,
2691            storage: IdxTensorStorage::from_storage(storage),
2692            eager_cache: Self::empty_eager_cache(),
2693        })
2694    }
2695
2696    /// Create a tensor from explicit structured storage.
2697    ///
2698    /// This is an alias for [`IdxTensor::from_storage`] with a name that
2699    /// emphasizes that compact structured metadata is preserved.
2700    ///
2701    /// # Errors
2702    /// Returns an error when the structured storage is invalid (an invalid-storage
2703    /// failure) or the index space is incompatible (a shape mismatch).
2704    /// # Examples
2705    ///
2706    /// ```
2707    /// use std::sync::Arc;
2708    /// use tensor4all_core::{DynIndex, IdxTensor};
2709    /// use tensor4all_tensorbackend::{Storage, StorageKind};
2710    ///
2711    /// let i = DynIndex::new_dyn(2);
2712    /// let j = DynIndex::new_dyn(2);
2713    /// let storage = Arc::new(Storage::from_diag_col_major(vec![1.0_f64, 2.0], 2).unwrap());
2714    /// let tensor = IdxTensor::from_structured_storage(vec![i, j], storage).unwrap();
2715    /// assert_eq!(tensor.storage().unwrap().storage_kind(), StorageKind::Diagonal);
2716    /// ```
2717    pub fn from_structured_storage(
2718        indices: Vec<DynIndex>,
2719        storage: Arc<Storage>,
2720    ) -> std::result::Result<Self, IdxTensorError> {
2721        Self::from_storage(indices, storage)
2722    }
2723
2724    /// Construct a compact copy tensor that selects one physical-site value.
2725    ///
2726    /// The returned rank-3 tensor has logical indices `[left, site, right]` and
2727    /// value `scale` exactly when `left == right` and `site == selected_value`;
2728    /// every other entry is zero. Its payload has `left.dim * site.dim`
2729    /// elements rather than `left.dim * site.dim * right.dim` dense elements.
2730    ///
2731    /// # Arguments
2732    ///
2733    /// - `left`: left copy axis; its dimension must be positive and equal to
2734    ///
2735    ///   `right.dim`.
2736    /// - `site`: physical axis whose selected coordinate remains active.
2737    /// - `right`: right copy axis paired with `left`.
2738    /// - `selected_value`: zero-based coordinate in `0..site.dim`.
2739    /// - `scale`: value stored on the selected copy diagonal.
2740    ///
2741    /// # Returns
2742    ///
2743    /// A structured tensor with axis classes `[0, 1, 0]`. For `f64` and
2744    /// `Complex64`, compact storage is retained; `f32` and `Complex32` keep
2745    /// an eager authoritative payload because compact storage has no 32-bit
2746    /// scalar representation.
2747    ///
2748    /// # Errors
2749    ///
2750    /// Returns [`StructuredSelectorError`] when dimensions are zero or
2751    /// inconsistent, the selected value is out of bounds, checked size or
2752    /// stride arithmetic overflows, allocation fails, or backend structured
2753    /// storage validation rejects the metadata.
2754    ///
2755    /// # Examples
2756    ///
2757    /// ```
2758    /// use tensor4all_core::{DynIndex, IdxTensor};
2759    /// use tensor4all_tensorbackend::StorageKind;
2760    ///
2761    /// let left = DynIndex::new_dyn(2);
2762    /// let site = DynIndex::new_dyn(3);
2763    /// let right = DynIndex::new_dyn(2);
2764    /// let tensor = IdxTensor::from_copy_selector(
2765    ///     left,
2766    ///     site,
2767    ///     right,
2768    ///     1,
2769    ///     2.5_f64,
2770    /// ).unwrap();
2771    ///
2772    /// assert_eq!(tensor.storage().unwrap().storage_kind(), StorageKind::Structured);
2773    /// assert_eq!(tensor.storage().unwrap().payload_len(), 6);
2774    /// assert_eq!(
2775    ///     tensor.to_vec::<f64>().unwrap(),
2776    ///     vec![0.0, 0.0, 2.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.5, 0.0, 0.0],
2777    /// );
2778    /// ```
2779    pub fn from_copy_selector<T>(
2780        left: DynIndex,
2781        site: DynIndex,
2782        right: DynIndex,
2783        selected_value: usize,
2784        scale: T,
2785    ) -> std::result::Result<Self, StructuredSelectorError>
2786    where
2787        T: TensorElement + Copy + Zero,
2788    {
2789        if left.dim == 0 {
2790            return Err(StructuredSelectorError::ZeroDimension { axis: "left" });
2791        }
2792        if site.dim == 0 {
2793            return Err(StructuredSelectorError::ZeroDimension { axis: "site" });
2794        }
2795        if right.dim == 0 {
2796            return Err(StructuredSelectorError::ZeroDimension { axis: "right" });
2797        }
2798        if left.dim != right.dim {
2799            return Err(StructuredSelectorError::BondDimensionMismatch {
2800                left: left.dim,
2801                right: right.dim,
2802            });
2803        }
2804        if selected_value >= site.dim {
2805            return Err(StructuredSelectorError::SelectedValueOutOfBounds {
2806                value: selected_value,
2807                site_dim: site.dim,
2808            });
2809        }
2810
2811        let payload_len =
2812            left.dim
2813                .checked_mul(site.dim)
2814                .ok_or(StructuredSelectorError::PayloadSizeOverflow {
2815                    bond_dim: left.dim,
2816                    site_dim: site.dim,
2817                })?;
2818        let _site_stride = isize::try_from(left.dim)
2819            .map_err(|_| StructuredSelectorError::StrideOverflow { bond_dim: left.dim })?;
2820        let selected_offset = left.dim.checked_mul(selected_value).ok_or(
2821            StructuredSelectorError::PayloadSizeOverflow {
2822                bond_dim: left.dim,
2823                site_dim: site.dim,
2824            },
2825        )?;
2826
2827        let mut payload = Vec::new();
2828        payload.try_reserve_exact(payload_len).map_err(|_| {
2829            StructuredSelectorError::AllocationFailed {
2830                elements: payload_len,
2831            }
2832        })?;
2833        payload.resize(payload_len, T::zero());
2834        for bond in 0..left.dim {
2835            payload[selected_offset + bond] = scale;
2836        }
2837        let payload_native = dense_native_tensor_from_col_major(&payload, &[left.dim, site.dim])
2838            .map_err(|error| StructuredSelectorError::InvalidStorage {
2839                message: error.to_string(),
2840            })?;
2841        let payload_dtype = payload_native.dtype();
2842        let payload_inner = EagerTensor::from_tensor_in(
2843            payload_native,
2844            default_eager_ctx().map_err(|error| StructuredSelectorError::InvalidStorage {
2845                message: error.to_string(),
2846            })?,
2847        )
2848        .map_err(|error| StructuredSelectorError::InvalidStorage {
2849            message: error.to_string(),
2850        })?;
2851        let payload_dims = vec![left.dim, site.dim];
2852        let indices = vec![left, site, right];
2853        if !matches!(
2854            payload_dtype,
2855            DType::F32 | DType::F64 | DType::C32 | DType::C64
2856        ) {
2857            return Err(StructuredSelectorError::InvalidStorage {
2858                message: format!("unsupported selector dtype {:?}", payload_dtype),
2859            });
2860        }
2861        Self::from_structured_payload_inner(indices, payload_inner, payload_dims, vec![0, 1, 0])
2862            .map_err(|error| StructuredSelectorError::InvalidStorage {
2863                message: error.to_string(),
2864            })
2865    }
2866
2867    /// Create a tensor from a native tenferro payload.
2868    pub(crate) fn from_native(indices: Vec<DynIndex>, native: NativeTensor) -> Result<Self> {
2869        let axis_classes = Self::dense_axis_classes(indices.len());
2870        Self::from_native_with_axis_classes(indices, native, axis_classes)
2871    }
2872
2873    pub(crate) fn from_native_with_axis_classes(
2874        indices: Vec<DynIndex>,
2875        native: NativeTensor,
2876        axis_classes: Vec<usize>,
2877    ) -> Result<Self> {
2878        Self::from_inner_with_axis_classes(
2879            indices,
2880            EagerTensor::from_tensor_in(native, default_eager_ctx()?)?,
2881            axis_classes,
2882        )
2883    }
2884
2885    /// Adopt a backend-derived value whose result indices and axis classes were
2886    /// validated by the contraction planner that produced it.
2887    pub(crate) fn from_untracked_native_with_axis_classes(
2888        indices: Vec<DynIndex>,
2889        native: NativeTensor,
2890        axis_classes: Vec<usize>,
2891    ) -> Result<Self> {
2892        Self::from_inner_with_validated_metadata(
2893            indices,
2894            adopt_untracked_eager_value(default_eager_ctx()?, TensorValue::from_tensor(native))?,
2895            axis_classes,
2896        )
2897    }
2898
2899    pub(crate) fn from_inner(indices: Vec<DynIndex>, inner: EagerTensor) -> Result<Self> {
2900        let axis_classes = Self::dense_axis_classes(indices.len());
2901        Self::from_inner_with_axis_classes(indices, inner, axis_classes)
2902    }
2903
2904    /// Compute the Hermitian eigendecomposition of a rank-2 tensor.
2905    ///
2906    /// The tensor must have two square matrix axes. The returned eigenvectors
2907    /// stay in [`IdxTensor`] form so downstream tensor algebra can preserve
2908    /// AD metadata where the backend supports it. Eigenvalues are returned as
2909    /// detached real primal values because truncation and rank selection are
2910    /// nonsmooth control-flow decisions.
2911    ///
2912    /// `hermitian_tol` controls the allowed imaginary part of complex
2913    /// eigenvalues after the backend solve; use a small non-negative value such
2914    /// as `1e-12` for numerically Hermitian inputs.
2915    ///
2916    /// # Errors
2917    /// Returns an error when the tensor is not rank-2, when the two indices have
2918    /// unequal dimensions (a shape or shape mismatch), or when the
2919    /// eigensolver fails to converge (a non-convergence failure).
2920    /// # Examples
2921    ///
2922    /// ```
2923    /// use tensor4all_core::{AnyScalar, DynIndex, TensorContractionLike, IdxTensor};
2924    ///
2925    /// let row = DynIndex::new_dyn(2);
2926    /// let col = DynIndex::new_dyn(2);
2927    /// let matrix = IdxTensor::from_dense(
2928    ///     vec![row.clone(), col.clone()],
2929    ///     vec![3.0_f64, 0.0, 0.0, 5.0],
2930    /// ).unwrap();
2931    ///
2932    /// let decomp = matrix.hermitian_eigendecomposition(1.0e-12).unwrap();
2933    /// let eigenvector = decomp
2934    ///     .eigenvectors
2935    ///     .select_indices(&[decomp.eigenvector_index.clone()], &[0])
2936    ///     .unwrap();
2937    /// let eigenvector_as_col = eigenvector.replaceind(&row, &col).unwrap();
2938    /// let applied = IdxTensor::contract(&[&matrix, &eigenvector_as_col]).unwrap();
2939    /// let expected = eigenvector.scale(AnyScalar::new_real(decomp.eigenvalues[0])).unwrap();
2940    ///
2941    /// assert!(applied.isapprox(&expected, 1.0e-12, 0.0).unwrap());
2942    /// ```
2943    pub fn hermitian_eigendecomposition(
2944        &self,
2945        hermitian_tol: f64,
2946    ) -> std::result::Result<TensorHermitianEigendecomposition, IdxTensorError> {
2947        if !(self.indices.len() == 2) {
2948            return Err(anyhow::anyhow!(
2949                "IdxTensor::hermitian_eigendecomposition requires a rank-2 tensor, got rank {}",
2950                self.indices.len()
2951            )
2952            .into());
2953        };
2954        let dims = self.dims();
2955        if !(dims[0] == dims[1]) {
2956            return Err(anyhow::anyhow!(
2957                "IdxTensor::hermitian_eigendecomposition requires a square matrix, got {}x{}",
2958                dims[0],
2959                dims[1]
2960            )
2961            .into());
2962        };
2963        if !(dims[0] > 0) {
2964            return Err(anyhow::anyhow!(
2965                "IdxTensor::hermitian_eigendecomposition requires a non-empty matrix"
2966            )
2967            .into());
2968        };
2969        if !(hermitian_tol.is_finite() && hermitian_tol >= 0.0) {
2970            return Err(anyhow::anyhow!(
2971                "IdxTensor::hermitian_eigendecomposition requires a finite non-negative tolerance"
2972            )
2973            .into());
2974        };
2975
2976        let input = self.try_materialized_inner()?;
2977        let (values, vectors) = input
2978            .eigh()
2979            .map_err(|source| anyhow::anyhow!("Hermitian eigendecomposition failed: {source}"))?;
2980
2981        let eigenvalue_index = DynIndex::new_dyn(dims[0]);
2982        let eigenvector_index = DynIndex::new_dyn(dims[0]);
2983        let eigenvalue_tensor = Self::from_inner(vec![eigenvalue_index], values)?;
2984        let eigenvalues = Self::read_real_eigenvalues(&eigenvalue_tensor, hermitian_tol)
2985            .with_context(|| {
2986                "IdxTensor::hermitian_eigendecomposition failed to read eigenvalues"
2987            })?;
2988        let eigenvectors = Self::from_inner(
2989            vec![self.indices[0].clone(), eigenvector_index.clone()],
2990            vectors,
2991        )?;
2992
2993        Ok(TensorHermitianEigendecomposition {
2994            eigenvalues,
2995            eigenvectors,
2996            eigenvector_index,
2997        })
2998    }
2999
3000    fn read_real_eigenvalues(values: &Self, hermitian_tol: f64) -> Result<Vec<f64>> {
3001        if values.is_complex() {
3002            values
3003                .to_vec::<Complex64>()?
3004                .into_iter()
3005                .enumerate()
3006                .map(|(index, value)| {
3007                    let imaginary = value.im.abs();
3008                    let allowed = hermitian_tol * value.norm().max(1.0);
3009                    if !matches!(
3010                        imaginary.partial_cmp(&allowed),
3011                        Some(std::cmp::Ordering::Less) | Some(std::cmp::Ordering::Equal)
3012                    ) {
3013            return Err(anyhow::anyhow!("Hermitian eigenvalue {index} has imaginary part {imaginary}, exceeding tolerance {allowed}"));
3014        };
3015                    Ok(value.re)
3016                })
3017                .collect()
3018        } else {
3019            values.to_vec::<f64>().map_err(anyhow::Error::from)
3020        }
3021    }
3022
3023    pub(crate) fn from_diag_inner(
3024        indices: Vec<DynIndex>,
3025        payload_inner: EagerTensor,
3026    ) -> Result<Self> {
3027        let dims = Self::expected_dims_from_indices(&indices);
3028        Self::validate_indices(&indices)?;
3029        Self::validate_diag_dims(&dims)?;
3030        let payload_len = checked_product(payload_inner.shape())?;
3031        Self::validate_diag_payload_len(payload_len, &dims)?;
3032        let axis_classes = Self::diag_axis_classes(dims.len());
3033        let diag_inner = payload_inner.embed_diag(0, 1)?;
3034        Self::from_inner_with_axis_classes(indices, diag_inner, axis_classes)
3035    }
3036
3037    fn compact_inner_from_logical(
3038        inner: &EagerTensor,
3039        axis_classes: &[usize],
3040    ) -> Result<EagerTensor> {
3041        let mut payload = inner.clone();
3042        let mut classes = axis_classes.to_vec();
3043        while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(&classes) {
3044            payload = payload.extract_diag(axis_a, axis_b)?;
3045            classes.remove(axis_b);
3046        }
3047        Ok(payload)
3048    }
3049
3050    pub(crate) fn from_inner_with_axis_classes(
3051        indices: Vec<DynIndex>,
3052        inner: EagerTensor,
3053        axis_classes: Vec<usize>,
3054    ) -> Result<Self> {
3055        Self::from_inner_with_axis_classes_impl(indices, inner, axis_classes, false)
3056    }
3057
3058    fn from_inner_with_validated_metadata(
3059        indices: Vec<DynIndex>,
3060        inner: EagerTensor,
3061        axis_classes: Vec<usize>,
3062    ) -> Result<Self> {
3063        Self::from_inner_with_axis_classes_impl(indices, inner, axis_classes, true)
3064    }
3065
3066    fn from_inner_with_axis_classes_impl(
3067        indices: Vec<DynIndex>,
3068        inner: EagerTensor,
3069        axis_classes: Vec<usize>,
3070        validated: bool,
3071    ) -> Result<Self> {
3072        let dims = profile_pairwise_contract_section("from_inner_expected_dims", || {
3073            Self::expected_dims_from_indices(&indices)
3074        });
3075        let dense_axis_classes = axis_classes.iter().copied().eq(0..indices.len());
3076        if validated {
3077            debug_assert!(Self::validate_indices(&indices).is_ok());
3078            if !dense_axis_classes {
3079                Self::validate_axis_classes(&axis_classes, indices.len())?;
3080            }
3081        } else {
3082            profile_pairwise_contract_section("from_inner_validate_indices", || {
3083                Self::validate_indices(&indices)
3084            })?;
3085            Self::validate_axis_classes(&axis_classes, indices.len())?;
3086        }
3087        if dims != inner.shape() {
3088            return Err(anyhow::anyhow!(
3089                "native payload dims {:?} do not match indices dims {:?}",
3090                inner.shape(),
3091                dims
3092            ));
3093        }
3094        if Self::is_diag_axis_classes(&axis_classes) {
3095            profile_pairwise_contract_section("from_inner_validate_diag_dims", || {
3096                Self::validate_diag_dims(&dims)
3097            })?;
3098        }
3099        let storage = if dense_axis_classes {
3100            IdxTensorStorage::from_eager_dense(inner, indices.len())
3101        } else {
3102            let payload = Self::compact_inner_from_logical(&inner, &axis_classes)?;
3103            let payload_dims = payload.shape().to_vec();
3104            IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3105                payload: Arc::new(payload),
3106                payload_dims,
3107                axis_classes,
3108            }))
3109        };
3110        Ok(Self {
3111            indices,
3112            storage,
3113            eager_cache: Self::empty_eager_cache(),
3114        })
3115    }
3116
3117    /// Borrow the indices.
3118    pub fn indices(&self) -> &[DynIndex] {
3119        &self.indices
3120    }
3121
3122    pub(crate) fn axis_classes(&self) -> &[usize] {
3123        self.storage.axis_classes()
3124    }
3125
3126    #[cfg(feature = "tenferro-cuda")]
3127    pub(crate) fn deferred_storage_error(&self) -> Option<&TensorStorageError> {
3128        self.storage.deferred_error()
3129    }
3130
3131    #[cfg(feature = "tenferro-cuda")]
3132    pub(crate) fn cuda_eager_inner(&self) -> Option<&EagerTensor> {
3133        self.storage
3134            .eager()
3135            .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3136    }
3137
3138    /// Borrow the owning eager runtime without materializing or transferring.
3139    ///
3140    /// Returns `None` for tensors with no eager value. Used to keep
3141    /// contraction results in their operands' context.
3142    pub(crate) fn eager_runtime(&self) -> Option<Arc<EagerRuntime>> {
3143        self.storage
3144            .eager()
3145            .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3146            .map(|inner| Arc::clone(inner.runtime()))
3147    }
3148
3149    /// Check whether this tensor's eager value is device-resident.
3150    ///
3151    /// Metadata-only: inspects placement without reading data, transferring,
3152    /// or consulting any execution context. Used to route factorization to the
3153    /// context-scoped path and to reject context-free calls on resident inputs
3154    /// before algorithm work begins.
3155    #[cfg(feature = "tenferro-cuda")]
3156    pub(crate) fn is_cuda_resident(&self) -> bool {
3157        self.cuda_eager_inner().is_some_and(|inner| {
3158            inner.tensor_read().placement().memory_kind == tenferro_tensor::MemoryKind::Device
3159        })
3160    }
3161
3162    /// Check whether this tensor's eager value is device-resident.
3163    ///
3164    /// Without the CUDA feature no tensor can be device-resident.
3165    #[cfg(not(feature = "tenferro-cuda"))]
3166    pub(crate) fn is_cuda_resident(&self) -> bool {
3167        false
3168    }
3169
3170    #[cfg(feature = "tenferro-cuda")]
3171    pub(crate) fn cuda_duplicate_native(&self) -> Result<NativeTensor> {
3172        Ok(self.try_materialized_inner()?.duplicate_value()?)
3173    }
3174
3175    /// Enable reverse-mode AD tracking on this tensor by creating a tracked leaf.
3176    /// # Errors
3177    /// Returns an error when the tensor is not a scalar (a rank mismatch) or the
3178    /// AD backend cannot track the tensor's dtype.
3179    ///
3180    pub fn enable_grad(self) -> std::result::Result<Self, IdxTensorError> {
3181        self.ensure_storage_ready()?;
3182        // Keep the eager payload when available: compact Storage currently
3183        // stores only f64/C64 and must not promote f32/C32 leaves before AD.
3184        let eager_payload = self
3185            .storage
3186            .eager()
3187            .or_else(|| self.eager_cache.get().map(AsRef::as_ref))
3188            .filter(|inner| inner.shape() == self.storage.payload_dims());
3189        let payload = match eager_payload {
3190            Some(inner) => inner.duplicate_value()?,
3191            None => {
3192                let materialized = self.storage.materialize(self.indices.len())?;
3193                storage_payload_native(materialized.as_ref())
3194                    .context("IdxTensor::enable_grad failed")?
3195            }
3196        };
3197        let payload_dims = self.storage.payload_dims().to_vec();
3198        let axis_classes = self.storage.axis_classes().to_vec();
3199        let tracked = Arc::new(EagerTensor::requires_grad_in(
3200            payload,
3201            default_eager_ctx()?,
3202        )?);
3203        let storage = if axis_classes == Self::dense_axis_classes(self.indices.len()) {
3204            IdxTensorStorage::Eager {
3205                inner: tracked,
3206                axis_classes,
3207            }
3208        } else {
3209            IdxTensorStorage::Compact(Arc::new(StructuredPayload {
3210                payload: tracked,
3211                payload_dims,
3212                axis_classes,
3213            }))
3214        };
3215        Ok(Self {
3216            indices: self.indices,
3217            storage,
3218            eager_cache: Self::empty_eager_cache(),
3219        })
3220    }
3221
3222    /// Report whether this tensor participates in gradient tracking.
3223    pub fn tracks_grad(&self) -> bool {
3224        self.storage.eager().is_some_and(EagerTensor::tracks_grad)
3225            || self
3226                .eager_cache
3227                .get()
3228                .is_some_and(|inner| inner.tracks_grad())
3229    }
3230
3231    /// Return the accumulated gradient, if one has been stored.
3232    /// # Errors
3233    /// Returns an error when the tensor is not a tracked leaf or the gradient is
3234    /// unavailable for the tensor's dtype (an unavailable-gradient failure).
3235    ///
3236    pub fn grad(&self) -> std::result::Result<Option<Self>, IdxTensorError> {
3237        if let Some(value) = self.tracked_compact_payload_value() {
3238            let Some(gradient) = value.payload.grad()? else {
3239                return Ok(None);
3240            };
3241            let gradient_shape = gradient.shape().to_vec();
3242            let gradient_tensor = gradient.to_tensor()?;
3243            if self.compact_payload_is_logical_dense(&value.payload_dims) {
3244                return Ok(Some(Self::from_native_with_axis_classes(
3245                    self.indices.clone(),
3246                    gradient_tensor,
3247                    value.axis_classes.clone(),
3248                )?));
3249            }
3250            if gradient_shape != value.payload_dims {
3251                return Err(anyhow::anyhow!(
3252                    "gradient payload dims {:?} do not match {:?}",
3253                    gradient_shape,
3254                    value.payload_dims
3255                )
3256                .into());
3257            }
3258            let gradient = EagerTensor::from_tensor_in(gradient_tensor, default_eager_ctx()?)?;
3259            return Ok(Some(Self::from_structured_payload_inner(
3260                self.indices.clone(),
3261                gradient,
3262                value.payload_dims.clone(),
3263                value.axis_classes.clone(),
3264            )?));
3265        }
3266
3267        let Some(gradient) = self.try_materialized_inner()?.grad()? else {
3268            return Ok(None);
3269        };
3270        Ok(Some(Self::from_native_with_axis_classes(
3271            self.indices.clone(),
3272            gradient.to_tensor()?,
3273            self.storage.axis_classes().to_vec(),
3274        )?))
3275    }
3276
3277    /// Clear the accumulated gradient stored for this tensor.
3278    /// # Errors
3279    /// Returns an error when the tensor is not a tracked leaf (a missing-graph
3280    /// failure).
3281    ///
3282    pub fn clear_grad(&self) -> std::result::Result<(), IdxTensorError> {
3283        self.ensure_storage_ready()?;
3284        if let Some(value) = self.tracked_compact_payload_value() {
3285            value.payload.clear_grad()?;
3286        }
3287        if let Some(inner) = self.storage.eager() {
3288            inner.clear_grad()?;
3289        }
3290        if let Some(inner) = self.eager_cache.get() {
3291            inner.clear_grad()?;
3292        }
3293        Ok(())
3294    }
3295
3296    /// Run reverse-mode autodiff from this scalar tensor.
3297    /// # Errors
3298    /// Returns an error when the tensor is not a scalar (a rank mismatch) or the
3299    /// reverse pass fails (a graph failure).
3300    ///
3301    pub fn backward(&self) -> std::result::Result<(), IdxTensorError> {
3302        if let Some(value) = self.tracked_compact_payload_value() {
3303            return value.payload.backward().map(|_| ()).map_err(|e| {
3304                IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}"))
3305            });
3306        }
3307        self.try_materialized_inner()?
3308            .backward()
3309            .map(|_| ())
3310            .map_err(|e| IdxTensorError::from(anyhow::anyhow!("IdxTensor::backward failed: {e}")))
3311    }
3312
3313    /// Detach this tensor from the reverse graph.
3314    /// # Errors
3315    /// Returns an error when the tensor is not a tracked leaf (a missing-graph
3316    /// failure).
3317    ///
3318    pub fn detach(&self) -> std::result::Result<Self, IdxTensorError> {
3319        Self::from_inner_with_axis_classes(
3320            self.indices.clone(),
3321            self.try_materialized_inner()?.detach(),
3322            self.storage.axis_classes().to_vec(),
3323        )
3324        .map_err(IdxTensorError::from)
3325    }
3326
3327    /// Check if this tensor is already in canonical form.
3328    pub fn is_simple(&self) -> bool {
3329        true
3330    }
3331
3332    /// Materialize the primal payload as a compact storage snapshot.
3333    ///
3334    /// The eager payload remains authoritative for `f32`/`c32` and tracked
3335    /// structured tensors; this method is a fallible bridge to compact storage.
3336    ///
3337    /// # Errors
3338    ///
3339    /// Returns [`TensorStorageError`] when an eager backend payload cannot be
3340    /// converted to compact storage, when its dtype is `f32`/`c32`, or when a
3341    /// deferred eager operation failed.
3342    pub fn to_storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3343        self.storage.materialize(self.indices.len())
3344    }
3345
3346    /// Materializes and returns a compact storage snapshot.
3347    ///
3348    /// # Errors
3349    /// Returns an error when the compact storage cannot be materialized (a
3350    /// backend failure).
3351    /// # Examples
3352    ///
3353    /// ```
3354    /// use tensor4all_core::{DynIndex, IdxTensor};
3355    /// use tensor4all_tensorbackend::StorageKind;
3356    ///
3357    /// let tensor = IdxTensor::from_dense(
3358    ///     vec![DynIndex::new_dyn(2)],
3359    ///     vec![1.0_f64, 2.0],
3360    /// )
3361    /// .unwrap();
3362    /// assert_eq!(tensor.storage().unwrap().storage_kind(), StorageKind::Dense);
3363    /// ```
3364    pub fn storage(&self) -> std::result::Result<Arc<Storage>, TensorStorageError> {
3365        self.storage.materialize(self.indices.len())
3366    }
3367
3368    /// Return the logical storage layout without materializing compact storage.
3369    ///
3370    /// For `f32` and `c32`, the eager representation is authoritative because
3371    /// compact [`Storage`] supports only `f64` and `c64`.
3372    ///
3373    /// # Examples
3374    ///
3375    /// ```
3376    /// use tensor4all_core::{DynIndex, IdxTensor};
3377    /// use tensor4all_tensorbackend::StorageKind;
3378    ///
3379    /// let tensor = IdxTensor::from_diag(
3380    ///     vec![DynIndex::new_dyn(2), DynIndex::new_dyn(2)],
3381    ///     vec![1.0_f32, 2.0],
3382    /// )
3383    /// .unwrap();
3384    /// assert_eq!(tensor.storage_kind(), StorageKind::Diagonal);
3385    /// ```
3386    pub fn storage_kind(&self) -> StorageKind {
3387        self.storage.storage_kind()
3388    }
3389
3390    /// Return the exact eager scalar dtype without reading tensor values.
3391    ///
3392    /// This feature-gated accessor is used by the CUDA TreeTN boundary to
3393    /// reject mixed dtypes before the first device contraction.
3394    ///
3395    /// # Errors
3396    ///
3397    /// Returns [`IdxTensorError`] when the dtype cannot be determined from the
3398    /// tensor metadata.
3399    ///
3400    /// # Examples
3401    ///
3402    /// ```
3403    /// # #[cfg(feature = "tenferro-cuda")]
3404    /// # {
3405    /// use tensor4all_core::{DynIndex, IdxTensor, IdxTensorError};
3406    ///
3407    /// let tensor = IdxTensor::from_dense(vec![DynIndex::new_dyn(1)], vec![1.0_f32]).unwrap();
3408    /// let dtype = tensor.cuda_dtype().unwrap();
3409    /// assert_eq!(dtype, tenferro::DType::F32);
3410    /// let accessor: fn(&IdxTensor) -> Result<tenferro::DType, IdxTensorError> =
3411    ///     IdxTensor::cuda_dtype;
3412    /// assert_eq!(
3413    ///     std::mem::size_of_val(&accessor),
3414    ///     std::mem::size_of::<fn(&IdxTensor) -> Result<tenferro::DType, IdxTensorError>>(),
3415    /// );
3416    /// # }
3417    /// ```
3418    #[cfg(feature = "tenferro-cuda")]
3419    pub fn cuda_dtype(&self) -> std::result::Result<DType, IdxTensorError> {
3420        self.scalar_dtype().map_err(IdxTensorError::from)
3421    }
3422
3423    /// Sum all elements, returning `AnyScalar`.
3424    ///
3425    /// # Errors
3426    /// Returns an error when the reduction fails (a backend or scalar-extraction
3427    /// failure).
3428    /// # Examples
3429    ///
3430    /// ```
3431    /// use tensor4all_core::{DynIndex, IdxTensor};
3432    ///
3433    /// let i = DynIndex::new_dyn(3);
3434    /// let t = IdxTensor::from_dense(vec![i], vec![1.0, 2.0, 3.0]).unwrap();
3435    /// let s = t.sum().unwrap();
3436    /// assert!((s.real() - 6.0).abs() < 1e-12);
3437    /// ```
3438    pub fn sum(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3439        self.ensure_storage_ready()?;
3440        if self.indices.is_empty() {
3441            return AnyScalar::from_tensor(self.clone()).map_err(IdxTensorError::from);
3442        }
3443        if let Some(payload) = self.storage.eager().filter(|payload| payload.tracks_grad()) {
3444            let axes: Vec<usize> = (0..payload.shape().len()).collect();
3445            let reduced = payload.reduce_sum(Some(&axes))?;
3446            return AnyScalar::from_tensor(Self::from_inner(Vec::new(), reduced)?)
3447                .map_err(IdxTensorError::from);
3448        }
3449        self.storage.sum_scalar().map_err(IdxTensorError::from)
3450    }
3451
3452    /// Extract the scalar value from a 0-dimensional tensor (or 1-element tensor).
3453    ///
3454    /// This is similar to Julia's `only()` function.
3455    ///
3456    /// # Errors
3457    /// Returns an error when the tensor is not rank-0 and does not contain exactly
3458    /// one element (a rank mismatch).
3459    /// # Panics
3460    ///
3461    /// Panics if the tensor has more than one element.
3462    ///
3463    /// # Example
3464    ///
3465    /// ```
3466    /// use tensor4all_core::{IdxTensor, AnyScalar};
3467    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
3468    ///
3469    /// // Create a scalar tensor (0 dimensions, 1 element)
3470    /// let indices: Vec<Index<DynId>> = vec![];
3471    /// let tensor: IdxTensor = IdxTensor::from_dense(indices, vec![42.0]).unwrap();
3472    ///
3473    /// assert_eq!(tensor.only().unwrap().real(), 42.0);
3474    /// ```
3475    pub fn only(&self) -> std::result::Result<AnyScalar, IdxTensorError> {
3476        let dims = self.dims();
3477        let total_size = checked_product(&dims)?;
3478        if !(total_size == 1 || dims.is_empty()) {
3479            return Err(anyhow::anyhow!(
3480                "only() requires a scalar tensor (1 element), got {} elements with dims {:?}",
3481                if dims.is_empty() { 1 } else { total_size },
3482                dims
3483            )
3484            .into());
3485        };
3486        self.sum()
3487    }
3488
3489    /// Permute the tensor dimensions using the given new indices order.
3490    ///
3491    /// This is the main permutation method that takes the desired new indices
3492    /// and automatically computes the corresponding permutation of dimensions
3493    /// and data. The new indices must be a permutation of the original indices
3494    /// (matched by full index identity).
3495    ///
3496    /// # Arguments
3497    /// * `new_indices` - The desired new indices order. Must be a permutation
3498    ///
3499    ///   of `self.indices` (matched by full index identity).
3500    ///
3501    /// # Errors
3502    /// Returns an error when `new_order` does not contain exactly the tensor's
3503    /// indices (an index-set mismatch or a missing-index failure).
3504    /// # Panics
3505    /// Panics if `new_indices.len() != self.indices.len()`, if any full index
3506    /// identity doesn't match, or if there are duplicate indices.
3507    ///
3508    /// # Example
3509    /// ```
3510    /// use tensor4all_core::IdxTensor;
3511    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
3512    ///
3513    /// // Create a 2×3 tensor
3514    /// let i = Index::new_dyn(2);
3515    /// let j = Index::new_dyn(3);
3516    /// let indices = vec![i.clone(), j.clone()];
3517    /// let tensor: IdxTensor = IdxTensor::from_dense(indices, vec![0.0; 6]).unwrap();
3518    ///
3519    /// // Permute to 3×2: swap the two dimensions by providing new indices order
3520    /// let permuted = tensor.permute_indices(&[j, i]).unwrap();
3521    /// assert_eq!(permuted.dims(), vec![3, 2]);
3522    /// ```
3523    pub fn permute_indices(
3524        &self,
3525        new_indices: &[DynIndex],
3526    ) -> std::result::Result<Self, IdxTensorError> {
3527        // Compute permutation by full index equality
3528        let perm = compute_permutation_from_indices(&self.indices, new_indices)?;
3529        if perm.iter().copied().eq(0..perm.len()) {
3530            return Ok(Self {
3531                indices: new_indices.to_vec(),
3532                storage: self.storage.clone(),
3533                eager_cache: Arc::clone(&self.eager_cache),
3534            });
3535        }
3536
3537        let permuted = self.try_materialized_inner()?.transpose(&perm)?;
3538        let axis_classes = self.permute_axis_classes(&perm);
3539        Self::from_inner_with_axis_classes(new_indices.to_vec(), permuted, axis_classes)
3540            .map_err(IdxTensorError::from)
3541    }
3542
3543    /// Permute the tensor dimensions, returning a new tensor.
3544    ///
3545    /// This method reorders the indices, dimensions, and data according to the
3546    /// given permutation. The permutation specifies which old axis each new
3547    /// axis corresponds to: `new_axis[i] = old_axis[perm[i]]`.
3548    ///
3549    /// # Arguments
3550    /// * `perm` - The permutation: `perm[i]` is the old axis index for new axis `i`
3551    ///
3552    /// # Errors
3553    /// Returns an error when `new_order` does not contain exactly the tensor's
3554    /// indices (an index-set mismatch or a missing-index failure).
3555    /// # Panics
3556    /// Panics if `perm.len() != self.indices.len()` or if the permutation is invalid.
3557    ///
3558    /// # Example
3559    /// ```
3560    /// use tensor4all_core::IdxTensor;
3561    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
3562    ///
3563    /// // Create a 2×3 tensor
3564    /// let indices = vec![
3565    ///     Index::new_dyn(2),
3566    ///     Index::new_dyn(3),
3567    /// ];
3568    /// let tensor: IdxTensor = IdxTensor::from_dense(indices, vec![0.0; 6]).unwrap();
3569    ///
3570    /// // Permute to 3×2: swap the two dimensions
3571    /// let permuted = tensor.permute(&[1, 0]).unwrap();
3572    /// assert_eq!(permuted.dims(), vec![3, 2]);
3573    /// ```
3574    pub fn permute(&self, perm: &[usize]) -> std::result::Result<Self, IdxTensorError> {
3575        if !(perm.len() == self.indices.len()) {
3576            return Err(anyhow::anyhow!("permutation length must match tensor rank").into());
3577        };
3578        let mut seen = HashSet::new();
3579        for &axis in perm {
3580            if !(axis < self.indices.len()) {
3581                return Err(anyhow::anyhow!("permutation axis {axis} out of range").into());
3582            };
3583            if !(seen.insert(axis)) {
3584                return Err(anyhow::anyhow!("duplicate axis {axis} in permutation").into());
3585            };
3586        }
3587        if perm.iter().copied().eq(0..perm.len()) {
3588            return Ok(self.clone());
3589        }
3590
3591        // Permute indices
3592        let new_indices: Vec<DynIndex> = perm.iter().map(|&i| self.indices[i].clone()).collect();
3593        let permuted = self.try_materialized_inner()?.transpose(perm)?;
3594        let axis_classes = self.permute_axis_classes(perm);
3595        Self::from_inner_with_axis_classes(new_indices, permuted, axis_classes)
3596            .map_err(IdxTensorError::from)
3597    }
3598
3599    pub(crate) fn try_contract_pairwise_default(&self, other: &Self) -> Result<Self> {
3600        self.try_contract_pairwise_default_with_options(other, PairwiseContractionOptions::new())
3601    }
3602
3603    pub(crate) fn try_contract_pairwise_default_with_options(
3604        &self,
3605        other: &Self,
3606        options: PairwiseContractionOptions,
3607    ) -> Result<Self> {
3608        // Device-resident structured/diagonal operands cannot enter the native
3609        // host session used below; route them through the owning-runtime eager
3610        // plan executor (which also serves N-ary resident contraction).
3611        // Dense-dense resident pairs keep the pairwise eager path, and host
3612        // operands keep this function unchanged.
3613        #[cfg(feature = "tenferro-cuda")]
3614        if self.is_cuda_resident()
3615            && other.is_cuda_resident()
3616            && self.should_use_structured_payload_contract(other)
3617        {
3618            return super::contract::contract_pair_via_plan(self, other, options);
3619        }
3620        let self_indices = profile_pairwise_contract_section("operand_indices", || {
3621            self.operand_indices_for_contraction(options.lhs_conj)
3622        });
3623        let other_indices = profile_pairwise_contract_section("operand_indices", || {
3624            other.operand_indices_for_contraction(options.rhs_conj)
3625        });
3626        let self_dims = profile_pairwise_contract_section("expected_dims", || {
3627            Self::expected_dims_from_indices(&self_indices)
3628        });
3629        let other_dims = profile_pairwise_contract_section("expected_dims", || {
3630            Self::expected_dims_from_indices(&other_indices)
3631        });
3632        let spec = profile_pairwise_contract_section("prepare_contraction", || {
3633            prepare_contraction(&self_indices, &self_dims, &other_indices, &other_dims)
3634        })
3635        .context("contraction preparation failed")?;
3636        let result_axis_classes = profile_pairwise_contract_section("result_axis_classes", || {
3637            Self::binary_contraction_axis_classes(
3638                self.storage.axis_classes(),
3639                &spec.axes_a,
3640                other.storage.axis_classes(),
3641                &spec.axes_b,
3642            )
3643        })?;
3644
3645        if profile_pairwise_contract_section("structured_check", || {
3646            self.should_use_structured_payload_contract(other)
3647        }) {
3648            if options.has_conj() {
3649                let lhs = if options.lhs_conj {
3650                    self.conj()
3651                } else {
3652                    self.clone()
3653                };
3654                let rhs = if options.rhs_conj {
3655                    other.conj()
3656                } else {
3657                    other.clone()
3658                };
3659                return profile_pairwise_contract_section("structured_conj_fallback", || {
3660                    lhs.try_contract_pairwise_default(&rhs)
3661                });
3662            }
3663            return profile_pairwise_contract_section("structured_payload_contract", || {
3664                self.contract_structured_payloads(
3665                    other,
3666                    spec.result_indices.into_vec(),
3667                    &spec.axes_a,
3668                    &spec.axes_b,
3669                )
3670            });
3671        }
3672
3673        if self.indices.is_empty() && other.indices.is_empty() {
3674            if options.has_conj() {
3675                let lhs = if options.lhs_conj {
3676                    self.conj()
3677                } else {
3678                    self.clone()
3679                };
3680                let rhs = if options.rhs_conj {
3681                    other.conj()
3682                } else {
3683                    other.clone()
3684                };
3685                return lhs.try_contract_pairwise_default(&rhs);
3686            }
3687            let result = profile_pairwise_contract_section("scalar_mul", || {
3688                Ok::<_, anyhow::Error>(
3689                    self.try_materialized_inner()?
3690                        .mul(other.try_materialized_inner()?)?,
3691                )
3692            })?;
3693            return profile_pairwise_contract_section("from_inner", || {
3694                Self::from_inner(spec.result_indices.into_vec(), result)
3695            });
3696        }
3697
3698        let self_dtype = self.try_materialized_inner()?.dtype();
3699        let other_dtype = other.try_materialized_inner()?.dtype();
3700        if self_dtype != other_dtype {
3701            if options.has_conj() {
3702                let lhs = if options.lhs_conj {
3703                    self.conj()
3704                } else {
3705                    self.clone()
3706                };
3707                let rhs = if options.rhs_conj {
3708                    other.conj()
3709                } else {
3710                    other.clone()
3711                };
3712                return lhs.try_contract_pairwise_default(&rhs);
3713            }
3714            let self_native = self.try_materialized_inner()?.duplicate_value()?;
3715            let other_native = other.try_materialized_inner()?.duplicate_value()?;
3716            let result_native = profile_pairwise_contract_section("native_contract", || {
3717                contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)
3718            })?;
3719            return profile_pairwise_contract_section("from_native", || {
3720                Self::from_untracked_native_with_axis_classes(
3721                    spec.result_indices.into_vec(),
3722                    result_native,
3723                    result_axis_classes,
3724                )
3725            });
3726        }
3727
3728        let config = profile_pairwise_contract_section("build_dot_general_config", || {
3729            Self::binary_dot_general_config(&spec.axes_a, &spec.axes_b)
3730        })?;
3731        let result = profile_pairwise_contract_section("dot_general_with_conj", || {
3732            let lhs = profile_pairwise_contract_section("lhs_try_materialized_inner", || {
3733                self.try_materialized_inner()
3734            })?;
3735            let rhs = profile_pairwise_contract_section("rhs_try_materialized_inner", || {
3736                other.try_materialized_inner()
3737            })?;
3738            profile_pairwise_contract_section("dot_general_execute", || {
3739                lhs.dot_general_with_conj(rhs, config, options.lhs_conj, options.rhs_conj)
3740            })
3741            .map_err(anyhow::Error::from)
3742        })?;
3743        record_pairwise_contract_profile_bytes(
3744            "dot_general_output",
3745            tensor_profile_bytes(result.dtype(), result.shape()),
3746        );
3747        profile_pairwise_contract_section("from_inner_axis_classes", || {
3748            Self::from_inner_with_axis_classes(
3749                spec.result_indices.into_vec(),
3750                result,
3751                result_axis_classes,
3752            )
3753        })
3754    }
3755
3756    // Provenance: this retained-axis dense path is a tensor4all-specific
3757    // implementation of the batch needed by SRC. Its output ordering follows
3758    // `defaults::contract::build_contraction_plan`; it is not present in
3759    // RandomMPOMPS and is labelled `[AI-Supplied]` in the SRC audit.
3760    pub(crate) fn try_contract_pairwise_retaining(
3761        &self,
3762        other: &Self,
3763        retained_indices: &[DynIndex],
3764    ) -> Result<Self> {
3765        if retained_indices.is_empty()
3766            || self.should_use_structured_payload_contract(other)
3767            || self.try_materialized_inner()?.dtype() != other.try_materialized_inner()?.dtype()
3768        {
3769            let options = ContractionOptions::new().with_retain_indices(retained_indices);
3770            return super::contract::contract_with_options(&[self, other], options)
3771                .map_err(anyhow::Error::from);
3772        }
3773
3774        let common = common_ind_positions(&self.indices, &other.indices);
3775        if common.is_empty() {
3776            let options = ContractionOptions::new().with_retain_indices(retained_indices);
3777            return super::contract::contract_with_options(&[self, other], options)
3778                .map_err(anyhow::Error::from);
3779        }
3780
3781        let mut retained_pairs = Vec::with_capacity(retained_indices.len());
3782        for retained in retained_indices {
3783            let Some(pos_a) = self.indices.iter().position(|index| index == retained) else {
3784                let options = ContractionOptions::new().with_retain_indices(retained_indices);
3785                return super::contract::contract_with_options(&[self, other], options)
3786                    .map_err(anyhow::Error::from);
3787            };
3788            let Some(pos_b) = other.indices.iter().position(|index| index == retained) else {
3789                let options = ContractionOptions::new().with_retain_indices(retained_indices);
3790                return super::contract::contract_with_options(&[self, other], options)
3791                    .map_err(anyhow::Error::from);
3792            };
3793            if !self.indices[pos_a].is_contractable(&other.indices[pos_b]) {
3794                let options = ContractionOptions::new().with_retain_indices(retained_indices);
3795                return super::contract::contract_with_options(&[self, other], options)
3796                    .map_err(anyhow::Error::from);
3797            }
3798            retained_pairs.push((pos_a, pos_b));
3799        }
3800
3801        let retained_pairs_set: HashSet<(usize, usize)> = retained_pairs.iter().copied().collect();
3802        let mut contracting_a = Vec::with_capacity(common.len());
3803        let mut contracting_b = Vec::with_capacity(common.len());
3804        for &(pos_a, pos_b) in &common {
3805            if !retained_pairs_set.contains(&(pos_a, pos_b)) {
3806                contracting_a.push(pos_a);
3807                contracting_b.push(pos_b);
3808            }
3809        }
3810
3811        let config = DotGeneralConfig {
3812            lhs_contracting_dims: contracting_a.clone(),
3813            rhs_contracting_dims: contracting_b.clone(),
3814            lhs_batch_dims: retained_pairs.iter().map(|&(pos_a, _)| pos_a).collect(),
3815            rhs_batch_dims: retained_pairs.iter().map(|&(_, pos_b)| pos_b).collect(),
3816        };
3817        let result = self
3818            .try_materialized_inner()?
3819            .dot_general_with_conj(other.try_materialized_inner()?, config, false, false)
3820            .map_err(|error| anyhow::anyhow!("retained pairwise contraction failed: {error}"))?;
3821
3822        // dot_general emits [lhs_free, rhs_free, batch]. Restore the index
3823        // order used by contract_with_options, where retained axes occur at
3824        // their first operand position. The permutation is a single dense
3825        // copy after the batched GEMM and is much cheaper than one einsum per
3826        // probe column.
3827        let contracting_a_set: HashSet<usize> = contracting_a.into_iter().collect();
3828        let contracting_b_set: HashSet<usize> = contracting_b.into_iter().collect();
3829        let retained_a_set: HashSet<usize> = retained_pairs.iter().map(|&(a, _)| a).collect();
3830        let retained_b_set: HashSet<usize> = retained_pairs.iter().map(|&(_, b)| b).collect();
3831        let mut current_indices = self
3832            .indices
3833            .iter()
3834            .enumerate()
3835            .filter(|(axis, _)| !retained_a_set.contains(axis) && !contracting_a_set.contains(axis))
3836            .map(|(_, index)| index.clone())
3837            .chain(
3838                other
3839                    .indices
3840                    .iter()
3841                    .enumerate()
3842                    .filter(|(axis, _)| {
3843                        !retained_b_set.contains(axis) && !contracting_b_set.contains(axis)
3844                    })
3845                    .map(|(_, index)| index.clone()),
3846            )
3847            .collect::<Vec<_>>();
3848        current_indices.extend(
3849            retained_pairs
3850                .iter()
3851                .map(|&(pos_a, _)| self.indices[pos_a].clone()),
3852        );
3853
3854        let mut desired_indices = self
3855            .indices
3856            .iter()
3857            .enumerate()
3858            .filter(|(axis, _)| !contracting_a_set.contains(axis))
3859            .map(|(_, index)| index.clone())
3860            .chain(
3861                other
3862                    .indices
3863                    .iter()
3864                    .enumerate()
3865                    .filter(|(axis, _)| {
3866                        !contracting_b_set.contains(axis) && !retained_b_set.contains(axis)
3867                    })
3868                    .map(|(_, index)| index.clone()),
3869            )
3870            .collect::<Vec<_>>();
3871        if desired_indices.is_empty() {
3872            desired_indices = current_indices.clone();
3873        }
3874        let result = if current_indices == desired_indices {
3875            result
3876        } else {
3877            let permutation = desired_indices
3878                .iter()
3879                .map(|desired| {
3880                    current_indices
3881                        .iter()
3882                        .position(|current| current == desired)
3883                        .ok_or_else(|| {
3884                            anyhow::anyhow!(
3885                                "retained pairwise result is missing index {:?}",
3886                                desired
3887                            )
3888                        })
3889                })
3890                .collect::<Result<Vec<_>>>()?;
3891            result.transpose(&permutation)?
3892        };
3893        Self::from_inner(desired_indices, result)
3894    }
3895
3896    pub(crate) fn try_tensordot_pairwise_explicit(
3897        &self,
3898        other: &Self,
3899        pairs: &[(DynIndex, DynIndex)],
3900    ) -> Result<Self> {
3901        use crate::index_ops::ContractionError;
3902
3903        let self_dims = Self::expected_dims_from_indices(&self.indices);
3904        let other_dims = Self::expected_dims_from_indices(&other.indices);
3905        let spec = prepare_contraction_pairs(
3906            &self.indices,
3907            &self_dims,
3908            &other.indices,
3909            &other_dims,
3910            pairs,
3911        )
3912        .map_err(|e| match e {
3913            ContractionError::NoCommonIndices => {
3914                anyhow::anyhow!("tensordot: No pairs specified for contraction")
3915            }
3916            ContractionError::BatchContractionNotImplemented => anyhow::anyhow!(
3917                "tensordot: Common index found but not in contraction pairs. \
3918                         Batch contraction is not yet implemented."
3919            ),
3920            ContractionError::IndexNotFound { tensor } => {
3921                anyhow::anyhow!("tensordot: Index not found in {} tensor", tensor)
3922            }
3923            ContractionError::DimensionMismatch {
3924                pos_a,
3925                pos_b,
3926                dim_a,
3927                dim_b,
3928            } => anyhow::anyhow!(
3929                "tensordot: Dimension mismatch: self[{}]={} != other[{}]={}",
3930                pos_a,
3931                dim_a,
3932                pos_b,
3933                dim_b
3934            ),
3935            ContractionError::DuplicateAxis { tensor, pos } => {
3936                anyhow::anyhow!("tensordot: Duplicate axis {} in {} tensor", pos, tensor)
3937            }
3938        })?;
3939        let result_axis_classes = Self::binary_contraction_axis_classes(
3940            self.storage.axis_classes(),
3941            &spec.axes_a,
3942            other.storage.axis_classes(),
3943            &spec.axes_b,
3944        )?;
3945
3946        if self.should_use_structured_payload_contract(other) {
3947            return self.contract_structured_payloads(
3948                other,
3949                spec.result_indices.into_vec(),
3950                &spec.axes_a,
3951                &spec.axes_b,
3952            );
3953        }
3954
3955        if self.indices.is_empty() && other.indices.is_empty() {
3956            let result = self
3957                .try_materialized_inner()?
3958                .mul(other.try_materialized_inner()?)
3959                .map_err(|e| anyhow::anyhow!("tensordot scalar multiply failed: {e}"))?;
3960            return Self::from_inner(spec.result_indices.into_vec(), result);
3961        }
3962
3963        let self_dtype = self.try_materialized_inner()?.dtype();
3964        let other_dtype = other.try_materialized_inner()?.dtype();
3965        if self_dtype != other_dtype {
3966            let self_native = self.try_materialized_inner()?.duplicate_value()?;
3967            let other_native = other.try_materialized_inner()?.duplicate_value()?;
3968            let result_native =
3969                contract_native_tensor(&self_native, &spec.axes_a, &other_native, &spec.axes_b)?;
3970            return Self::from_untracked_native_with_axis_classes(
3971                spec.result_indices.into_vec(),
3972                result_native,
3973                result_axis_classes,
3974            );
3975        }
3976
3977        let subscripts = Self::build_binary_einsum_subscripts(
3978            self.indices.len(),
3979            &spec.axes_a,
3980            other.indices.len(),
3981            &spec.axes_b,
3982        )?;
3983        let result = [
3984            self.try_materialized_inner()?,
3985            other.try_materialized_inner()?,
3986        ]
3987        .einsum_subscripts(&subscripts)
3988        .map_err(|e| anyhow::anyhow!("tensordot failed: {e}"))?;
3989        Self::from_inner_with_axis_classes(
3990            spec.result_indices.into_vec(),
3991            result,
3992            result_axis_classes,
3993        )
3994    }
3995
3996    pub(crate) fn try_outer_product_pairwise(&self, other: &Self) -> Result<Self> {
3997        use anyhow::Context;
3998
3999        // Check for common indices - outer product should have none
4000        let common_positions = common_ind_positions(&self.indices, &other.indices);
4001        if !common_positions.is_empty() {
4002            let common_ids: Vec<_> = common_positions
4003                .iter()
4004                .map(|(pos_a, _)| self.indices[*pos_a].id())
4005                .collect();
4006            return Err(anyhow::anyhow!(
4007                "outer_product: tensors have common indices {:?}. \
4008                 Use tensordot to contract common indices, or use sim() to replace \
4009                 indices with fresh IDs before computing outer product.",
4010                common_ids
4011            ))
4012            .context("outer_product: common indices found");
4013        }
4014
4015        // Build result indices and dimensions
4016        let mut result_indices = self.indices.clone();
4017        result_indices.extend(other.indices.iter().cloned());
4018        let result_axis_classes = Self::binary_contraction_axis_classes(
4019            self.storage.axis_classes(),
4020            &[],
4021            other.storage.axis_classes(),
4022            &[],
4023        )?;
4024        if self.should_use_structured_payload_contract(other) {
4025            return self.contract_structured_payloads(other, result_indices, &[], &[]);
4026        }
4027        let self_dtype = self.try_materialized_inner()?.dtype();
4028        let other_dtype = other.try_materialized_inner()?.dtype();
4029        if self_dtype != other_dtype {
4030            let self_native = self.try_materialized_inner()?.duplicate_value()?;
4031            let other_native = other.try_materialized_inner()?.duplicate_value()?;
4032            let result_native = contract_native_tensor(&self_native, &[], &other_native, &[])?;
4033            return Self::from_untracked_native_with_axis_classes(
4034                result_indices,
4035                result_native,
4036                result_axis_classes,
4037            );
4038        }
4039
4040        let subscripts = Self::build_binary_einsum_subscripts(
4041            self.indices.len(),
4042            &[],
4043            other.indices.len(),
4044            &[],
4045        )?;
4046        let result = [
4047            self.try_materialized_inner()?,
4048            other.try_materialized_inner()?,
4049        ]
4050        .einsum_subscripts(&subscripts)
4051        .map_err(|e| anyhow::anyhow!("outer_product failed: {e}"))?;
4052        Self::from_inner_with_axis_classes(result_indices, result, result_axis_classes)
4053    }
4054}
4055
4056// ============================================================================
4057// Random tensor generation
4058// ============================================================================
4059
4060impl IdxTensor {
4061    /// Create a random tensor with values from standard normal distribution (generic over scalar type).
4062    ///
4063    /// For `f64`, each element is drawn from the standard normal distribution.
4064    /// For `Complex64`, both real and imaginary parts are drawn independently.
4065    ///
4066    /// # Type Parameters
4067    /// * `T` - The scalar element type (must implement [`RandomScalar`])
4068    /// * `R` - The random number generator type
4069    ///
4070    /// # Arguments
4071    /// * `rng` - Random number generator
4072    /// * `indices` - The indices for the tensor
4073    ///
4074    /// # Errors
4075    /// Returns an error when the dimension product overflows (an overflow failure)
4076    /// or the backend cannot generate the requested scalar type.
4077    /// # Example
4078    /// ```
4079    /// use tensor4all_core::IdxTensor;
4080    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4081    /// use rand::SeedableRng;
4082    /// use rand_chacha::ChaCha8Rng;
4083    ///
4084    /// let mut rng = ChaCha8Rng::seed_from_u64(42);
4085    /// let i = Index::new_dyn(2);
4086    /// let j = Index::new_dyn(3);
4087    /// let tensor: IdxTensor = IdxTensor::random::<f64, _>(&mut rng, vec![i, j]).unwrap();
4088    /// assert_eq!(tensor.dims(), vec![2, 3]);
4089    /// ```
4090    pub fn random<T: RandomScalar, R: Rng>(
4091        rng: &mut R,
4092        indices: Vec<DynIndex>,
4093    ) -> std::result::Result<Self, IdxTensorError> {
4094        let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
4095        let size = checked_product(&dims)?;
4096        let data: Vec<T> = (0..size).map(|_| T::random_value(rng)).collect();
4097        Self::from_dense(indices, data)
4098    }
4099}
4100
4101impl IdxTensor {
4102    /// Add two tensors element-wise.
4103    ///
4104    /// The tensors must have the same full index set (including tags and prime
4105    /// levels). If the indices are in a different order, the other tensor will
4106    /// be permuted to match `self`.
4107    ///
4108    /// # Arguments
4109    /// * `other` - The tensor to add
4110    ///
4111    /// # Returns
4112    /// A new tensor representing `self + other`, or an error if:
4113    /// - The tensors have different index sets
4114    /// - The dimensions don't match
4115    /// - Storage types are incompatible
4116    ///
4117    /// # Errors
4118    /// Returns an error when the two tensors have different index sets (an
4119    /// index-set mismatch) or the arithmetic reports a failure.
4120    /// # Example
4121    /// ```
4122    /// use tensor4all_core::IdxTensor;
4123    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4124    ///
4125    /// let i = Index::new_dyn(2);
4126    /// let j = Index::new_dyn(3);
4127    ///
4128    /// let indices_a = vec![i.clone(), j.clone()];
4129    /// let data_a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
4130    /// let tensor_a: IdxTensor = IdxTensor::from_dense(indices_a, data_a).unwrap();
4131    ///
4132    /// let indices_b = vec![i.clone(), j.clone()];
4133    /// let data_b = vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
4134    /// let tensor_b: IdxTensor = IdxTensor::from_dense(indices_b, data_b).unwrap();
4135    ///
4136    /// let sum = tensor_a.add(&tensor_b).unwrap();
4137    /// // sum = [[2, 3, 4], [5, 6, 7]]
4138    /// ```
4139    pub fn add(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
4140        // Validate that both tensors have the same number of indices
4141        if self.indices.len() != other.indices.len() {
4142            return Err(anyhow::anyhow!(
4143                "Index count mismatch: self has {} indices, other has {}",
4144                self.indices.len(),
4145                other.indices.len()
4146            )
4147            .into());
4148        }
4149
4150        // Validate that both tensors have the same set of indices
4151        let self_set: HashSet<_> = self.indices.iter().collect();
4152        let other_set: HashSet<_> = other.indices.iter().collect();
4153
4154        if self_set != other_set {
4155            return Err(
4156                anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
4157            );
4158        }
4159
4160        // Permute other to match self's index order (no-op if already aligned)
4161        let other_aligned = other.permute_indices(&self.indices)?;
4162
4163        // Validate dimensions match after alignment
4164        let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
4165        let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
4166        if self_expected_dims != other_expected_dims {
4167            use crate::TagSetLike;
4168            let fmt = |indices: &[DynIndex]| -> Vec<String> {
4169                indices
4170                    .iter()
4171                    .map(|idx| {
4172                        let tags: Vec<String> = idx.tags().iter().collect();
4173                        format!("{:?}(dim={},tags={:?})", idx.id(), idx.dim(), tags)
4174                    })
4175                    .collect()
4176            };
4177            return Err(anyhow::anyhow!(
4178                "Dimension mismatch after alignment.\n\
4179                 self: dims={:?}, indices(order)={:?}\n\
4180                 other_aligned: dims={:?}, indices(order)={:?}",
4181                self_expected_dims,
4182                fmt(&self.indices),
4183                other_expected_dims,
4184                fmt(&other_aligned.indices)
4185            )
4186            .into());
4187        }
4188
4189        self.axpby(
4190            AnyScalar::new_real(1.0),
4191            &other_aligned,
4192            AnyScalar::new_real(1.0),
4193        )
4194    }
4195
4196    /// Compute a linear combination: `a * self + b * other`.
4197    ///
4198    /// Both tensors must have the same full set of indices (including tags and
4199    /// prime levels). If indices are in a different order, `other` is automatically permuted
4200    /// to match `self`.
4201    ///
4202    /// # Errors
4203    /// Returns an error when the tensors have different index sets (an index-set
4204    /// mismatch) or the arithmetic reports a failure.
4205    /// # Examples
4206    ///
4207    /// ```
4208    /// use tensor4all_core::{AnyScalar, DynIndex, IdxTensor};
4209    ///
4210    /// let i = DynIndex::new_dyn(2);
4211    /// let a = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
4212    /// let b = IdxTensor::from_dense(vec![i.clone()], vec![3.0, 4.0]).unwrap();
4213    ///
4214    /// // 2*a + 3*b = [2+9, 4+12] = [11, 16]
4215    /// let result = a.axpby(AnyScalar::new_real(2.0), &b, AnyScalar::new_real(3.0)).unwrap();
4216    /// let data = result.to_vec::<f64>().unwrap();
4217    /// assert!((data[0] - 11.0).abs() < 1e-12);
4218    /// assert!((data[1] - 16.0).abs() < 1e-12);
4219    /// ```
4220    pub fn axpby(
4221        &self,
4222        a: AnyScalar,
4223        other: &Self,
4224        b: AnyScalar,
4225    ) -> std::result::Result<Self, IdxTensorError> {
4226        // Validate that both tensors have the same number of indices.
4227        if self.indices.len() != other.indices.len() {
4228            return Err(anyhow::anyhow!(
4229                "Index count mismatch: self has {} indices, other has {}",
4230                self.indices.len(),
4231                other.indices.len()
4232            )
4233            .into());
4234        }
4235
4236        // Validate that both tensors have the same set of indices.
4237        let self_set: HashSet<_> = self.indices.iter().collect();
4238        let other_set: HashSet<_> = other.indices.iter().collect();
4239        if self_set != other_set {
4240            return Err(
4241                anyhow::anyhow!("Index set mismatch: tensors must have the same indices").into(),
4242            );
4243        }
4244
4245        // Align other tensor axis order to self.
4246        let other_aligned = other.permute_indices(&self.indices)?;
4247
4248        // Validate dimensions match after alignment.
4249        let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
4250        let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
4251        if self_expected_dims != other_expected_dims {
4252            return Err(anyhow::anyhow!(
4253                "Dimension mismatch after alignment: self={:?}, other_aligned={:?}",
4254                self_expected_dims,
4255                other_expected_dims
4256            )
4257            .into());
4258        }
4259
4260        let axis_classes = if self.storage.axis_classes() == other_aligned.storage.axis_classes() {
4261            self.storage.axis_classes().to_vec()
4262        } else {
4263            Self::dense_axis_classes(self.indices.len())
4264        };
4265
4266        let same_compact_layout = self.storage.payload_dims()
4267            == other_aligned.storage.payload_dims()
4268            && self.storage.payload_strides_vec() == other_aligned.storage.payload_strides_vec()
4269            && self.storage.axis_classes() == other_aligned.storage.axis_classes();
4270        if same_compact_layout
4271            && matches!(&self.storage, IdxTensorStorage::Materialized(_))
4272            && matches!(&other_aligned.storage, IdxTensorStorage::Materialized(_))
4273            && !self.tracks_grad()
4274            && !other_aligned.tracks_grad()
4275            && !a.tracks_grad()
4276            && !b.tracks_grad()
4277        {
4278            let lhs_storage = self.storage.materialize(self.indices.len())?;
4279            let rhs_storage = other_aligned
4280                .storage
4281                .materialize(other_aligned.indices.len())?;
4282            let combined = lhs_storage
4283                .axpby(
4284                    &a.to_backend_scalar(),
4285                    rhs_storage.as_ref(),
4286                    &b.to_backend_scalar(),
4287                )
4288                .map_err(|e| anyhow::anyhow!("storage axpby failed: {e}"))?;
4289            return Self::from_storage(self.indices.clone(), Arc::new(combined));
4290        }
4291
4292        let lhs = self.scale(a)?;
4293        let rhs = other_aligned.scale(b)?;
4294        let combined = lhs
4295            .try_materialized_inner()?
4296            .add(rhs.try_materialized_inner()?)
4297            .map_err(|e| anyhow::anyhow!("tensor addition failed: {e}"))?;
4298        Self::from_inner_with_axis_classes(self.indices.clone(), combined, axis_classes)
4299            .map_err(IdxTensorError::from)
4300    }
4301
4302    /// Scalar multiplication.
4303    ///
4304    /// Multiplies every element by `scalar`.
4305    ///
4306    /// # Errors
4307    /// Returns an error when the scalar coefficient is invalid for the tensor's
4308    /// scalar type (an invalid scalar dtype) or the backend reports a
4309    /// failure.
4310    /// # Examples
4311    ///
4312    /// ```
4313    /// use tensor4all_core::{AnyScalar, DynIndex, IdxTensor};
4314    ///
4315    /// let i = DynIndex::new_dyn(3);
4316    /// let t = IdxTensor::from_dense(vec![i], vec![1.0, 2.0, 3.0]).unwrap();
4317    /// let scaled = t.scale(AnyScalar::new_real(2.0)).unwrap();
4318    /// assert_eq!(scaled.to_vec::<f64>().unwrap(), vec![2.0, 4.0, 6.0]);
4319    /// ```
4320    pub fn scale(&self, scalar: AnyScalar) -> std::result::Result<Self, IdxTensorError> {
4321        if matches!(
4322            &self.storage,
4323            IdxTensorStorage::Eager { .. }
4324                | IdxTensorStorage::Compact(_)
4325                | IdxTensorStorage::Materialized(_)
4326        ) {
4327            // Scale via the compact payload only. Materialized structured
4328            // storage is converted payload-coordinate by payload-coordinate
4329            // (never the logical domain) and returned as compact storage, so
4330            // scaling never touches unreferenced strided-gap backing entries.
4331            let storage = self.storage.scale_eager_payload(&scalar)?;
4332            return Ok(Self {
4333                indices: self.indices.clone(),
4334                storage,
4335                eager_cache: Self::empty_eager_cache(),
4336            });
4337        }
4338
4339        let self_dtype = self.try_materialized_inner()?.dtype();
4340        let scalar_dtype = scalar.as_tensor()?.try_materialized_inner()?.dtype();
4341        if self_dtype != scalar_dtype {
4342            let target_dtype = Self::scale_target_dtype(self_dtype, scalar_dtype)?;
4343            let self_inner = self.try_materialized_inner()?;
4344            let self_inner = if self_inner.dtype() == target_dtype {
4345                self_inner.clone()
4346            } else {
4347                self_inner.cast(target_dtype)?
4348            };
4349            let scalar_inner = scalar.as_tensor()?.try_materialized_inner()?;
4350            let scalar_inner = if scalar_inner.dtype() == target_dtype {
4351                scalar_inner.clone()
4352            } else {
4353                scalar_inner.cast(target_dtype)?
4354            };
4355            let scaled = if self.indices.is_empty() {
4356                self_inner
4357                    .mul(&scalar_inner)
4358                    .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4359            } else {
4360                let subscripts = Self::scale_subscripts(self.indices.len())?;
4361                [&self_inner, &scalar_inner]
4362                    .einsum_subscripts(&subscripts)
4363                    .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4364            };
4365            return Self::from_inner_with_axis_classes(
4366                self.indices.clone(),
4367                scaled,
4368                self.storage.axis_classes().to_vec(),
4369            )
4370            .map_err(IdxTensorError::from);
4371        }
4372        let scaled = if self.indices.is_empty() {
4373            self.try_materialized_inner()?
4374                .mul(scalar.as_tensor()?.try_materialized_inner()?)
4375                .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
4376        } else {
4377            let subscripts = Self::scale_subscripts(self.indices.len())?;
4378            [
4379                self.try_materialized_inner()?,
4380                scalar.as_tensor()?.try_materialized_inner()?,
4381            ]
4382            .einsum_subscripts(&subscripts)
4383            .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
4384        };
4385        Self::from_inner_with_axis_classes(
4386            self.indices.clone(),
4387            scaled,
4388            self.storage.axis_classes().to_vec(),
4389        )
4390        .map_err(IdxTensorError::from)
4391    }
4392
4393    /// Inner product (dot product) of two tensors.
4394    ///
4395    /// Computes `⟨self, other⟩ = Σ conj(self)_i * other_i`.
4396    ///
4397    /// # Errors
4398    /// Returns an error when the tensors have different index sets (an index-set
4399    /// mismatch) or the contraction reports a failure.
4400    /// # Examples
4401    ///
4402    /// ```
4403    /// use tensor4all_core::{DynIndex, IdxTensor};
4404    ///
4405    /// let i = DynIndex::new_dyn(3);
4406    /// let a = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0, 3.0]).unwrap();
4407    /// let b = IdxTensor::from_dense(vec![i.clone()], vec![4.0, 5.0, 6.0]).unwrap();
4408    ///
4409    /// // <a, b> = 1*4 + 2*5 + 3*6 = 32
4410    /// let ip = a.inner_product(&b).unwrap();
4411    /// assert!((ip.real() - 32.0).abs() < 1e-12);
4412    /// ```
4413    pub fn inner_product(&self, other: &Self) -> std::result::Result<AnyScalar, IdxTensorError> {
4414        if self.indices.len() == other.indices.len() {
4415            let self_set: HashSet<_> = self.indices.iter().collect();
4416            let other_set: HashSet<_> = other.indices.iter().collect();
4417            if self_set == other_set {
4418                let other_aligned = other.permute_indices(&self.indices)?;
4419                let result = super::contract::contract_pair_with_operand_options(
4420                    self,
4421                    &other_aligned,
4422                    PairwiseContractionOptions::new().with_lhs_conj(true),
4423                )?;
4424                return result.sum();
4425            }
4426        }
4427
4428        // Contract self.conj() with other over all indices
4429        let result = super::contract::contract_pair_with_operand_options(
4430            self,
4431            other,
4432            PairwiseContractionOptions::new().with_lhs_conj(true),
4433        )?;
4434        // Result should be a scalar (no indices)
4435        result.sum()
4436    }
4437}
4438
4439// ============================================================================
4440// Index Replacement Methods
4441// ============================================================================
4442
4443impl IdxTensor {
4444    /// Replace an index in the tensor with a new index.
4445    ///
4446    /// This replaces every index equal to `old_index` (full index equality,
4447    /// including id, prime level, and tags) with `new_index`.
4448    /// The storage data is not modified, only the index metadata is changed.
4449    ///
4450    /// # Arguments
4451    /// * `old_index` - The index to replace (matched by full index equality)
4452    /// * `new_index` - The new index to use
4453    ///
4454    /// # Returns
4455    /// A new tensor with the index replaced. If no index matches `old_index`,
4456    /// returns a clone of the original tensor.
4457    ///
4458    /// # Errors
4459    /// Returns an error when the new index has an incompatible dimension
4460    /// (a shape mismatch: the replacement dimension must equal the original).
4461    /// # Example
4462    /// ```
4463    /// use tensor4all_core::IdxTensor;
4464    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4465    ///
4466    /// let i = Index::new_dyn(2);
4467    /// let j = Index::new_dyn(3);
4468    /// let new_i = Index::new_dyn(2);  // Same dimension, different ID
4469    ///
4470    /// let indices = vec![i.clone(), j.clone()];
4471    /// let tensor: IdxTensor = IdxTensor::from_dense(indices, vec![0.0; 6]).unwrap();
4472    ///
4473    /// // Replace index i with new_i
4474    /// let replaced = tensor.replaceind(&i, &new_i).unwrap();
4475    /// assert_eq!(replaced.indices[0].id, new_i.id);
4476    /// assert_eq!(replaced.indices[1].id, j.id);
4477    /// ```
4478    pub fn replaceind(
4479        &self,
4480        old_index: &DynIndex,
4481        new_index: &DynIndex,
4482    ) -> std::result::Result<Self, IdxTensorError> {
4483        // Validate dimension match
4484        if old_index.dim() != new_index.dim() {
4485            return Err(IdxTensorError::ShapeMismatch {
4486                operation: "replaceind",
4487                expected: format!("dimension {}", old_index.dim()),
4488                actual: format!("dimension {}", new_index.dim()),
4489            });
4490        }
4491
4492        let new_indices: Vec<_> = self
4493            .indices
4494            .iter()
4495            .map(|idx| {
4496                if *idx == *old_index {
4497                    new_index.clone()
4498                } else {
4499                    idx.clone()
4500                }
4501            })
4502            .collect();
4503
4504        Ok(Self {
4505            indices: new_indices,
4506            storage: self.storage.clone(),
4507            eager_cache: Arc::clone(&self.eager_cache),
4508        })
4509    }
4510
4511    /// Replace multiple indices in the tensor.
4512    ///
4513    /// This replaces each index in `old_indices` (matched by full index equality,
4514    /// including id, prime level, and tags) with the corresponding index in
4515    /// `new_indices`. The storage data is not modified.
4516    ///
4517    /// # Arguments
4518    /// * `old_indices` - The indices to replace (matched by full index equality)
4519    /// * `new_indices` - The new indices to use
4520    ///
4521    /// # Returns
4522    /// A new tensor with the indices replaced. Indices not found in `old_indices`
4523    /// are kept unchanged.
4524    ///
4525    /// # Errors
4526    /// Returns an error when `old_indices` and `new_indices` differ in length
4527    /// (a shape mismatch), or when any replacement index has an incompatible
4528    /// dimension (a shape mismatch: the replacement dimension must equal the
4529    /// original).
4530    /// # Example
4531    /// ```
4532    /// use tensor4all_core::IdxTensor;
4533    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4534    ///
4535    /// let i = Index::new_dyn(2);
4536    /// let j = Index::new_dyn(3);
4537    /// let new_i = Index::new_dyn(2);
4538    /// let new_j = Index::new_dyn(3);
4539    ///
4540    /// let indices = vec![i.clone(), j.clone()];
4541    /// let tensor: IdxTensor = IdxTensor::from_dense(indices, vec![0.0; 6]).unwrap();
4542    ///
4543    /// // Replace both indices
4544    /// let replaced = tensor
4545    ///     .replace_indices(&[i.clone(), j.clone()], &[new_i.clone(), new_j.clone()])
4546    ///     .unwrap();
4547    /// assert_eq!(replaced.indices[0].id, new_i.id);
4548    /// assert_eq!(replaced.indices[1].id, new_j.id);
4549    /// ```
4550    pub fn replace_indices(
4551        &self,
4552        old_indices: &[DynIndex],
4553        new_indices: &[DynIndex],
4554    ) -> std::result::Result<Self, IdxTensorError> {
4555        if old_indices.len() != new_indices.len() {
4556            return Err(IdxTensorError::ShapeMismatch {
4557                operation: "replace_indices",
4558                expected: format!("{} indices", old_indices.len()),
4559                actual: format!("{} indices", new_indices.len()),
4560            });
4561        }
4562
4563        // Validate dimension matches for all replacements
4564        for (old, new) in old_indices.iter().zip(new_indices.iter()) {
4565            if old.dim() != new.dim() {
4566                return Err(IdxTensorError::ShapeMismatch {
4567                    operation: "replace_indices",
4568                    expected: format!("dimension {}", old.dim()),
4569                    actual: format!("dimension {}", new.dim()),
4570                });
4571            }
4572        }
4573
4574        // Build a map from old indices to new indices
4575        let replacement_map: std::collections::HashMap<_, _> =
4576            old_indices.iter().zip(new_indices.iter()).collect();
4577
4578        let new_indices_vec: Vec<_> = self
4579            .indices
4580            .iter()
4581            .map(|idx| {
4582                if let Some(new_idx) = replacement_map.get(idx) {
4583                    (*new_idx).clone()
4584                } else {
4585                    idx.clone()
4586                }
4587            })
4588            .collect();
4589
4590        Ok(Self {
4591            indices: new_indices_vec,
4592            storage: self.storage.clone(),
4593            eager_cache: Arc::clone(&self.eager_cache),
4594        })
4595    }
4596}
4597
4598// ============================================================================
4599// Complex Conjugation
4600// ============================================================================
4601
4602impl IdxTensor {
4603    /// Complex conjugate of all tensor elements.
4604    ///
4605    /// For real (`f32`/`f64`) tensors, returns a copy (conjugate of real is
4606    /// identity). For complex (`Complex32`/`Complex64`) tensors, conjugates
4607    /// each element.
4608    ///
4609    /// The indices and dimensions remain unchanged. If an eager backend cannot
4610    /// perform the conjugation, the failure is retained and reported by the
4611    /// next fallible materialization or AD-sensitive operation.
4612    ///
4613    /// This is inspired by the `conj` operation in ITensorMPS.jl.
4614    ///
4615    /// # Example
4616    /// ```
4617    /// use tensor4all_core::IdxTensor;
4618    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4619    /// use num_complex::Complex64;
4620    ///
4621    /// let i = Index::new_dyn(2);
4622    /// let data = vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, -4.0)];
4623    /// let tensor: IdxTensor = IdxTensor::from_dense(vec![i], data).unwrap();
4624    ///
4625    /// let conj_tensor = tensor.conj();
4626    /// assert_eq!(
4627    ///     conj_tensor.to_vec::<Complex64>().unwrap(),
4628    ///     vec![Complex64::new(1.0, -2.0), Complex64::new(3.0, 4.0)]
4629    /// );
4630    /// ```
4631    pub fn conj(&self) -> Self {
4632        self.conj_with(&conjugate_eager)
4633    }
4634
4635    fn conj_with<F>(&self, conjugate: &F) -> Self
4636    where
4637        F: Fn(
4638            &EagerTensor,
4639        ) -> std::result::Result<
4640            EagerTensor,
4641            Arc<dyn std::error::Error + Send + Sync + 'static>,
4642        >,
4643    {
4644        // Conjugate tensor storage and map indices via IndexLike::conj(). For
4645        // default undirected indices, conj() is a no-op; this remains future-
4646        // proof for QSpace-compatible directed indices.
4647        let new_indices: Vec<DynIndex> = self.indices.iter().map(|idx| idx.conj()).collect();
4648        let mut storage = match self.storage.conjugate_with(conjugate) {
4649            Ok(storage) => storage,
4650            Err(error) => self.storage.clone().with_deferred_error(error),
4651        };
4652        let mut eager_cache = if storage.deferred_error().is_some() {
4653            Arc::clone(&self.eager_cache)
4654        } else {
4655            Self::empty_eager_cache()
4656        };
4657
4658        if storage.deferred_error().is_none() {
4659            if let Some(inner) = self.eager_cache.get() {
4660                match conjugate(inner.as_ref()) {
4661                    Ok(conjugated) => eager_cache = Self::eager_cache_with(conjugated),
4662                    Err(source) => {
4663                        storage =
4664                            storage.with_deferred_error(TensorStorageError::Conjugation { source });
4665                        // Keep the original cache alive so a deferred failure
4666                        // retains its graph until a fallible consumer reports it.
4667                        eager_cache = Arc::clone(&self.eager_cache);
4668                    }
4669                }
4670            }
4671        }
4672
4673        Self {
4674            indices: new_indices,
4675            storage,
4676            eager_cache,
4677        }
4678    }
4679}
4680
4681#[derive(Debug, Default)]
4682struct Lassq {
4683    scale: f64,
4684    sumsq: f64,
4685    infinite: bool,
4686}
4687
4688impl Lassq {
4689    fn add_component(&mut self, value: f64) {
4690        let value = value.abs();
4691        if value == 0.0 {
4692            return;
4693        }
4694        if value.is_infinite() {
4695            self.infinite = true;
4696            return;
4697        }
4698        if self.scale < value {
4699            if self.scale == 0.0 {
4700                self.sumsq = 1.0;
4701            } else {
4702                let ratio = self.scale / value;
4703                self.sumsq = 1.0 + self.sumsq * ratio * ratio;
4704            }
4705            self.scale = value;
4706        } else {
4707            let ratio = value / self.scale;
4708            self.sumsq += ratio * ratio;
4709        }
4710    }
4711
4712    fn add_complex(&mut self, value: Complex64) {
4713        self.add_component(value.re);
4714        self.add_component(value.im);
4715    }
4716
4717    fn add_scaled(&mut self, scale: f64, coefficient: f64) {
4718        if scale == 0.0 || coefficient == 0.0 {
4719            return;
4720        }
4721        if self.scale < scale {
4722            if self.scale == 0.0 {
4723                self.sumsq = coefficient * coefficient;
4724            } else {
4725                let ratio = self.scale / scale;
4726                self.sumsq = coefficient * coefficient + self.sumsq * ratio * ratio;
4727            }
4728            self.scale = scale;
4729        } else {
4730            let ratio = scale / self.scale * coefficient;
4731            self.sumsq += ratio * ratio;
4732        }
4733    }
4734
4735    fn add_component_difference(&mut self, lhs: f64, rhs: f64) {
4736        if lhs == rhs {
4737            return;
4738        }
4739        let scale = lhs.abs().max(rhs.abs());
4740        if scale != 0.0 {
4741            self.add_scaled(scale, (lhs / scale - rhs / scale).abs());
4742        }
4743    }
4744
4745    fn add_complex_difference(&mut self, lhs: Complex64, rhs: Complex64) {
4746        self.add_component_difference(lhs.re, rhs.re);
4747        self.add_component_difference(lhs.im, rhs.im);
4748    }
4749
4750    fn is_zero(&self) -> bool {
4751        !self.infinite && self.scale == 0.0
4752    }
4753
4754    fn norm(&self) -> f64 {
4755        if self.infinite {
4756            f64::INFINITY
4757        } else if self.scale == 0.0 {
4758            0.0
4759        } else {
4760            self.scale * self.sumsq.sqrt()
4761        }
4762    }
4763
4764    fn norm_squared(&self) -> f64 {
4765        let norm = self.norm();
4766        norm * norm
4767    }
4768
4769    fn log_norm(&self) -> f64 {
4770        if self.infinite {
4771            f64::INFINITY
4772        } else if self.scale == 0.0 {
4773            f64::NEG_INFINITY
4774        } else {
4775            self.scale.ln() + 0.5 * self.sumsq.ln()
4776        }
4777    }
4778}
4779
4780// ============================================================================
4781// Norm Computation
4782// ============================================================================
4783
4784impl IdxTensor {
4785    /// Compute the squared Frobenius norm of the tensor: ||T||² = Σ|T_ijk...|²
4786    ///
4787    /// For real tensors: sum of squares of all elements.
4788    /// For complex tensors: sum of `|z|²` over the compact payload.
4789    /// The reduction promotes source values to `f64` and uses a stable LASSQ
4790    /// accumulator, so it does not form a source-dtype `self * conj(self)`.
4791    ///
4792    /// # Errors
4793    /// Returns [`IdxTensorError`] when storage/materialization or scalar
4794    /// extraction fails, or when the input produces NaN. The result is
4795    /// accumulated from squared magnitudes, so it is never negative;
4796    /// positive infinity is preserved.
4797    ///
4798    /// # Example
4799    /// ```
4800    /// use tensor4all_core::IdxTensor;
4801    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4802    ///
4803    /// let i = Index::new_dyn(2);
4804    /// let j = Index::new_dyn(3);
4805    /// let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];  // 1² + 2² + ... + 6² = 91
4806    /// let tensor: IdxTensor = IdxTensor::from_dense(vec![i, j], data).unwrap();
4807    ///
4808    /// assert!((tensor.norm_squared().unwrap() - 91.0).abs() < 1e-10);
4809    /// ```
4810    pub fn norm_squared(&self) -> std::result::Result<f64, IdxTensorError> {
4811        let dtype = self
4812            .scalar_dtype()
4813            .map_err(IdxTensorError::scalar_extraction)?;
4814        if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4815            return Err(IdxTensorError::ScalarTypeMismatch {
4816                expected: "f32, f64, c32, or c64",
4817                actual: Self::dtype_name(dtype).to_string(),
4818            });
4819        }
4820        let (has_nan, _) = self
4821            .compact_nonfinite_flags()
4822            .map_err(IdxTensorError::materialization)?;
4823        if has_nan {
4824            return Err(IdxTensorError::NaNInput {
4825                operation: "norm_squared",
4826            });
4827        }
4828
4829        let mut norm = Lassq::default();
4830        self.storage
4831            .for_each_payload_value(|value| norm.add_complex(value))
4832            .map_err(IdxTensorError::materialization)?;
4833        let value = norm.norm_squared();
4834        if value.is_nan() {
4835            return Err(IdxTensorError::NaNInput {
4836                operation: "norm_squared",
4837            });
4838        }
4839        Ok(value)
4840    }
4841
4842    /// Compute the Frobenius norm of the tensor: ||T|| = sqrt(Σ|T_ijk...|²)
4843    ///
4844    /// # Errors
4845    /// Returns [`IdxTensorError`] when norm evaluation fails or when the
4846    /// input contains NaN. Positive infinity is preserved.
4847    ///
4848    /// # Example
4849    /// ```
4850    /// use tensor4all_core::IdxTensor;
4851    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
4852    ///
4853    /// let i = Index::new_dyn(2);
4854    /// let data = vec![3.0, 4.0];  // sqrt(9 + 16) = 5
4855    /// let tensor: IdxTensor = IdxTensor::from_dense(vec![i], data).unwrap();
4856    ///
4857    /// assert!((tensor.norm().unwrap() - 5.0).abs() < 1e-10);
4858    /// ```
4859    pub fn norm(&self) -> std::result::Result<f64, IdxTensorError> {
4860        Ok(self.norm_squared()?.sqrt())
4861    }
4862
4863    /// Maximum absolute value of all elements (L-infinity norm).
4864    ///
4865    /// # Errors
4866    /// Returns [`IdxTensorError`] when authoritative storage or eager
4867    /// materialization cannot be read, or when the input contains NaN.
4868    ///
4869    /// # Examples
4870    ///
4871    /// ```
4872    /// use tensor4all_core::{DynIndex, IdxTensor};
4873    ///
4874    /// let i = DynIndex::new_dyn(4);
4875    /// let t = IdxTensor::from_dense(vec![i], vec![-5.0, 1.0, 3.0, -2.0]).unwrap();
4876    /// assert!((t.maxabs().unwrap() - 5.0).abs() < 1e-12);
4877    /// ```
4878    pub fn maxabs(&self) -> std::result::Result<f64, IdxTensorError> {
4879        if let Some(error) = self.storage.deferred_error() {
4880            return Err(IdxTensorError::Storage {
4881                source: error.clone(),
4882            });
4883        }
4884        let dtype = self
4885            .storage
4886            .dtype()
4887            .ok_or_else(|| IdxTensorError::ScalarTypeMismatch {
4888                expected: "f32, f64, c32, or c64",
4889                actual: "unknown".to_string(),
4890            })?;
4891        if !matches!(dtype, DType::F32 | DType::F64 | DType::C32 | DType::C64) {
4892            return Err(IdxTensorError::ScalarTypeMismatch {
4893                expected: "f32, f64, c32, or c64",
4894                actual: Self::dtype_name(dtype).to_string(),
4895            });
4896        }
4897        let (has_nan, _) = self
4898            .compact_nonfinite_flags()
4899            .map_err(IdxTensorError::materialization)?;
4900        if has_nan {
4901            return Err(IdxTensorError::NaNInput {
4902                operation: "maxabs",
4903            });
4904        }
4905        let mut value = 0.0_f64;
4906        self.storage
4907            .for_each_payload_value(|scalar| {
4908                let magnitude = scalar.re.hypot(scalar.im);
4909                value = value.max(magnitude);
4910            })
4911            .map_err(IdxTensorError::materialization)?;
4912        Ok(value)
4913    }
4914
4915    fn native_complex_payload_value_at(
4916        native: &EagerTensor,
4917        payload_coords: &[usize],
4918    ) -> Result<Complex64> {
4919        if !(payload_coords.len() == native.shape().len()) {
4920            return Err(anyhow::anyhow!(
4921                "payload coordinate rank {} does not match payload rank {}",
4922                payload_coords.len(),
4923                native.shape().len()
4924            ));
4925        };
4926        for (&coordinate, &dim) in payload_coords.iter().zip(native.shape().iter()) {
4927            if coordinate >= dim {
4928                return Err(anyhow::anyhow!(
4929                    "payload coordinate {coordinate} is out of bounds for dim {dim}"
4930                ));
4931            }
4932        }
4933        let value = native.value()?;
4934        match value.as_tensor_view() {
4935            TensorView::F32(view) => view
4936                .get(payload_coords)
4937                .copied()
4938                .map(|value| Complex64::new(f64::from(value), 0.0))
4939                .ok_or_else(|| anyhow::anyhow!("failed to read f32 payload value")),
4940            TensorView::F64(view) => view
4941                .get(payload_coords)
4942                .copied()
4943                .map(|value| Complex64::new(value, 0.0))
4944                .ok_or_else(|| anyhow::anyhow!("failed to read f64 payload value")),
4945            TensorView::C32(view) => view
4946                .get(payload_coords)
4947                .copied()
4948                .map(|value| Complex64::new(f64::from(value.re), f64::from(value.im)))
4949                .ok_or_else(|| anyhow::anyhow!("failed to read c32 payload value")),
4950            TensorView::C64(view) => view
4951                .get(payload_coords)
4952                .copied()
4953                .ok_or_else(|| anyhow::anyhow!("failed to read c64 payload value")),
4954            view => Err(anyhow::anyhow!(
4955                "unsupported payload dtype {:?}",
4956                view.dtype()
4957            )),
4958        }
4959    }
4960
4961    fn native_sum_scalar(native: &EagerTensor) -> Result<AnyScalar> {
4962        let value = native.value()?;
4963        match native.dtype() {
4964            DType::F32 => Ok(AnyScalar::from_value(
4965                value.as_slice::<f32>()?.iter().copied().sum::<f32>(),
4966            )),
4967            DType::F64 => Ok(AnyScalar::from_value(
4968                value.as_slice::<f64>()?.iter().copied().sum::<f64>(),
4969            )),
4970            DType::C32 => Ok(AnyScalar::from_value(
4971                value
4972                    .as_slice::<Complex32>()?
4973                    .iter()
4974                    .copied()
4975                    .sum::<Complex32>(),
4976            )),
4977            DType::C64 => Ok(AnyScalar::from_value(
4978                value
4979                    .as_slice::<Complex64>()?
4980                    .iter()
4981                    .copied()
4982                    .sum::<Complex64>(),
4983            )),
4984            dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
4985        }
4986    }
4987
4988    fn native_nonfinite_flags_typed<T: TensorElement>(
4989        native: &EagerTensor,
4990        classify: impl Fn(T) -> (bool, bool),
4991    ) -> Result<(bool, bool)> {
4992        let value = native.value()?;
4993        if let Ok(values) = value.as_slice::<T>() {
4994            return Ok(values.iter().copied().map(&classify).fold(
4995                (false, false),
4996                |(has_nan, has_infinity), (is_nan, is_infinite)| {
4997                    (has_nan | is_nan, has_infinity | is_infinite)
4998                },
4999            ));
5000        }
5001        drop(value);
5002        let value = IdxTensorStorage::materialize_eager_payload(native)?;
5003        Ok(value.as_slice::<T>()?.iter().copied().map(classify).fold(
5004            (false, false),
5005            |(has_nan, has_infinity), (is_nan, is_infinite)| {
5006                (has_nan | is_nan, has_infinity | is_infinite)
5007            },
5008        ))
5009    }
5010
5011    fn native_nonfinite_flags(native: &EagerTensor) -> Result<(bool, bool)> {
5012        match native.dtype() {
5013            DType::F32 => Self::native_nonfinite_flags_typed(native, |value: f32| {
5014                (value.is_nan(), value.is_infinite())
5015            }),
5016            DType::F64 => Self::native_nonfinite_flags_typed(native, |value: f64| {
5017                (value.is_nan(), value.is_infinite())
5018            }),
5019            DType::C32 => Self::native_nonfinite_flags_typed(native, |value: Complex32| {
5020                (
5021                    value.re.is_nan() || value.im.is_nan(),
5022                    value.re.is_infinite() || value.im.is_infinite(),
5023                )
5024            }),
5025            DType::C64 => Self::native_nonfinite_flags_typed(native, |value: Complex64| {
5026                (
5027                    value.re.is_nan() || value.im.is_nan(),
5028                    value.re.is_infinite() || value.im.is_infinite(),
5029                )
5030            }),
5031            dtype => Err(anyhow::anyhow!("unsupported dtype {dtype:?}")),
5032        }
5033    }
5034
5035    fn compact_nonfinite_flags(&self) -> Result<(bool, bool)> {
5036        self.storage.nonfinite_flags()
5037    }
5038
5039    /// Element-wise subtraction with index alignment.
5040    ///
5041    /// This computes `self - other` using the same vector-space semantics as
5042    /// [`TensorVectorSpace`](crate::TensorVectorSpace).
5043    ///
5044    /// # Errors
5045    /// Returns an error when the tensors have different index sets (an index-set
5046    /// mismatch) or the arithmetic reports a failure.
5047    ///
5048    pub fn sub(&self, other: &Self) -> std::result::Result<Self, IdxTensorError> {
5049        self.axpby(AnyScalar::new_real(1.0), other, AnyScalar::new_real(-1.0))
5050    }
5051
5052    /// Negate all elements.
5053    ///
5054    /// # Errors
5055    /// Returns an error when scalar multiplication fails for the tensor storage
5056    /// (a dtype mismatch) or the backend reports a failure.
5057    ///
5058    pub fn neg(&self) -> std::result::Result<Self, IdxTensorError> {
5059        self.scale(AnyScalar::new_real(-1.0))
5060    }
5061
5062    /// Approximate equality check using Julia `isapprox`-style semantics.
5063    ///
5064    /// Values are aligned by index identity and streamed from each tensor's
5065    /// compact support. Exact zero-tolerance comparisons use exact scalar
5066    /// equality; nonzero tolerances use scaled sum-of-squares accumulation,
5067    /// avoiding logical-dense traversal and avoidable underflow/overflow.
5068    ///
5069    /// # Errors
5070    /// Returns [`IdxTensorError`] when tolerances are invalid, the index
5071    /// spaces cannot be aligned, storage cannot be read, or either input
5072    /// contains NaN.
5073    pub fn isapprox(
5074        &self,
5075        other: &Self,
5076        atol: f64,
5077        rtol: f64,
5078    ) -> std::result::Result<bool, IdxTensorError> {
5079        for (name, value) in [("atol", atol), ("rtol", rtol)] {
5080            if !value.is_finite() || value < 0.0 {
5081                return Err(IdxTensorError::InvalidTolerance { name, value });
5082            }
5083        }
5084        if self.indices.len() != other.indices.len() {
5085            return Err(IdxTensorError::ShapeMismatch {
5086                operation: "isapprox",
5087                expected: format!("indices {:?}", self.indices),
5088                actual: format!("indices {:?}", other.indices),
5089            });
5090        }
5091
5092        let other_axis_by_index = other
5093            .indices
5094            .iter()
5095            .cloned()
5096            .enumerate()
5097            .map(|(axis, index)| (index, axis))
5098            .collect::<HashMap<_, _>>();
5099        let other_positions = self
5100            .indices
5101            .iter()
5102            .map(|index| {
5103                other_axis_by_index.get(index).copied().ok_or_else(|| {
5104                    IdxTensorError::ShapeMismatch {
5105                        operation: "isapprox",
5106                        expected: format!("indices {:?}", self.indices),
5107                        actual: format!("indices {:?}", other.indices),
5108                    }
5109                })
5110            })
5111            .collect::<std::result::Result<Vec<_>, _>>()?;
5112        let self_dims = self.dims();
5113        let other_dims = other.dims();
5114        for (axis, &other_axis) in other_positions.iter().enumerate() {
5115            if self_dims[axis] != other_dims[other_axis] {
5116                return Err(IdxTensorError::ShapeMismatch {
5117                    operation: "isapprox",
5118                    expected: format!("dims {:?}", self_dims),
5119                    actual: format!("dims {:?}", other_dims),
5120                });
5121            }
5122        }
5123        for tensor in [self, other] {
5124            if tensor
5125                .compact_nonfinite_flags()
5126                .map_err(IdxTensorError::materialization)?
5127                .0
5128            {
5129                return Err(IdxTensorError::NaNInput {
5130                    operation: "isapprox",
5131                });
5132            }
5133        }
5134
5135        let exact = atol == 0.0 && rtol == 0.0;
5136        let lhs_payload_dims = self.storage.payload_dims().to_vec();
5137        let rhs_payload_dims = other.storage.payload_dims().to_vec();
5138        let lhs_payload_len =
5139            checked_product(&lhs_payload_dims).map_err(IdxTensorError::materialization)?;
5140        let rhs_payload_len =
5141            checked_product(&rhs_payload_dims).map_err(IdxTensorError::materialization)?;
5142        let lhs_axis_classes = self.storage.axis_classes();
5143        let rhs_axis_classes = other.storage.axis_classes();
5144        let self_to_other = other_positions.clone();
5145        let mut other_to_self = vec![0usize; self_to_other.len()];
5146        for (self_axis, &other_axis) in self_to_other.iter().enumerate() {
5147            other_to_self[other_axis] = self_axis;
5148        }
5149
5150        let mut diff = Lassq::default();
5151        let mut lhs_norm = Lassq::default();
5152        let mut rhs_norm = Lassq::default();
5153        let mut compare = |lhs: Complex64, rhs: Complex64| -> bool {
5154            if exact {
5155                return lhs == rhs;
5156            }
5157            let lhs_infinite = lhs.re.is_infinite() || lhs.im.is_infinite();
5158            let rhs_infinite = rhs.re.is_infinite() || rhs.im.is_infinite();
5159            if lhs_infinite || rhs_infinite {
5160                return lhs == rhs;
5161            }
5162            lhs_norm.add_complex(lhs);
5163            rhs_norm.add_complex(rhs);
5164            diff.add_complex_difference(lhs, rhs);
5165            true
5166        };
5167
5168        // Each compact payload coordinate identifies exactly one logical
5169        // support point. Map it through the aligned logical axes instead of
5170        // traversing structural zeros in the logical tensor.
5171        let mut lhs_coords = vec![0usize; lhs_payload_dims.len()];
5172        let mut rhs_from_lhs = vec![0usize; rhs_payload_dims.len()];
5173        let mut rhs_seen = vec![false; rhs_payload_dims.len()];
5174        for _ in 0..lhs_payload_len {
5175            let lhs = self
5176                .storage
5177                .payload_value_at(&lhs_coords)
5178                .map_err(IdxTensorError::materialization)?;
5179            if map_payload_support_coordinate(
5180                lhs_axis_classes,
5181                rhs_axis_classes,
5182                &self_to_other,
5183                &lhs_coords,
5184                &mut rhs_from_lhs,
5185                &mut rhs_seen,
5186            ) {
5187                let rhs = other
5188                    .storage
5189                    .payload_value_at(&rhs_from_lhs)
5190                    .map_err(IdxTensorError::materialization)?;
5191                if !compare(lhs, rhs) {
5192                    return Ok(false);
5193                }
5194            } else if !compare(lhs, Complex64::new(0.0, 0.0)) {
5195                return Ok(false);
5196            }
5197
5198            increment_col_major_coordinate(&mut lhs_coords, &lhs_payload_dims);
5199        }
5200
5201        // The reverse pass accounts for support points that exist only in the
5202        // right tensor. Overlap points were compared in the first pass and are
5203        // therefore skipped here without a payload-sized visited set.
5204        let mut rhs_coords = vec![0usize; rhs_payload_dims.len()];
5205        let mut lhs_from_rhs = vec![0usize; lhs_payload_dims.len()];
5206        let mut lhs_seen = vec![false; lhs_payload_dims.len()];
5207        for _ in 0..rhs_payload_len {
5208            let rhs = other
5209                .storage
5210                .payload_value_at(&rhs_coords)
5211                .map_err(IdxTensorError::materialization)?;
5212            if !map_payload_support_coordinate(
5213                rhs_axis_classes,
5214                lhs_axis_classes,
5215                &other_to_self,
5216                &rhs_coords,
5217                &mut lhs_from_rhs,
5218                &mut lhs_seen,
5219            ) && !compare(Complex64::new(0.0, 0.0), rhs)
5220            {
5221                return Ok(false);
5222            }
5223            increment_col_major_coordinate(&mut rhs_coords, &rhs_payload_dims);
5224        }
5225
5226        if exact {
5227            return Ok(true);
5228        }
5229        let absolute_ok = if diff.is_zero() {
5230            true
5231        } else if atol == 0.0 || diff.infinite {
5232            false
5233        } else {
5234            diff.log_norm() <= atol.ln()
5235        };
5236        let relative_ok = if rtol == 0.0 || diff.is_zero() {
5237            diff.is_zero()
5238        } else if diff.infinite {
5239            lhs_norm.infinite || rhs_norm.infinite
5240        } else if lhs_norm.infinite || rhs_norm.infinite {
5241            true
5242        } else {
5243            let reference_log = lhs_norm.log_norm().max(rhs_norm.log_norm());
5244            diff.log_norm() <= rtol.ln() + reference_log
5245        };
5246        Ok(absolute_ok || relative_ok)
5247    }
5248
5249    /// Create a diagonal Kronecker-delta tensor for one input/output index pair.
5250    ///
5251    /// # Errors
5252    /// Returns an error when the two indices have different dimensions (an
5253    /// index shape mismatch).
5254    ///
5255    pub fn diagonal(
5256        input_index: &DynIndex,
5257        output_index: &DynIndex,
5258    ) -> std::result::Result<Self, IdxTensorError> {
5259        <Self as TensorConstructionLike>::diagonal(input_index, output_index)
5260    }
5261
5262    /// Create a product of Kronecker-delta tensors for paired index lists.
5263    ///
5264    /// # Errors
5265    /// Returns an error if the index lists have different lengths or paired
5266    /// dimensions do not match.
5267    pub fn delta(
5268        input_indices: &[DynIndex],
5269        output_indices: &[DynIndex],
5270    ) -> std::result::Result<Self, IdxTensorError> {
5271        <Self as TensorConstructionLike>::delta(input_indices, output_indices)
5272    }
5273
5274    /// Create a scalar tensor equal to one.
5275    ///
5276    /// # Errors
5277    /// Returns an error when dense scalar construction fails for the element type
5278    /// (an invalid scalar dtype or a construction failure).
5279    ///
5280    pub fn scalar_one() -> std::result::Result<Self, IdxTensorError> {
5281        <Self as TensorConstructionLike>::scalar_one()
5282    }
5283
5284    /// Create a tensor filled with ones over the given indices.
5285    ///
5286    /// # Errors
5287    /// Returns an error when the tensor size overflows (an overflow failure) or
5288    /// dense construction fails.
5289    ///
5290    pub fn ones(indices: &[DynIndex]) -> std::result::Result<Self, IdxTensorError> {
5291        <Self as TensorConstructionLike>::ones(indices)
5292    }
5293
5294    /// Create a one-hot tensor with value one at the specified index positions.
5295    ///
5296    /// # Errors
5297    /// Returns an error when any coordinate is outside its index dimension (an
5298    /// out of bounds failure).
5299    ///
5300    pub fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, IdxTensorError> {
5301        <Self as TensorConstructionLike>::onehot(index_vals)
5302    }
5303
5304    /// Keep one coordinate along an index while retaining that index axis.
5305    ///
5306    /// This is the differentiable masking counterpart to [`Self::select_indices`].
5307    /// It selects the requested slice, forms a one-hot tensor over the removed
5308    /// axis in the source dtype, and takes an explicit tensor product to restore
5309    /// the original index order. The implementation stays in the tensor backend,
5310    /// so structured storage and reverse-mode metadata are preserved whenever
5311    /// the backend can represent the operation.
5312    ///
5313    /// # Arguments
5314    ///
5315    /// * `index` - Existing tensor index to mask.
5316    /// * `position` - Zero-based coordinate to keep; all other coordinates become
5317    ///
5318    ///   zero.
5319    ///
5320    /// # Errors
5321    /// Returns an error when the coordinate is outside the index dimension (an
5322    /// out of bounds failure) or the mask construction fails.
5323    /// # Examples
5324    ///
5325    /// ```
5326    /// use tensor4all_core::{DynIndex, IdxTensor};
5327    ///
5328    /// let i = DynIndex::new_dyn(2);
5329    /// let tensor = IdxTensor::from_dense(vec![i.clone()], vec![3.0_f64, 4.0]).unwrap();
5330    /// let masked = tensor.mask_index(&i, 1).unwrap();
5331    ///
5332    /// assert_eq!(masked.indices(), &[i]);
5333    /// assert_eq!(masked.to_vec::<f64>().unwrap(), vec![0.0, 4.0]);
5334    /// assert!(IdxTensor::from_dense(
5335    ///     vec![DynIndex::new_dyn(2)],
5336    ///     vec![1.0_f64, 2.0],
5337    /// )
5338    /// .unwrap()
5339    /// .mask_index(&DynIndex::new_dyn(2), 0)
5340    /// .is_err());
5341    /// ```
5342    pub fn mask_index(
5343        &self,
5344        index: &DynIndex,
5345        position: usize,
5346    ) -> std::result::Result<Self, IdxTensorError> {
5347        if !(self.indices.iter().any(|candidate| candidate == index)) {
5348            return Err(anyhow::anyhow!("mask_index: index is not present in tensor").into());
5349        };
5350        if !(position < index.dim()) {
5351            return Err(anyhow::anyhow!(
5352                "mask_index: position {position} is out of range for dimension {}",
5353                index.dim()
5354            )
5355            .into());
5356        };
5357
5358        // Retaining the shared index turns contraction into a backend-level
5359        // elementwise product instead of materializing a host mask. Construct
5360        // the constant mask in the input dtype so f32/c32 values and AD graphs
5361        // are not promoted or detached.
5362        let mask = match self.scalar_dtype()? {
5363            DType::F32 => Self::from_dense(
5364                vec![index.clone()],
5365                (0..index.dim())
5366                    .map(|value| if value == position { 1.0_f32 } else { 0.0 })
5367                    .collect(),
5368            ),
5369            DType::F64 => Self::from_dense(
5370                vec![index.clone()],
5371                (0..index.dim())
5372                    .map(|value| if value == position { 1.0_f64 } else { 0.0 })
5373                    .collect(),
5374            ),
5375            DType::C32 => Self::from_dense(
5376                vec![index.clone()],
5377                (0..index.dim())
5378                    .map(|value| {
5379                        if value == position {
5380                            num_complex::Complex32::new(1.0, 0.0)
5381                        } else {
5382                            num_complex::Complex32::new(0.0, 0.0)
5383                        }
5384                    })
5385                    .collect(),
5386            ),
5387            DType::C64 => Self::from_dense(
5388                vec![index.clone()],
5389                (0..index.dim())
5390                    .map(|value| {
5391                        if value == position {
5392                            Complex64::new(1.0, 0.0)
5393                        } else {
5394                            Complex64::new(0.0, 0.0)
5395                        }
5396                    })
5397                    .collect(),
5398            ),
5399            dtype => {
5400                return Err(anyhow::anyhow!("mask_index does not support dtype {dtype:?}").into())
5401            }
5402        }?;
5403        super::contract::contract_pair_with_options(
5404            self,
5405            &mask,
5406            super::contract::ContractionOptions::new()
5407                .with_retain_indices(std::slice::from_ref(index)),
5408        )
5409    }
5410
5411    /// Compute the relative distance between two tensors.
5412    ///
5413    /// Returns `||A - B|| / ||A||` (Frobenius norm).
5414    /// If `||A|| = 0`, returns `||B||` instead to avoid division by zero.
5415    ///
5416    /// This is the ITensor-style distance function useful for comparing tensors.
5417    ///
5418    /// # Arguments
5419    /// * `other` - The other tensor to compare with
5420    ///
5421    /// # Errors
5422    /// Returns [`IdxTensorError`] when either norm contains NaN, or when
5423    /// scaling and subtracting the tensors fails.
5424    ///
5425    /// # Returns
5426    /// The relative distance as a f64 value.
5427    ///
5428    /// # Note
5429    /// The indices of both tensors must be permutable to each other.
5430    /// The result tensor (A - B) uses the index ordering from self.
5431    ///
5432    /// # Example
5433    /// ```
5434    /// use tensor4all_core::IdxTensor;
5435    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
5436    ///
5437    /// let i = Index::new_dyn(2);
5438    /// let data_a = vec![1.0, 0.0];
5439    /// let data_b = vec![1.0, 0.0];  // Same tensor
5440    /// let tensor_a: IdxTensor = IdxTensor::from_dense(vec![i.clone()], data_a).unwrap();
5441    /// let tensor_b: IdxTensor = IdxTensor::from_dense(vec![i.clone()], data_b).unwrap();
5442    ///
5443    /// assert!(tensor_a.distance(&tensor_b).unwrap() < 1e-10);  // Zero distance
5444    /// ```
5445    pub fn distance(&self, other: &Self) -> std::result::Result<f64, IdxTensorError> {
5446        let norm_self = self.norm()?;
5447
5448        // Compute A - B = A + (-1) * B
5449        let neg_other = other.scale(AnyScalar::new_real(-1.0))?;
5450        let diff = self.add(&neg_other)?;
5451        let norm_diff = diff.norm()?;
5452
5453        if norm_self > 0.0 {
5454            Ok(norm_diff / norm_self)
5455        } else {
5456            Ok(norm_diff)
5457        }
5458    }
5459}
5460
5461impl std::fmt::Debug for IdxTensor {
5462    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5463        f.debug_struct("IdxTensor")
5464            .field("indices", &self.indices)
5465            .field("dims", &self.dims())
5466            .field("is_diag", &self.is_diag())
5467            .finish()
5468    }
5469}
5470
5471/// Create a diagonal tensor with dynamic rank from diagonal data.
5472/// # Arguments
5473/// * `indices` - The indices for the tensor (all must have the same dimension)
5474/// * `diag_data` - The diagonal elements (length must equal the dimension of indices)
5475///
5476/// The returned tensor preserves compact diagonal payload metadata; use
5477/// [`IdxTensor::is_diag`] or [`IdxTensor::storage`] to inspect that
5478/// representation.
5479///
5480/// # Errors
5481/// Returns an error when the index dimensions are unequal (a dimension
5482/// mismatch) or the diagonal construction fails.
5483/// # Panics
5484/// Panics if indices have different dimensions, or if diag_data length doesn't match.
5485/// # Examples
5486/// ```
5487/// use tensor4all_core::{DynIndex, diag_idx_tensor};
5488/// let i = DynIndex::new_dyn(3);
5489/// let j = DynIndex::new_dyn(3);
5490/// let t = diag_idx_tensor(vec![i, j], vec![1.0, 2.0, 3.0]).unwrap();
5491/// assert_eq!(t.dims(), vec![3, 3]);
5492/// assert!(t.is_diag());
5493/// ```
5494pub fn diag_idx_tensor(
5495    indices: Vec<DynIndex>,
5496    diag_data: Vec<f64>,
5497) -> std::result::Result<IdxTensor, IdxTensorError> {
5498    IdxTensor::from_diag(indices, diag_data)
5499}
5500
5501#[allow(clippy::type_complexity)]
5502pub(crate) type UnfoldSplitInnerResult = (
5503    EagerTensor,
5504    usize,
5505    usize,
5506    usize,
5507    Vec<DynIndex>,
5508    Vec<DynIndex>,
5509);
5510
5511/// Unfold a tensor into a matrix by splitting indices into left and right groups.
5512/// This function validates the split, permutes the tensor so that left indices
5513/// come first, and returns a rank-2 native tenferro tensor along with metadata.
5514/// # Arguments
5515/// * `t` - Input tensor
5516/// * `left_inds` - Indices to place on the left (row) side of the matrix
5517/// # Returns
5518/// A tuple `(matrix_tensor, left_len, m, n, left_indices, right_indices)` where:
5519/// - `matrix_tensor` is a rank-2 `tenferro::Tensor` with shape `[m, n]`
5520/// - `left_len` is the number of left indices
5521/// - `m` is the product of left index dimensions
5522/// - `n` is the product of right index dimensions
5523/// - `left_indices` is the vector of left indices (cloned)
5524/// - `right_indices` is the vector of right indices (cloned)
5525/// # Errors
5526///
5527/// Returns an error when the tensor rank is less than 2 (a rank mismatch),
5528/// when `left_inds` is empty or contains all indices (an invalid split), when
5529/// `left_inds` contains indices not present in the tensor (a missing-index
5530/// failure) or duplicates, or when the native reshape fails (a backend
5531/// failure).
5532/// # Examples
5533/// ```
5534/// use tensor4all_core::{DynIndex, IdxTensor, unfold_split};
5535/// let i = DynIndex::new_dyn(2);
5536/// let j = DynIndex::new_dyn(3);
5537/// // 2x3 dense tensor with data [1..6]
5538/// let t = IdxTensor::from_dense(
5539///     vec![i.clone(), j.clone()],
5540///     vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
5541/// ).unwrap();
5542/// let (matrix, left_len, m, n, left_indices, right_indices) =
5543///     unfold_split(&t, &[i]).unwrap();
5544/// assert_eq!(left_len, 1);
5545/// assert_eq!(m, 2);
5546/// assert_eq!(n, 3);
5547/// assert_eq!(left_indices.len(), 1);
5548/// assert_eq!(right_indices.len(), 1);
5549/// ```
5550#[allow(clippy::type_complexity)]
5551pub fn unfold_split(
5552    t: &IdxTensor,
5553    left_inds: &[DynIndex],
5554) -> std::result::Result<
5555    (
5556        NativeTensor,
5557        usize,
5558        usize,
5559        usize,
5560        Vec<DynIndex>,
5561        Vec<DynIndex>,
5562    ),
5563    IdxTensorError,
5564> {
5565    let (matrix_inner, left_len, m, n, left_indices, right_indices) =
5566        unfold_split_inner(t, left_inds)?;
5567
5568    Ok((
5569        matrix_inner.duplicate_value()?,
5570        left_len,
5571        m,
5572        n,
5573        left_indices,
5574        right_indices,
5575    ))
5576}
5577
5578pub(crate) fn unfold_split_inner(
5579    t: &IdxTensor,
5580    left_inds: &[DynIndex],
5581) -> Result<UnfoldSplitInnerResult> {
5582    let rank = t.indices.len();
5583
5584    // Validate rank
5585    if !(rank >= 2) {
5586        return Err(anyhow::anyhow!(
5587            "Tensor must have rank >= 2, got rank {}",
5588            rank
5589        ));
5590    };
5591
5592    let left_len = left_inds.len();
5593
5594    // Validate split: must be a proper subset
5595    if !(left_len > 0 && left_len < rank) {
5596        return Err(anyhow::anyhow!("Left indices must be a non-empty proper subset of tensor indices (0 < left_len < rank), got left_len={}, rank={}",
5597        left_len,
5598        rank));
5599    };
5600
5601    // Validate that all left_inds are in the tensor and there are no duplicates
5602    let tensor_set: HashSet<_> = t.indices.iter().collect();
5603    let mut left_set = HashSet::new();
5604
5605    for left_idx in left_inds {
5606        if !(tensor_set.contains(left_idx)) {
5607            return Err(anyhow::anyhow!("Index in left_inds not found in tensor"));
5608        };
5609        if !(left_set.insert(left_idx)) {
5610            return Err(anyhow::anyhow!("Duplicate index in left_inds"));
5611        };
5612    }
5613
5614    // Build right_inds: all indices not in left_inds, in original order
5615    let mut right_inds = Vec::new();
5616    for idx in &t.indices {
5617        if !left_set.contains(idx) {
5618            right_inds.push(idx.clone());
5619        }
5620    }
5621
5622    // Build new_indices: left_inds first, then right_inds
5623    let mut new_indices = Vec::with_capacity(rank);
5624    new_indices.extend_from_slice(left_inds);
5625    new_indices.extend_from_slice(&right_inds);
5626
5627    // Permute tensor to have left indices first, then right indices
5628    let unfolded = t.permute_indices(&new_indices)?;
5629
5630    // Compute matrix dimensions
5631    let unfolded_dims = unfolded.dims();
5632    let m = checked_product(&unfolded_dims[..left_len])?;
5633    let n = checked_product(&unfolded_dims[left_len..])?;
5634
5635    let matrix_tensor = unfolded.try_materialized_inner()?.reshape(&[m, n])?;
5636
5637    Ok((
5638        matrix_tensor,
5639        left_len,
5640        m,
5641        n,
5642        left_inds.to_vec(),
5643        right_inds,
5644    ))
5645}
5646
5647// ============================================================================
5648// TensorIndex implementation for IdxTensor
5649// ============================================================================
5650
5651use crate::tensor_index::TensorIndex;
5652
5653impl TensorIndex for IdxTensor {
5654    type Index = DynIndex;
5655    type Error = IdxTensorError;
5656
5657    fn external_indices(&self) -> Vec<DynIndex> {
5658        // For IdxTensor, all indices are external.
5659        self.indices.clone()
5660    }
5661
5662    fn num_external_indices(&self) -> usize {
5663        self.indices.len()
5664    }
5665
5666    fn replaceind(
5667        &self,
5668        old_index: &DynIndex,
5669        new_index: &DynIndex,
5670    ) -> std::result::Result<Self, Self::Error> {
5671        // Delegate to the inherent method.
5672        IdxTensor::replaceind(self, old_index, new_index)
5673    }
5674
5675    fn replace_indices(
5676        &self,
5677        old_indices: &[DynIndex],
5678        new_indices: &[DynIndex],
5679    ) -> std::result::Result<Self, Self::Error> {
5680        // Delegate to the inherent method.
5681        IdxTensor::replace_indices(self, old_indices, new_indices)
5682    }
5683}
5684
5685// ============================================================================
5686// TensorLike implementation for IdxTensor
5687// ============================================================================
5688
5689use crate::tensor_like::{
5690    FactorizeError, FactorizeOptions, FactorizeResult, IncrementalQrState, TensorConstructionLike,
5691    TensorContractionLike, TensorFactorizationLike, TensorVectorSpace,
5692};
5693
5694impl TensorVectorSpace for IdxTensor {
5695    fn norm_squared(&self) -> std::result::Result<f64, Self::Error> {
5696        IdxTensor::norm_squared(self)
5697    }
5698
5699    fn maxabs(&self) -> std::result::Result<f64, Self::Error> {
5700        IdxTensor::maxabs(self)
5701    }
5702
5703    fn isapprox(
5704        &self,
5705        other: &Self,
5706        atol: f64,
5707        rtol: f64,
5708    ) -> std::result::Result<bool, Self::Error> {
5709        IdxTensor::isapprox(self, other, atol, rtol)
5710    }
5711
5712    fn axpby(
5713        &self,
5714        a: crate::AnyScalar,
5715        other: &Self,
5716        b: crate::AnyScalar,
5717    ) -> std::result::Result<Self, Self::Error> {
5718        IdxTensor::axpby(self, a, other, b)
5719    }
5720
5721    fn scale(&self, scalar: crate::AnyScalar) -> std::result::Result<Self, Self::Error> {
5722        IdxTensor::scale(self, scalar)
5723    }
5724
5725    fn scale_in(
5726        &self,
5727        factor: f64,
5728        context: &ExecutionContext,
5729    ) -> std::result::Result<Self, Self::Error> {
5730        IdxTensor::scale_in(self, factor, context)
5731    }
5732
5733    fn norm_in(&self, context: &ExecutionContext) -> std::result::Result<f64, Self::Error> {
5734        IdxTensor::norm_in(self, context)
5735    }
5736
5737    fn inner_product(&self, other: &Self) -> std::result::Result<crate::AnyScalar, Self::Error> {
5738        IdxTensor::inner_product(self, other)
5739    }
5740}
5741
5742impl TensorFactorizationLike for IdxTensor {
5743    fn factorize(
5744        &self,
5745        left_inds: &[DynIndex],
5746        options: &FactorizeOptions,
5747    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5748        crate::factorize::factorize(self, left_inds, options)
5749    }
5750
5751    fn factorize_in(
5752        &self,
5753        left_inds: &[DynIndex],
5754        options: &FactorizeOptions,
5755        context: &ExecutionContext,
5756    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5757        crate::factorize::factorize_in(self, left_inds, options, context)
5758    }
5759
5760    fn factorize_auto(
5761        &self,
5762        left_inds: &[DynIndex],
5763        options: &FactorizeOptions,
5764    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5765        crate::factorize::factorize_auto(self, left_inds, options)
5766    }
5767
5768    fn factorize_full_rank(
5769        &self,
5770        left_inds: &[DynIndex],
5771        alg: crate::FactorizeAlg,
5772        canonical: crate::Canonical,
5773    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5774        crate::factorize::factorize_full_rank(self, left_inds, alg, canonical)
5775    }
5776
5777    fn factorize_full_rank_in(
5778        &self,
5779        left_inds: &[DynIndex],
5780        alg: crate::FactorizeAlg,
5781        canonical: crate::Canonical,
5782        context: &ExecutionContext,
5783    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5784        crate::factorize::factorize_full_rank_in(self, left_inds, alg, canonical, context)
5785    }
5786
5787    fn factorize_probe_columns_incremental(
5788        previous: Option<&FactorizeResult<Self>>,
5789        all_columns: &[&Self],
5790        appended_columns: &[&Self],
5791        left_inds: &[DynIndex],
5792    ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
5793        factorize_probe_columns_incremental(previous, all_columns, appended_columns, left_inds)
5794    }
5795
5796    fn factorize_probe_batch_incremental(
5797        previous: Option<&FactorizeResult<IdxTensor>>,
5798        batch_tensor: &IdxTensor,
5799        batch_index: &DynIndex,
5800        left_inds: &[DynIndex],
5801    ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5802        factorize_probe_batch_incremental_impl(previous, batch_tensor, batch_index, left_inds)
5803    }
5804
5805    fn factorize_probe_batch_incremental_in(
5806        previous: Option<&FactorizeResult<IdxTensor>>,
5807        batch_tensor: &IdxTensor,
5808        batch_index: &DynIndex,
5809        left_inds: &[DynIndex],
5810        context: &ExecutionContext,
5811    ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5812        if let Some(previous) = previous {
5813            previous
5814                .left
5815                .validate_context(context)
5816                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5817            previous
5818                .right
5819                .validate_context(context)
5820                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5821        }
5822        batch_tensor
5823            .validate_context(context)
5824            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5825        // Compatibility boundary: the process-global CPU context keeps the
5826        // historical host `IncrementalQr` matrix path (bitwise-identical
5827        // numerics for legacy callers). Every other explicit context uses the
5828        // scoped recompute path with no host round-trip.
5829        if context.is_global_default_cpu() {
5830            return factorize_probe_batch_incremental_impl(
5831                previous,
5832                batch_tensor,
5833                batch_index,
5834                left_inds,
5835            );
5836        }
5837        Self::resident_probe_batch_qr(previous, batch_tensor, batch_index, left_inds, context)
5838    }
5839
5840    fn src_error_estimate(
5841        &self,
5842    ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5843        if self.is_cuda_resident() {
5844            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5845                "CUDA-resident QR factor requires src_error_estimate_in with the owning execution context"
5846            )));
5847        }
5848        let indices = self.indices();
5849        if indices.len() != 2 {
5850            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5851                "SRC QR factor must have rank 2, got {}",
5852                indices.len()
5853            )));
5854        }
5855        let nrows = indices[0].dim();
5856        let ncols = indices[1].dim();
5857        let result = if self.is_f64() {
5858            let matrix = Matrix::from_col_major_vec(
5859                nrows,
5860                ncols,
5861                self.to_vec::<f64>()
5862                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5863            );
5864            backend_src_error_estimate(&matrix)
5865        } else if self.is_c64() {
5866            let matrix = Matrix::from_col_major_vec(
5867                nrows,
5868                ncols,
5869                self.to_vec::<Complex64>()
5870                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5871            );
5872            backend_src_error_estimate(&matrix)
5873        } else {
5874            return Err(FactorizeError::UnsupportedStorage(
5875                "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
5876            ));
5877        };
5878        result.map_err(|error| FactorizeError::ComputationError(anyhow::anyhow!(error)))
5879    }
5880
5881    fn src_error_estimate_in(
5882        &self,
5883        context: &ExecutionContext,
5884    ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5885        self.validate_context(context)
5886            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
5887        #[cfg(feature = "tenferro-cuda")]
5888        if matches!(context, ExecutionContext::Cuda(_)) {
5889            return self.resident_src_error_estimate(context);
5890        }
5891        if context.is_global_default_cpu() {
5892            return TensorFactorizationLike::src_error_estimate(self);
5893        }
5894        self.host_src_error_estimate_general()
5895    }
5896}
5897
5898impl IdxTensor {
5899    /// Evaluate the SRC estimator for an explicit CPU context's restored factor.
5900    fn host_src_error_estimate_general(
5901        &self,
5902    ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
5903        let indices = self.indices();
5904        if indices.len() != 2 {
5905            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5906                "SRC QR factor must have rank 2, got {}",
5907                indices.len()
5908            )));
5909        }
5910        let nrows = indices[0].dim();
5911        let ncols = indices[1].dim();
5912        let result = if self.is_f64() {
5913            let matrix = Matrix::from_col_major_vec(
5914                nrows,
5915                ncols,
5916                self.to_vec::<f64>()
5917                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5918            );
5919            backend_src_error_estimate_general(&matrix)
5920        } else if self.is_c64() {
5921            let matrix = Matrix::from_col_major_vec(
5922                nrows,
5923                ncols,
5924                self.to_vec::<Complex64>()
5925                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
5926            );
5927            backend_src_error_estimate_general(&matrix)
5928        } else {
5929            return Err(FactorizeError::UnsupportedStorage(
5930                "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
5931            ));
5932        };
5933        result.map_err(|error| FactorizeError::ComputationError(anyhow::anyhow!(error)))
5934    }
5935
5936    /// Resident batch-native probe factorization by full RRQR recompute.
5937    ///
5938    /// Reconstructs the full sketch block `[Q_prev R_prev, A_new]` with resident
5939    /// matmul/concatenation and runs column-pivoted rank-revealing QR in the
5940    /// owning runtime. Only the scalar rank metadata is explicitly read back;
5941    /// matrix payloads and permutation metadata remain resident. The returned
5942    /// right factor is projected back into original sketch-column order, so
5943    /// `left * right` reconstructs the unpermuted sketch. Since that right
5944    /// factor is not generally triangular, the adaptive estimator uses a
5945    /// general solve of the small square sketch factor.
5946    fn resident_probe_batch_qr(
5947        previous: Option<&FactorizeResult<IdxTensor>>,
5948        batch_tensor: &IdxTensor,
5949        batch_index: &DynIndex,
5950        left_inds: &[DynIndex],
5951        context: &ExecutionContext,
5952    ) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
5953        use crate::defaults::idx_tensor::unfold_split_inner;
5954
5955        // The host path consumes `batch_index` as the appended-axis label;
5956        // reject a mismatched label before any arithmetic, as the device path
5957        // otherwise never observes it.
5958        if !batch_tensor.indices().contains(batch_index) {
5959            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5960                "resident probe batch is missing its batch index"
5961            )));
5962        }
5963        let (appended_inner, _, m, b, _, _) = unfold_split_inner(batch_tensor, left_inds)
5964            .context("resident probe batch unfold failed")
5965            .map_err(FactorizeError::ComputationError)?;
5966        if !matches!(appended_inner.dtype(), DType::F64 | DType::C64) {
5967            return Err(FactorizeError::UnsupportedStorage(
5968                "incremental SRC factorization currently supports f64 and Complex64 tensors",
5969            ));
5970        }
5971        if b == 0 {
5972            if let Some(previous) = previous {
5973                return Ok(previous.clone());
5974            }
5975            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5976                "incremental SRC factorization requires at least one column"
5977            )));
5978        }
5979        let full_inner = if let Some(previous) = previous {
5980            let (q_inner, _, previous_rows, _, _, _) =
5981                unfold_split_inner(&previous.left, left_inds)
5982                    .context("resident previous Q unfold failed")
5983                    .map_err(FactorizeError::ComputationError)?;
5984            if previous_rows != m {
5985                return Err(FactorizeError::ComputationError(anyhow::anyhow!(
5986                    "previous and appended sketch row dimensions differ: {previous_rows} vs {m}"
5987                )));
5988            }
5989            let (r_inner, _, _, _, _, _) =
5990                unfold_split_inner(&previous.right, std::slice::from_ref(&previous.bond_index))
5991                    .context("resident previous R unfold failed")
5992                    .map_err(FactorizeError::ComputationError)?;
5993            let reconstructed = q_inner.matmul(&r_inner).map_err(|error| {
5994                FactorizeError::ComputationError(
5995                    anyhow::Error::new(error).context("resident sketch reconstruction failed"),
5996                )
5997            })?;
5998            EagerTensor::concatenate(&[&reconstructed, &appended_inner], 1).map_err(|error| {
5999                FactorizeError::ComputationError(
6000                    anyhow::Error::new(error).context("resident sketch concatenation failed"),
6001                )
6002            })?
6003        } else {
6004            appended_inner
6005        };
6006        let total_width = *full_inner.shape().get(1).ok_or_else(|| {
6007            FactorizeError::ComputationError(anyhow::anyhow!(
6008                "resident sketch factor is not rank-2"
6009            ))
6010        })?;
6011        let rank_tolerance = 32.0 * f64::EPSILON * m.max(total_width) as f64;
6012        let decomposition = full_inner
6013            .rank_revealing_qr(RankRevealingQrOptions::default().rtol(rank_tolerance))
6014            .map_err(|error| {
6015                FactorizeError::ComputationError(
6016                    anyhow::Error::new(error).context("resident probe batch RRQR failed"),
6017                )
6018            })?;
6019        let rank = Self::read_resident_rank(&decomposition.rank, context)?;
6020        let maximum_rank = m.min(total_width);
6021        if rank > maximum_rank {
6022            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6023                "resident RRQR returned rank {rank}, exceeding maximum {maximum_rank}"
6024            )));
6025        }
6026        if let Some(previous) = previous {
6027            if rank < previous.rank {
6028                return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6029                    "resident SRC RRQR rank decreased from {} to {rank}",
6030                    previous.rank
6031                )));
6032            }
6033        }
6034        let q_full = decomposition
6035            .q
6036            .slice_axis(1, 0..rank)
6037            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6038        // RRQR factors the pivoted sketch A[:, permutation]. Restore the
6039        // original sketch-column order without reading permutation metadata:
6040        // Q_rank^H A computes the equivalent right factor entirely resident.
6041        let q_adjoint = q_full
6042            .transpose(&[1, 0])
6043            .and_then(|transposed| transposed.conj())
6044            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6045        let r_full = q_adjoint.matmul(&full_inner).map_err(|error| {
6046            FactorizeError::ComputationError(
6047                anyhow::Error::new(error).context("resident RRQR right-factor restoration failed"),
6048            )
6049        })?;
6050        let cap = DynIndex::new_bond(rank)
6051            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6052        let batch = DynIndex::new_link(total_width)
6053            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6054        let mut q_indices = left_inds.to_vec();
6055        q_indices.push(cap.clone());
6056        let q_dims: Vec<usize> = q_indices.iter().map(|index| index.dim()).collect();
6057        let left = Self::from_inner(
6058            q_indices,
6059            q_full.reshape(&q_dims).map_err(|error| {
6060                FactorizeError::ComputationError(
6061                    anyhow::Error::new(error).context("resident probe batch Q reshape failed"),
6062                )
6063            })?,
6064        )
6065        .map_err(FactorizeError::ComputationError)?;
6066        let right = Self::from_inner(
6067            vec![cap.clone(), batch],
6068            r_full.reshape(&[rank, total_width]).map_err(|error| {
6069                FactorizeError::ComputationError(
6070                    anyhow::Error::new(error).context("resident probe batch R reshape failed"),
6071                )
6072            })?,
6073        )
6074        .map_err(FactorizeError::ComputationError)?;
6075        Ok(FactorizeResult::new(left, right, cap, None, rank))
6076    }
6077
6078    /// Device-side Appendix-C estimator: solve and reduce on-device, read back
6079    /// `ncols + 1` decision scalars.
6080    ///
6081    /// Singularity and non-finiteness surface as typed solve/validation errors;
6082    /// unlike the host path there is no separate diagonal pre-check, but every
6083    /// invalid input still fails instead of producing an estimate.
6084    #[cfg(feature = "tenferro-cuda")]
6085    fn resident_src_error_estimate(
6086        &self,
6087        context: &ExecutionContext,
6088    ) -> std::result::Result<tensor4all_tensorbackend::SrcErrorEstimate, FactorizeError> {
6089        use tensor4all_tensorbackend::SrcErrorEstimate;
6090
6091        let indices = self.indices();
6092        if indices.len() != 2 {
6093            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6094                "SRC QR factor must have rank 2, got {}",
6095                indices.len()
6096            )));
6097        }
6098        let nrows = indices[0].dim();
6099        let ncols = indices[1].dim();
6100        if nrows != ncols {
6101            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6102                "SRC estimator requires a square R, got {nrows}x{ncols}"
6103            )));
6104        }
6105        if nrows == 0 {
6106            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6107                "SRC estimator requires a non-empty R"
6108            )));
6109        }
6110        let inner = self.cuda_eager_inner().ok_or_else(|| {
6111            FactorizeError::ComputationError(anyhow::anyhow!(
6112                "SRC estimator requires a resident eager QR factor"
6113            ))
6114        })?;
6115        let ExecutionContext::Cuda(cuda) = context else {
6116            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6117                "SRC device estimator requires a CUDA execution context"
6118            )));
6119        };
6120        // RRQR restores original sketch-column order, so the factor is square
6121        // but not generally triangular. Solve its adjoint as a general system.
6122        let adjoint = inner
6123            .transpose(&[1, 0])
6124            .and_then(|transposed| transposed.conj())
6125            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6126        let identity = Self::resident_identity(cuda, nrows, inner.dtype())?;
6127        // Solve F^H X = I for X = F^{-†}.
6128        let solved = adjoint.solve(&identity).map_err(|error| {
6129            FactorizeError::ComputationError(
6130                anyhow::Error::new(error).context("SRC inverse-adjoint general solve failed"),
6131            )
6132        })?;
6133        let column_sums = solved
6134            .abs()
6135            .and_then(|magnitudes| magnitudes.reduce_sum_squares(&[0]))
6136            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6137        let total = inner
6138            .abs()
6139            .and_then(|magnitudes| magnitudes.reduce_sum_squares(&[0, 1]))
6140            .and_then(|scalar| scalar.reshape(&[1]))
6141            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6142        let column_norms_sq = Self::read_resident_vector(&column_sums, ncols, context)?;
6143        let total_norm_sq = Self::read_resident_vector(&total, 1, context)?;
6144        Self::src_estimate_from_decision_scalars(&column_norms_sq, total_norm_sq[0], ncols)
6145            .map(|(error, norm)| SrcErrorEstimate { error, norm })
6146            .map_err(FactorizeError::ComputationError)
6147    }
6148
6149    /// Read the scalar integer rank returned by RRQR through its owning context.
6150    fn read_resident_rank(
6151        rank: &EagerTensor,
6152        context: &ExecutionContext,
6153    ) -> std::result::Result<usize, FactorizeError> {
6154        if !rank.shape().is_empty() {
6155            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6156                "resident RRQR rank metadata has shape {:?}, expected []",
6157                rank.shape()
6158            )));
6159        }
6160        let resident = rank.to_tensor().map_err(|error| {
6161            FactorizeError::ComputationError(
6162                anyhow::Error::new(error).context("resident RRQR rank materialization failed"),
6163            )
6164        })?;
6165        #[cfg(feature = "tenferro-cuda")]
6166        let host = match context {
6167            ExecutionContext::Cpu(_) => resident,
6168            ExecutionContext::Cuda(cuda) => cuda.download(&resident).map_err(|error| {
6169                FactorizeError::ComputationError(
6170                    anyhow::Error::new(error).context("resident RRQR rank download failed"),
6171                )
6172            })?,
6173        };
6174        #[cfg(not(feature = "tenferro-cuda"))]
6175        let host = match context {
6176            ExecutionContext::Cpu(_) => resident,
6177        };
6178        let values = host.as_slice::<i64>().map_err(|error| {
6179            FactorizeError::ComputationError(
6180                anyhow::Error::new(error).context("resident RRQR rank decode failed"),
6181            )
6182        })?;
6183        let value = *values.first().ok_or_else(|| {
6184            FactorizeError::ComputationError(anyhow::anyhow!(
6185                "resident RRQR rank metadata is empty"
6186            ))
6187        })?;
6188        usize::try_from(value).map_err(|error| {
6189            FactorizeError::ComputationError(
6190                anyhow::Error::new(error).context("resident RRQR rank is negative"),
6191            )
6192        })
6193    }
6194
6195    /// Build an `n x n` identity in the operand's owning runtime.
6196    ///
6197    /// Host-originated construction data takes the same explicit upload path
6198    /// as [`IdxTensor::from_dense_in`]; the transfer is part of the estimator's
6199    /// decision payload, not hidden arithmetic.
6200    #[cfg(feature = "tenferro-cuda")]
6201    fn resident_identity(
6202        cuda: &Arc<tensor4all_tensorbackend::CudaExecutionContext>,
6203        n: usize,
6204        dtype: DType,
6205    ) -> std::result::Result<EagerTensor, FactorizeError> {
6206        let native = match dtype {
6207            DType::F64 => {
6208                let mut data = vec![0.0_f64; n * n];
6209                for diagonal in 0..n {
6210                    data[diagonal + diagonal * n] = 1.0;
6211                }
6212                NativeTensor::from_vec_col_major(vec![n, n], data)
6213            }
6214            DType::C64 => {
6215                let mut data = vec![Complex64::new(0.0, 0.0); n * n];
6216                for diagonal in 0..n {
6217                    data[diagonal + diagonal * n] = Complex64::new(1.0, 0.0);
6218                }
6219                NativeTensor::from_vec_col_major(vec![n, n], data)
6220            }
6221            _other => {
6222                return Err(FactorizeError::UnsupportedStorage(
6223                    "SRC adaptive estimation currently supports f64 and Complex64 QR factors",
6224                ));
6225            }
6226        }
6227        .map_err(|error| {
6228            FactorizeError::ComputationError(
6229                anyhow::Error::new(error).context("SRC estimator identity construction failed"),
6230            )
6231        })?;
6232        let uploaded = cuda.upload_cuda(&native).map_err(|error| {
6233            FactorizeError::ComputationError(
6234                anyhow::Error::new(error).context("SRC estimator identity upload failed"),
6235            )
6236        })?;
6237        let runtime = cuda.eager_runtime().map_err(|error| {
6238            FactorizeError::ComputationError(
6239                anyhow::Error::new(error).context("SRC estimator eager runtime failed"),
6240            )
6241        })?;
6242        // `validate_context` already pinned the factor to this exact CUDA
6243        // context, so the uploaded identity lands in the factor's runtime.
6244        EagerTensor::from_tensor_in(uploaded, runtime).map_err(|error| {
6245            FactorizeError::ComputationError(
6246                anyhow::Error::new(error).context("SRC estimator identity wrapping failed"),
6247            )
6248        })
6249    }
6250
6251    /// Wrap a resident decision vector and read it back through `context`.
6252    #[cfg(feature = "tenferro-cuda")]
6253    fn read_resident_vector(
6254        eager: &EagerTensor,
6255        len: usize,
6256        context: &ExecutionContext,
6257    ) -> std::result::Result<Vec<f64>, FactorizeError> {
6258        if eager.shape() != [len] {
6259            return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6260                "device SRC decision vector has shape {:?}, expected [{len}]",
6261                eager.shape()
6262            )));
6263        }
6264        let decision = Self::from_inner(vec![DynIndex::new_dyn(len)], eager.clone())
6265            .map_err(FactorizeError::ComputationError)?;
6266        decision
6267            .read_decision_data(context)
6268            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))
6269    }
6270
6271    /// Scalar Appendix-C formula from the `ncols + 1` decision values.
6272    ///
6273    /// Mirrors `src_error_estimate_from_inverse_adjoint`: column norms must be
6274    /// finite and nonzero, and both estimates must be finite.
6275    #[cfg(feature = "tenferro-cuda")]
6276    fn src_estimate_from_decision_scalars(
6277        column_norms_sq: &[f64],
6278        total_norm_sq: f64,
6279        ncols: usize,
6280    ) -> std::result::Result<(f64, f64), anyhow::Error> {
6281        if !total_norm_sq.is_finite() {
6282            return Err(anyhow::anyhow!(
6283                "SRC estimator requires finite entries in R"
6284            ));
6285        }
6286        let mut inverse_column_error_sq = 0.0_f64;
6287        for (col, &column_norm_sq) in column_norms_sq.iter().enumerate() {
6288            if !column_norm_sq.is_finite() || column_norm_sq == 0.0 {
6289                return Err(anyhow::anyhow!(
6290                    "SRC inverse-adjoint solve returned an invalid column norm at column {col}"
6291                ));
6292            }
6293            inverse_column_error_sq += 1.0 / column_norm_sq;
6294        }
6295        let sketch_width = ncols as f64;
6296        let error_sq = inverse_column_error_sq / sketch_width;
6297        let norm_estimate_sq = total_norm_sq / sketch_width;
6298        if !error_sq.is_finite() || !norm_estimate_sq.is_finite() {
6299            return Err(anyhow::anyhow!(
6300                "SRC estimator produced a non-finite estimate"
6301            ));
6302        }
6303        Ok((error_sq.sqrt(), norm_estimate_sq.sqrt()))
6304    }
6305}
6306
6307// Provenance: the public seam mirrors the append contract in
6308// `chriscamano/RandomMPOMPS/code/tensornetwork/incrementalqr.py::IncrementalQR.append`
6309// and `incrementalqr.cpp::add_cols`; the tensor/index bridge and private state
6310// are tensor4all-specific and are labelled `[AI-Supplied]` in the audit.
6311fn factorize_probe_columns_incremental(
6312    previous: Option<&FactorizeResult<IdxTensor>>,
6313    all_columns: &[&IdxTensor],
6314    appended_columns: &[&IdxTensor],
6315    left_inds: &[DynIndex],
6316) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
6317    let first = all_columns.first().ok_or_else(|| {
6318        FactorizeError::ComputationError(anyhow::anyhow!(
6319            "incremental SRC factorization requires at least one column"
6320        ))
6321    })?;
6322    if first.is_f64() {
6323        incremental_probe_factorize_typed::<f64>(previous, all_columns, appended_columns, left_inds)
6324    } else if first.is_c64() {
6325        incremental_probe_factorize_typed::<Complex64>(
6326            previous,
6327            all_columns,
6328            appended_columns,
6329            left_inds,
6330        )
6331    } else {
6332        Err(FactorizeError::UnsupportedStorage(
6333            "incremental SRC factorization currently supports f64 and Complex64 tensors",
6334        ))
6335    }
6336}
6337
6338trait IncrementalQrStateScalar: IncrementalQrScalar + TensorElement {
6339    fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>>;
6340
6341    fn store(state: IncrementalQr<Self>) -> IncrementalQrState;
6342}
6343
6344impl IncrementalQrStateScalar for f64 {
6345    fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>> {
6346        match previous.incremental_qr_state()? {
6347            IncrementalQrState::F64(state) => Some(state.clone()),
6348            IncrementalQrState::C64(_) => None,
6349        }
6350    }
6351
6352    fn store(state: IncrementalQr<Self>) -> IncrementalQrState {
6353        IncrementalQrState::F64(state)
6354    }
6355}
6356
6357impl IncrementalQrStateScalar for Complex64 {
6358    fn resume(previous: &FactorizeResult<IdxTensor>) -> Option<IncrementalQr<Self>> {
6359        match previous.incremental_qr_state()? {
6360            IncrementalQrState::F64(_) => None,
6361            IncrementalQrState::C64(state) => Some(state.clone()),
6362        }
6363    }
6364
6365    fn store(state: IncrementalQr<Self>) -> IncrementalQrState {
6366        IncrementalQrState::C64(state)
6367    }
6368}
6369
6370fn incremental_probe_factorize_from_matrices<S>(
6371    previous: Option<&FactorizeResult<IdxTensor>>,
6372    appended_is_empty: bool,
6373    left_inds: &[DynIndex],
6374    initial_matrix: impl FnOnce() -> std::result::Result<Matrix<S>, FactorizeError>,
6375    appended_matrix: impl FnOnce() -> std::result::Result<Matrix<S>, FactorizeError>,
6376) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6377where
6378    S: IncrementalQrStateScalar,
6379{
6380    let (mut state, previous_left, previous_cap, previous_rank) = if let Some(previous) = previous {
6381        if appended_is_empty {
6382            return Ok(previous.clone());
6383        }
6384        let previous_rank = previous.rank;
6385        let resumed = S::resume(previous);
6386        let reused_previous_q = resumed.is_some();
6387        let state = if let Some(state) = resumed {
6388            state
6389        } else {
6390            let q = previous.left.clone();
6391            let cap = previous.bond_index.clone();
6392            let mut q_indices = left_inds.to_vec();
6393            q_indices.push(cap.clone());
6394            let q = q
6395                .permute_indices(&q_indices)
6396                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6397            let q_dims = q_indices.iter().map(IndexLike::dim).collect::<Vec<_>>();
6398            let q_rows = q_dims[..q_dims.len() - 1]
6399                .iter()
6400                .try_fold(1usize, |size, &dim| size.checked_mul(dim))
6401                .ok_or_else(|| {
6402                    FactorizeError::ComputationError(anyhow::anyhow!(
6403                        "incremental SRC Q row dimension overflows usize"
6404                    ))
6405                })?;
6406            let q_matrix = Matrix::from_col_major_vec(
6407                q_rows,
6408                cap.dim(),
6409                q.to_vec::<S>()
6410                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6411            );
6412
6413            let batch = previous.right.indices().last().cloned().ok_or_else(|| {
6414                FactorizeError::ComputationError(anyhow::anyhow!(
6415                    "incremental SRC R factor has no batch index"
6416                ))
6417            })?;
6418            let r = previous
6419                .right
6420                .permute_indices(&[cap.clone(), batch.clone()])
6421                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6422            let r_matrix = Matrix::from_col_major_vec(
6423                cap.dim(),
6424                batch.dim(),
6425                r.to_vec::<S>()
6426                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6427            );
6428            IncrementalQr::from_factors(q_matrix, r_matrix)
6429                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6430        };
6431        (
6432            state,
6433            if reused_previous_q {
6434                Some(previous.left.clone())
6435            } else {
6436                None
6437            },
6438            if reused_previous_q {
6439                Some(previous.bond_index.clone())
6440            } else {
6441                None
6442            },
6443            previous_rank,
6444        )
6445    } else {
6446        let initial_matrix = initial_matrix()?;
6447        (
6448            IncrementalQr::new(initial_matrix)
6449                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6450            None,
6451            None,
6452            0,
6453        )
6454    };
6455    if !appended_is_empty && previous.is_some() {
6456        let appended_matrix = appended_matrix()?;
6457        state
6458            .append(&appended_matrix)
6459            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6460    }
6461
6462    let rank = state.rank();
6463    if previous_left.is_some() && rank < previous_rank {
6464        return Err(FactorizeError::ComputationError(anyhow::anyhow!(
6465            "incremental SRC QR rank decreased from {previous_rank} to {rank}"
6466        )));
6467    }
6468    let r = state.r();
6469    let cap = if rank == previous_rank {
6470        match previous_cap.as_ref() {
6471            Some(cap) => cap.clone(),
6472            None => DynIndex::new_bond(rank)
6473                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6474        }
6475    } else {
6476        DynIndex::new_bond(rank)
6477            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6478    };
6479    let left = if let Some(previous_left) = previous_left {
6480        if rank == previous_rank {
6481            previous_left
6482        } else {
6483            let appended_rank = rank - previous_rank;
6484            let appended_cap = DynIndex::new_bond(appended_rank)
6485                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6486            let appended_q = state
6487                .q_columns(previous_rank, appended_rank)
6488                .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6489            let mut appended_indices = left_inds.to_vec();
6490            appended_indices.push(appended_cap.clone());
6491            let appended_left =
6492                IdxTensor::from_dense(appended_indices, appended_q.as_col_major_slice().to_vec())
6493                    .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6494            IdxTensor::concatenate_along_new_index(
6495                &[&previous_left, &appended_left],
6496                &[
6497                    previous_cap.ok_or_else(|| {
6498                        FactorizeError::ComputationError(anyhow::anyhow!(
6499                            "incremental SRC previous QR state has no bond index"
6500                        ))
6501                    })?,
6502                    appended_cap,
6503                ],
6504                cap.clone(),
6505            )
6506            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6507        }
6508    } else {
6509        let q = state.q();
6510        let mut q_indices = left_inds.to_vec();
6511        q_indices.push(cap.clone());
6512        IdxTensor::from_dense(q_indices, q.as_col_major_slice().to_vec())
6513            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?
6514    };
6515    let batch = DynIndex::new_link(r.ncols())
6516        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6517    let right = IdxTensor::from_dense(vec![cap.clone(), batch], r.as_col_major_slice().to_vec())
6518        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6519    Ok(FactorizeResult::new(left, right, cap, None, rank)
6520        .with_incremental_qr_state(S::store(state)))
6521}
6522
6523fn incremental_probe_factorize_typed<S>(
6524    previous: Option<&FactorizeResult<IdxTensor>>,
6525    all_columns: &[&IdxTensor],
6526    appended_columns: &[&IdxTensor],
6527    left_inds: &[DynIndex],
6528) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6529where
6530    S: IncrementalQrStateScalar,
6531{
6532    incremental_probe_factorize_from_matrices::<S>(
6533        previous,
6534        appended_columns.is_empty(),
6535        left_inds,
6536        || probe_columns_matrix::<S>(all_columns, left_inds),
6537        || probe_columns_matrix::<S>(appended_columns, left_inds),
6538    )
6539}
6540
6541fn incremental_probe_factorize_batch_typed<S>(
6542    previous: Option<&FactorizeResult<IdxTensor>>,
6543    batch_tensor: &IdxTensor,
6544    batch_index: &DynIndex,
6545    left_inds: &[DynIndex],
6546) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError>
6547where
6548    S: IncrementalQrStateScalar,
6549{
6550    incremental_probe_factorize_from_matrices::<S>(
6551        previous,
6552        batch_index.dim() == 0,
6553        left_inds,
6554        || probe_batch_matrix::<S>(batch_tensor, batch_index, left_inds),
6555        || probe_batch_matrix::<S>(batch_tensor, batch_index, left_inds),
6556    )
6557}
6558
6559fn factorize_probe_batch_incremental_impl(
6560    previous: Option<&FactorizeResult<IdxTensor>>,
6561    batch_tensor: &IdxTensor,
6562    batch_index: &DynIndex,
6563    left_inds: &[DynIndex],
6564) -> std::result::Result<FactorizeResult<IdxTensor>, FactorizeError> {
6565    if batch_tensor.is_f64() {
6566        incremental_probe_factorize_batch_typed::<f64>(
6567            previous,
6568            batch_tensor,
6569            batch_index,
6570            left_inds,
6571        )
6572    } else if batch_tensor.is_c64() {
6573        incremental_probe_factorize_batch_typed::<Complex64>(
6574            previous,
6575            batch_tensor,
6576            batch_index,
6577            left_inds,
6578        )
6579    } else {
6580        Err(FactorizeError::UnsupportedStorage(
6581            "incremental SRC factorization currently supports f64 and Complex64 tensors",
6582        ))
6583    }
6584}
6585
6586fn probe_columns_matrix<S>(
6587    columns: &[&IdxTensor],
6588    left_inds: &[DynIndex],
6589) -> std::result::Result<Matrix<S>, FactorizeError>
6590where
6591    S: TensorElement,
6592{
6593    let batch = DynIndex::new_link(columns.len())
6594        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6595    let stacked = IdxTensor::stack_along_new_index(columns, batch.clone(), -1)
6596        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6597    let mut ordered_indices = left_inds.to_vec();
6598    ordered_indices.push(batch);
6599    let ordered = stacked
6600        .permute_indices(&ordered_indices)
6601        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6602    let nrows = left_inds
6603        .iter()
6604        .try_fold(1usize, |size, index| size.checked_mul(index.dim()))
6605        .ok_or_else(|| {
6606            FactorizeError::ComputationError(anyhow::anyhow!(
6607                "incremental SRC sketch row dimension overflows usize"
6608            ))
6609        })?;
6610    Ok(Matrix::from_col_major_vec(
6611        nrows,
6612        columns.len(),
6613        ordered
6614            .to_vec::<S>()
6615            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6616    ))
6617}
6618
6619/// Like [`probe_columns_matrix`], but for a tensor that already carries a
6620/// batch axis (`batch_index`) instead of being split into separate column
6621/// tensors — skips the `stack_along_new_index` call `probe_columns_matrix`
6622/// needs to re-assemble one.
6623fn probe_batch_matrix<S>(
6624    batch_tensor: &IdxTensor,
6625    batch_index: &DynIndex,
6626    left_inds: &[DynIndex],
6627) -> std::result::Result<Matrix<S>, FactorizeError>
6628where
6629    S: TensorElement,
6630{
6631    let mut ordered_indices = left_inds.to_vec();
6632    ordered_indices.push(batch_index.clone());
6633    let ordered = batch_tensor
6634        .permute_indices(&ordered_indices)
6635        .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?;
6636    let nrows = left_inds
6637        .iter()
6638        .try_fold(1usize, |size, index| size.checked_mul(index.dim()))
6639        .ok_or_else(|| {
6640            FactorizeError::ComputationError(anyhow::anyhow!(
6641                "incremental SRC sketch row dimension overflows usize"
6642            ))
6643        })?;
6644    Ok(Matrix::from_col_major_vec(
6645        nrows,
6646        batch_index.dim(),
6647        ordered
6648            .to_vec::<S>()
6649            .map_err(|error| FactorizeError::ComputationError(anyhow::Error::new(error)))?,
6650    ))
6651}
6652
6653impl TensorContractionLike for IdxTensor {
6654    fn conj(&self) -> Self {
6655        // Delegate to the inherent method (complex conjugate for dense tensors)
6656        IdxTensor::conj(self)
6657    }
6658
6659    fn direct_sum(
6660        &self,
6661        other: &Self,
6662        pairs: &[(DynIndex, DynIndex)],
6663    ) -> std::result::Result<crate::tensor_like::DirectSumResult<Self>, Self::Error> {
6664        let (tensor, new_indices) = crate::direct_sum::direct_sum(self, other, pairs)?;
6665        Ok(crate::tensor_like::DirectSumResult {
6666            tensor,
6667            new_indices,
6668        })
6669    }
6670
6671    fn outer_product(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
6672        super::contract::outer_product(self, other)
6673    }
6674
6675    fn permuteinds(&self, new_order: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
6676        // Delegate to the inherent method
6677        IdxTensor::permute_indices(self, new_order)
6678    }
6679
6680    fn fuse_indices(
6681        &self,
6682        old_indices: &[DynIndex],
6683        new_index: DynIndex,
6684        order: LinearizationOrder,
6685    ) -> std::result::Result<Self, Self::Error> {
6686        IdxTensor::fuse_indices(self, old_indices, new_index, order)
6687    }
6688
6689    fn contract(tensors: &[&Self]) -> std::result::Result<Self, Self::Error> {
6690        super::contract::contract(tensors)
6691    }
6692
6693    fn contract_retaining_indices(
6694        tensors: &[&Self],
6695        retained_indices: &[DynIndex],
6696    ) -> std::result::Result<Self, Self::Error> {
6697        if tensors.len() == 2 {
6698            return tensors[0]
6699                .try_contract_pairwise_retaining(tensors[1], retained_indices)
6700                .map_err(IdxTensorError::from);
6701        }
6702        let options = ContractionOptions::new().with_retain_indices(retained_indices);
6703        super::contract::contract_with_options(tensors, options)
6704    }
6705
6706    fn contract_pair(&self, other: &Self) -> std::result::Result<Self, Self::Error> {
6707        super::contract::contract_pair(self, other)
6708    }
6709}
6710
6711impl TensorConstructionLike for IdxTensor {
6712    fn select_indices(
6713        &self,
6714        selected_indices: &[DynIndex],
6715        positions: &[usize],
6716    ) -> std::result::Result<Self, Self::Error> {
6717        IdxTensor::select_indices(self, selected_indices, positions)
6718    }
6719
6720    fn diagonal(
6721        input_index: &DynIndex,
6722        output_index: &DynIndex,
6723    ) -> std::result::Result<Self, Self::Error> {
6724        let dim = input_index.dim();
6725        if dim != output_index.dim() {
6726            return Err(anyhow::anyhow!(
6727                "Dimension mismatch: input index has dim {}, output has dim {}",
6728                dim,
6729                output_index.dim(),
6730            )
6731            .into());
6732        }
6733
6734        IdxTensor::from_diag(
6735            vec![input_index.clone(), output_index.clone()],
6736            vec![1.0_f64; dim],
6737        )
6738    }
6739
6740    fn scalar_one() -> std::result::Result<Self, Self::Error> {
6741        IdxTensor::from_dense(vec![], vec![1.0_f64])
6742    }
6743
6744    fn ones(indices: &[DynIndex]) -> std::result::Result<Self, Self::Error> {
6745        if indices.is_empty() {
6746            return <Self as TensorConstructionLike>::scalar_one();
6747        }
6748        let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
6749        let total_size = checked_total_size(&dims)?;
6750        IdxTensor::from_dense(indices.to_vec(), vec![1.0_f64; total_size])
6751    }
6752
6753    fn ones_in(
6754        context: &tensor4all_tensorbackend::ExecutionContext,
6755        indices: &[DynIndex],
6756    ) -> std::result::Result<Self, Self::Error> {
6757        IdxTensor::ones_in(context, indices)
6758    }
6759
6760    fn validate_context(
6761        &self,
6762        context: &tensor4all_tensorbackend::ExecutionContext,
6763    ) -> std::result::Result<(), Self::Error> {
6764        IdxTensor::validate_context(self, context)
6765    }
6766
6767    fn from_dense_any(
6768        indices: Vec<DynIndex>,
6769        data: Vec<AnyScalar>,
6770    ) -> std::result::Result<Self, Self::Error> {
6771        IdxTensor::from_dense_any(indices, data)
6772    }
6773
6774    fn from_dense<T: TensorElement + Into<AnyScalar>>(
6775        indices: Vec<DynIndex>,
6776        data: Vec<T>,
6777    ) -> std::result::Result<Self, Self::Error> {
6778        IdxTensor::from_dense(indices, data)
6779    }
6780
6781    fn from_dense_in<T: TensorElement + Into<AnyScalar>>(
6782        context: &tensor4all_tensorbackend::ExecutionContext,
6783        indices: Vec<DynIndex>,
6784        data: Vec<T>,
6785    ) -> std::result::Result<Self, Self::Error> {
6786        IdxTensor::from_dense_in(context, indices, data)
6787    }
6788
6789    fn stack_along_new_index(
6790        tensors: &[&Self],
6791        new_index: DynIndex,
6792        axis: isize,
6793    ) -> std::result::Result<Self, Self::Error> {
6794        IdxTensor::stack_along_new_index(tensors, new_index, axis)
6795    }
6796
6797    fn concatenate_along_new_index(
6798        tensors: &[&Self],
6799        source_indices: &[DynIndex],
6800        new_index: DynIndex,
6801    ) -> std::result::Result<Self, Self::Error> {
6802        let first = tensors
6803            .first()
6804            .ok_or_else(|| anyhow::anyhow!("concatenate requires at least one tensor"))?;
6805        if tensors.len() != source_indices.len() {
6806            return Err(anyhow::anyhow!(
6807                "concatenate tensor count {} does not match source-index count {}",
6808                tensors.len(),
6809                source_indices.len()
6810            )
6811            .into());
6812        }
6813        let first_axis = first
6814            .indices
6815            .iter()
6816            .position(|index| index == &source_indices[0])
6817            .ok_or_else(|| anyhow::anyhow!("concatenate source index is not present"))?;
6818        let mut total_dim = 0usize;
6819        let first_indices = first.indices();
6820        let mut inners = Vec::with_capacity(tensors.len());
6821        for (tensor, source) in tensors.iter().zip(source_indices) {
6822            let axis = tensor
6823                .indices
6824                .iter()
6825                .position(|index| index == source)
6826                .ok_or_else(|| anyhow::anyhow!("concatenate source index is not present"))?;
6827            if axis != first_axis {
6828                return Err(
6829                    anyhow::anyhow!("concatenate source axes must have the same position").into(),
6830                );
6831            }
6832            if tensor
6833                .indices()
6834                .iter()
6835                .zip(first_indices)
6836                .enumerate()
6837                .any(|(position, (actual, expected))| position != first_axis && actual != expected)
6838            {
6839                return Err(anyhow::anyhow!(
6840                    "concatenate tensors must match away from the source axis"
6841                )
6842                .into());
6843            }
6844            total_dim = total_dim
6845                .checked_add(source.dim())
6846                .ok_or_else(|| anyhow::anyhow!("concatenate dimension overflow"))?;
6847            tensor.ensure_shape_packing_preserves_ad("concatenate_along_new_index")?;
6848            inners.push(tensor.try_materialized_inner()?);
6849        }
6850        if new_index.dim() != total_dim {
6851            return Err(anyhow::anyhow!(
6852                "concatenate output dimension {} does not match source dimension sum {}",
6853                new_index.dim(),
6854                total_dim
6855            )
6856            .into());
6857        }
6858        let concatenated = EagerTensor::concatenate(&inners, first_axis)?;
6859        let mut indices = first_indices.to_vec();
6860        indices[first_axis] = new_index;
6861        Self::from_inner(indices, concatenated).map_err(IdxTensorError::from)
6862    }
6863
6864    fn onehot(index_vals: &[(DynIndex, usize)]) -> std::result::Result<Self, Self::Error> {
6865        if index_vals.is_empty() {
6866            return <Self as TensorConstructionLike>::scalar_one();
6867        }
6868        let indices: Vec<DynIndex> = index_vals.iter().map(|(idx, _)| idx.clone()).collect();
6869        let vals: Vec<usize> = index_vals.iter().map(|(_, v)| *v).collect();
6870        let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
6871
6872        for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
6873            if v >= d {
6874                return Err(anyhow::anyhow!(
6875                    "onehot: value {} at position {} is >= dimension {}",
6876                    v,
6877                    k,
6878                    d
6879                )
6880                .into());
6881            }
6882        }
6883
6884        let total_size = checked_total_size(&dims).map_err(Self::Error::from)?;
6885        let mut data = vec![0.0_f64; total_size];
6886
6887        let offset = column_major_offset(&dims, &vals).map_err(Self::Error::from)?;
6888        data[offset] = 1.0;
6889
6890        Self::from_dense(indices, data)
6891    }
6892
6893    // delta() uses the default implementation via diagonal() and outer_product()
6894}
6895
6896fn checked_total_size(dims: &[usize]) -> Result<usize> {
6897    dims.iter().try_fold(1_usize, |acc, &d| {
6898        if d == 0 {
6899            return Err(anyhow::anyhow!("invalid dimension 0"));
6900        }
6901        acc.checked_mul(d)
6902            .ok_or_else(|| anyhow::anyhow!("tensor size overflow"))
6903    })
6904}
6905
6906fn column_major_offset(dims: &[usize], vals: &[usize]) -> Result<usize> {
6907    if dims.len() != vals.len() {
6908        return Err(anyhow::anyhow!(
6909            "column_major_offset: dims.len() != vals.len()"
6910        ));
6911    }
6912    checked_total_size(dims)?;
6913
6914    let mut offset = 0usize;
6915    let mut stride = 1usize;
6916    for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
6917        if d == 0 {
6918            return Err(anyhow::anyhow!("invalid dimension 0 at position {}", k));
6919        }
6920        if v >= d {
6921            return Err(anyhow::anyhow!(
6922                "column_major_offset: value {} at position {} is >= dimension {}",
6923                v,
6924                k,
6925                d
6926            ));
6927        }
6928        let term = v
6929            .checked_mul(stride)
6930            .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6931        offset = offset
6932            .checked_add(term)
6933            .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6934        stride = stride
6935            .checked_mul(d)
6936            .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
6937    }
6938    Ok(offset)
6939}
6940
6941// ============================================================================
6942// High-level API for tensor construction (avoids direct Storage access)
6943// ============================================================================
6944
6945impl IdxTensor {
6946    fn any_scalar_payload_to_complex(data: Vec<AnyScalar>) -> Vec<Complex64> {
6947        data.into_iter()
6948            .map(|value| {
6949                value
6950                    .as_c64()
6951                    .unwrap_or_else(|| Complex64::new(value.real(), 0.0))
6952            })
6953            .collect()
6954    }
6955
6956    fn any_scalar_payload_to_real(data: Vec<AnyScalar>) -> Vec<f64> {
6957        data.into_iter().map(|value| value.real()).collect()
6958    }
6959
6960    fn validate_dense_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
6961        let expected_len = checked_total_size(dims)?;
6962        if !(data_len == expected_len) {
6963            return Err(anyhow::anyhow!(
6964                "dense payload length {} does not match dims {:?} (expected {})",
6965                data_len,
6966                dims,
6967                expected_len
6968            ));
6969        };
6970        Ok(())
6971    }
6972
6973    fn validate_diag_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
6974        if !(!dims.is_empty()) {
6975            return Err(anyhow::anyhow!(
6976                "diagonal tensor construction requires at least one index"
6977            ));
6978        };
6979        Self::validate_diag_dims(dims)?;
6980        if !(data_len == dims[0]) {
6981            return Err(anyhow::anyhow!(
6982                "diagonal payload length {} does not match diagonal dimension {}",
6983                data_len,
6984                dims[0]
6985            ));
6986        };
6987        Ok(())
6988    }
6989
6990    /// Create a tensor from dense data with explicit indices.
6991    ///
6992    /// This is the recommended high-level API for creating tensors from raw data.
6993    /// It avoids direct access to `Storage` internals.
6994    ///
6995    /// # Type Parameters
6996    /// * `T` - Scalar type (`f32`, `f64`, `Complex32`, or `Complex64`)
6997    ///
6998    /// # Arguments
6999    /// * `indices` - Vector of indices for the tensor
7000    /// * `data` - Tensor data in column-major order
7001    ///
7002    /// # Errors
7003    /// Returns an error when the data length does not match the index dimension
7004    /// product (a shape mismatch).
7005    /// # Example
7006    /// ```
7007    /// use tensor4all_core::IdxTensor;
7008    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7009    ///
7010    /// let i = Index::new_dyn(2);
7011    /// let j = Index::new_dyn(3);
7012    /// let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
7013    /// let tensor: IdxTensor = IdxTensor::from_dense(vec![i, j], data).unwrap();
7014    /// assert_eq!(tensor.dims(), vec![2, 3]);
7015    /// ```
7016    pub fn from_dense<T: TensorElement>(
7017        indices: Vec<DynIndex>,
7018        data: Vec<T>,
7019    ) -> std::result::Result<Self, IdxTensorError> {
7020        let dims = Self::expected_dims_from_indices(&indices);
7021        Self::validate_indices(&indices)?;
7022        Self::validate_dense_payload_len(data.len(), &dims)?;
7023        let native = dense_native_tensor_from_col_major_owned(data, &dims)?;
7024        Self::from_native(indices, native).map_err(IdxTensorError::from)
7025    }
7026
7027    /// Create a tensor from dense payload data provided as [`AnyScalar`] values.
7028    ///
7029    /// This is the preferred public API when the caller only knows the scalar
7030    /// type at runtime.
7031    ///
7032    /// # Errors
7033    /// Returns an error when the payload length does not match the index dimension
7034    /// product (a shape mismatch) or a scalar conversion fails.
7035    /// # Examples
7036    /// ```
7037    /// use tensor4all_core::{AnyScalar, IdxTensor};
7038    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7039    ///
7040    /// let i = Index::new_dyn(2);
7041    /// let j = Index::new_dyn(2);
7042    /// let tensor = IdxTensor::from_dense_any(
7043    ///     vec![i, j],
7044    ///     vec![
7045    ///         AnyScalar::new_real(1.0),
7046    ///         AnyScalar::new_complex(0.0, 1.0),
7047    ///         AnyScalar::new_real(2.0),
7048    ///         AnyScalar::new_real(3.0),
7049    ///     ],
7050    /// ).unwrap();
7051    ///
7052    /// assert!(tensor.is_complex());
7053    /// assert_eq!(tensor.dims(), vec![2, 2]);
7054    /// ```
7055    pub fn from_dense_any(
7056        indices: Vec<DynIndex>,
7057        data: Vec<AnyScalar>,
7058    ) -> std::result::Result<Self, IdxTensorError> {
7059        if data.iter().any(AnyScalar::is_complex) {
7060            Self::from_dense(indices, Self::any_scalar_payload_to_complex(data))
7061        } else {
7062            Self::from_dense(indices, Self::any_scalar_payload_to_real(data))
7063        }
7064    }
7065
7066    /// Create a diagonal tensor from diagonal payload data with explicit indices.
7067    ///
7068    /// All indices must have the same dimension, and `data.len()` must equal
7069    /// that dimension. The resulting tensor has nonzero entries only on
7070    /// the multi-index diagonal (`T[i,i,...,i] = data[i]`).
7071    ///
7072    /// The returned tensor preserves diagonal metadata; use
7073    /// [`IdxTensor::is_diag`] or [`IdxTensor::storage_kind`] to inspect
7074    /// that representation. `f32` and `Complex32` values remain eager and are
7075    /// never promoted into compact `f64`/`Complex64` storage.
7076    ///
7077    /// # Errors
7078    /// Returns an error when the index dimensions are unequal or the payload
7079    /// length does not match the diagonal dimension (a shape mismatch).
7080    /// # Examples
7081    ///
7082    /// ```
7083    /// use tensor4all_core::{DynIndex, IdxTensor};
7084    ///
7085    /// let i = DynIndex::new_dyn(3);
7086    /// let j = DynIndex::new_dyn(3);
7087    /// let diag = IdxTensor::from_diag(vec![i, j], vec![1.0, 2.0, 3.0]).unwrap();
7088    /// assert!(diag.is_diag());
7089    ///
7090    /// let data = diag.to_vec::<f64>().unwrap();
7091    /// // 3x3 identity-like: [1,0,0, 0,2,0, 0,0,3] in column-major
7092    /// assert!((data[0] - 1.0).abs() < 1e-12);
7093    /// assert!((data[4] - 2.0).abs() < 1e-12);
7094    /// assert!((data[8] - 3.0).abs() < 1e-12);
7095    /// assert!((data[1]).abs() < 1e-12);  // off-diagonal is zero
7096    /// ```
7097    pub fn from_diag<T: TensorElement>(
7098        indices: Vec<DynIndex>,
7099        data: Vec<T>,
7100    ) -> std::result::Result<Self, IdxTensorError> {
7101        let dims = Self::expected_dims_from_indices(&indices);
7102        Self::validate_indices(&indices)?;
7103        Self::validate_diag_payload_len(data.len(), &dims)?;
7104        let native = diag_native_tensor_from_col_major(&data, dims.len())?;
7105        Self::from_native_with_axis_classes(indices, native, Self::diag_axis_classes(dims.len()))
7106            .map_err(IdxTensorError::from)
7107    }
7108
7109    /// Create a diagonal tensor from diagonal payload data provided as
7110    /// [`AnyScalar`] values.
7111    ///
7112    /// This is the preferred public API when the caller only knows the scalar
7113    /// type at runtime.
7114    ///
7115    /// # Errors
7116    /// Returns an error when the index dimensions are unequal or the payload
7117    /// length does not match (a shape mismatch), or a scalar conversion
7118    /// fails.
7119    /// # Examples
7120    /// ```
7121    /// use tensor4all_core::{AnyScalar, IdxTensor};
7122    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7123    ///
7124    /// let i = Index::new_dyn(2);
7125    /// let j = Index::new_dyn(2);
7126    /// let tensor = IdxTensor::from_diag_any(
7127    ///     vec![i, j],
7128    ///     vec![AnyScalar::new_real(1.0), AnyScalar::new_complex(2.0, -1.0)],
7129    /// ).unwrap();
7130    ///
7131    /// assert!(tensor.is_complex());
7132    /// assert_eq!(tensor.dims(), vec![2, 2]);
7133    /// ```
7134    pub fn from_diag_any(
7135        indices: Vec<DynIndex>,
7136        data: Vec<AnyScalar>,
7137    ) -> std::result::Result<Self, IdxTensorError> {
7138        if data.iter().any(AnyScalar::is_complex) {
7139            Self::from_diag(indices, Self::any_scalar_payload_to_complex(data))
7140        } else {
7141            Self::from_diag(indices, Self::any_scalar_payload_to_real(data))
7142        }
7143    }
7144
7145    /// Create a copy tensor whose nonzero entries are `value` on the diagonal.
7146    ///
7147    /// For indices `[i, j, k]`, the returned tensor satisfies
7148    /// `T[i, j, k] = value` when `i = j = k`, and zero otherwise.
7149    ///
7150    /// # Errors
7151    /// Returns an error when the index dimensions are unequal (a dimension
7152    /// mismatch) or the construction fails.
7153    /// # Examples
7154    /// ```
7155    /// use tensor4all_core::{AnyScalar, IdxTensor};
7156    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7157    ///
7158    /// let i = Index::new_dyn(2);
7159    /// let j = Index::new_dyn(2);
7160    /// let k = Index::new_dyn(2);
7161    /// let tensor = IdxTensor::copy_tensor(
7162    ///     vec![i, j, k],
7163    ///     AnyScalar::new_real(1.0),
7164    /// ).unwrap();
7165    ///
7166    /// assert_eq!(tensor.dims(), vec![2, 2, 2]);
7167    /// ```
7168    pub fn copy_tensor(
7169        indices: Vec<DynIndex>,
7170        value: AnyScalar,
7171    ) -> std::result::Result<Self, IdxTensorError> {
7172        if indices.is_empty() {
7173            return Self::from_dense_any(vec![], vec![value]);
7174        }
7175        let dim = indices[0].dim();
7176        let data = vec![value; dim];
7177        Self::from_diag_any(indices, data)
7178    }
7179
7180    /// Replace multiple tensor indices with one fused index using an exact local reshape.
7181    ///
7182    /// The full indices in `old_indices` identify the axes to fuse and also
7183    /// define the coordinate order used inside `new_index`. The new fused index
7184    /// is inserted at the earliest axis position among the fused axes; all
7185    /// other axes keep their original relative order. Use
7186    /// [`LinearizationOrder::ColumnMajor`] to match tensor4all's dense vector
7187    /// layout, or [`LinearizationOrder::RowMajor`] when interoperating with
7188    /// row-major fused coordinates.
7189    ///
7190    /// # Arguments
7191    /// * `old_indices` - Non-empty list of existing tensor indices to replace.
7192    ///
7193    ///   Each index is matched by full identity, must appear exactly once in
7194    ///   the tensor, must have the same dimension as the matched tensor axis,
7195    ///   and must not be duplicated in this list.
7196    /// * `new_index` - Replacement index whose dimension must equal the product
7197    ///
7198    ///   of the dimensions in `old_indices`.
7199    /// * `order` - Linearization convention used to encode the old coordinates
7200    ///
7201    ///   into the single coordinate of `new_index`.
7202    ///
7203    /// # Returns
7204    /// A tensor with the same element type and values, but with `old_indices`
7205    /// replaced by `new_index`.
7206    ///
7207    /// # Errors
7208    /// Returns an error if `old_indices` is empty, contains duplicate IDs,
7209    /// references an index not present in the tensor, if the fused dimension
7210    /// does not match the product of the old dimensions, if the replacement
7211    /// would duplicate a kept index, or if the dense reshape cannot be
7212    /// represented without overflow.
7213    ///
7214    /// # Examples
7215    /// ```
7216    /// use tensor4all_core::{DynIndex, LinearizationOrder, IdxTensor};
7217    ///
7218    /// let i = DynIndex::new_dyn(2);
7219    /// let j = DynIndex::new_dyn(2);
7220    /// let fused = DynIndex::new_link(4).unwrap();
7221    /// let tensor = IdxTensor::from_dense(
7222    ///     vec![i.clone(), j.clone()],
7223    ///     vec![1.0, 2.0, 3.0, 4.0],
7224    /// ).unwrap();
7225    ///
7226    /// let fused_tensor = tensor
7227    ///     .fuse_indices(&[i.clone(), j.clone()], fused.clone(), LinearizationOrder::ColumnMajor)
7228    ///     .unwrap();
7229    /// assert_eq!(fused_tensor.dims(), vec![4]);
7230    ///
7231    /// let roundtrip = fused_tensor
7232    ///     .unfuse_index(&fused, &[i, j], LinearizationOrder::ColumnMajor)
7233    ///     .unwrap();
7234    /// assert!(roundtrip.isapprox(&tensor, 1e-12, 0.0).unwrap());
7235    /// ```
7236    pub fn fuse_indices(
7237        &self,
7238        old_indices: &[DynIndex],
7239        new_index: DynIndex,
7240        order: LinearizationOrder,
7241    ) -> std::result::Result<Self, IdxTensorError> {
7242        if !(!old_indices.is_empty()) {
7243            return Err(anyhow::anyhow!("fuse_indices requires at least one index to fuse").into());
7244        };
7245
7246        let old_dims = self.dims();
7247        let mut seen_indices = HashSet::new();
7248        let mut old_axes = Vec::with_capacity(old_indices.len());
7249        for old_index in old_indices {
7250            if !(seen_indices.insert(old_index)) {
7251                return Err(anyhow::anyhow!("duplicate index in old_indices").into());
7252            };
7253            let axis = self
7254                .indices
7255                .iter()
7256                .position(|idx| idx == old_index)
7257                .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
7258            if !(old_index.dim() == old_dims[axis]) {
7259                return Err(anyhow::anyhow!(
7260                    "old index dimension does not match tensor axis dimension"
7261                )
7262                .into());
7263            };
7264            old_axes.push(axis);
7265        }
7266
7267        let fused_dims: Vec<usize> = old_axes.iter().map(|&axis| old_dims[axis]).collect();
7268        let fused_product = checked_product(&fused_dims)?;
7269        if !(fused_product == new_index.dim()) {
7270            return Err(anyhow::anyhow!(
7271                "product of old index dimensions must match the replacement index dimension"
7272            )
7273            .into());
7274        };
7275
7276        let insertion_axis =
7277            old_axes.iter().copied().min().ok_or_else(|| {
7278                anyhow::anyhow!("fuse_indices requires at least one index to fuse")
7279            })?;
7280        let old_axis_set: HashSet<usize> = old_axes.iter().copied().collect();
7281
7282        let mut result_indices =
7283            Vec::with_capacity(self.indices.len() - old_indices.len() + 1usize);
7284        for (axis, index) in self.indices.iter().enumerate() {
7285            if axis == insertion_axis {
7286                result_indices.push(new_index.clone());
7287            }
7288            if !old_axis_set.contains(&axis) {
7289                result_indices.push(index.clone());
7290            }
7291        }
7292        let mut result_seen = HashSet::new();
7293        for index in &result_indices {
7294            if !(result_seen.insert(index)) {
7295                return Err(
7296                    anyhow::anyhow!("fuse_indices result would contain duplicate index").into(),
7297                );
7298            };
7299        }
7300        Self::validate_indices(&result_indices)?;
7301
7302        let mut new_dims = Vec::with_capacity(old_dims.len() - old_indices.len() + 1usize);
7303        for (axis, dim) in old_dims.iter().copied().enumerate() {
7304            if axis == insertion_axis {
7305                new_dims.push(new_index.dim());
7306            }
7307            if !old_axis_set.contains(&axis) {
7308                new_dims.push(dim);
7309            }
7310        }
7311
7312        self.ensure_shape_packing_preserves_ad("fuse_indices")?;
7313
7314        let mut grouped_axes = old_axes.clone();
7315        if matches!(order, LinearizationOrder::RowMajor) {
7316            grouped_axes.reverse();
7317        }
7318        let mut perm = Vec::with_capacity(self.indices.len());
7319        perm.extend((0..insertion_axis).filter(|axis| !old_axis_set.contains(axis)));
7320        perm.extend(grouped_axes);
7321        perm.extend(
7322            ((insertion_axis + 1)..self.indices.len()).filter(|axis| !old_axis_set.contains(axis)),
7323        );
7324        debug_assert_eq!(perm.len(), self.indices.len());
7325
7326        let packed = self.permute(&perm)?;
7327        let reshaped = packed.try_materialized_inner()?.reshape(&new_dims)?;
7328        Self::from_inner(result_indices, reshaped).map_err(IdxTensorError::from)
7329    }
7330
7331    /// Replace one fused index with multiple indices using an exact reshape.
7332    ///
7333    /// The caller must specify how the old fused index should be decoded into
7334    /// the new indices via `order`.
7335    ///
7336    /// # Errors
7337    /// Returns an error when the fused dimension does not equal the product of
7338    /// the new index dimensions (a shape mismatch).
7339    /// # Examples
7340    /// ```
7341    /// use tensor4all_core::{DynIndex, LinearizationOrder, IdxTensor};
7342    ///
7343    /// let fused = DynIndex::new_dyn(4);
7344    /// let i = DynIndex::new_dyn(2);
7345    /// let j = DynIndex::new_dyn(2);
7346    /// let tensor = IdxTensor::from_dense(vec![fused.clone()], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
7347    ///
7348    /// let unfused = tensor
7349    ///     .unfuse_index(&fused, &[i.clone(), j.clone()], LinearizationOrder::ColumnMajor)
7350    ///     .unwrap();
7351    ///
7352    /// let expected = IdxTensor::from_dense(vec![i, j], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
7353    /// assert!(unfused.isapprox(&expected, 1e-12, 0.0).unwrap());
7354    /// ```
7355    pub fn unfuse_index(
7356        &self,
7357        old_index: &DynIndex,
7358        new_indices: &[DynIndex],
7359        order: LinearizationOrder,
7360    ) -> std::result::Result<Self, IdxTensorError> {
7361        if !(!new_indices.is_empty()) {
7362            return Err(
7363                anyhow::anyhow!("unfuse_index requires at least one replacement index").into(),
7364            );
7365        };
7366
7367        let axis = self
7368            .indices
7369            .iter()
7370            .position(|idx| idx == old_index)
7371            .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
7372
7373        let replacement_dims: Vec<usize> = new_indices.iter().map(DynIndex::dim).collect();
7374        let replacement_product = checked_product(&replacement_dims)?;
7375        if !(replacement_product == old_index.dim()) {
7376            return Err(anyhow::anyhow!(
7377                "product of new index dimensions must match the replaced index dimension"
7378            )
7379            .into());
7380        };
7381
7382        let mut result_indices =
7383            Vec::with_capacity(self.indices.len() - 1usize + new_indices.len());
7384        result_indices.extend_from_slice(&self.indices[..axis]);
7385        result_indices.extend(new_indices.iter().cloned());
7386        result_indices.extend_from_slice(&self.indices[axis + 1..]);
7387        Self::validate_indices(&result_indices)?;
7388
7389        let old_dims = self.dims();
7390        let mut new_dims = Vec::with_capacity(old_dims.len() - 1usize + replacement_dims.len());
7391        new_dims.extend_from_slice(&old_dims[..axis]);
7392        new_dims.extend_from_slice(&replacement_dims);
7393        new_dims.extend_from_slice(&old_dims[axis + 1..]);
7394
7395        self.ensure_shape_packing_preserves_ad("unfuse_index")?;
7396
7397        let mut grouped_indices = new_indices.to_vec();
7398        let mut grouped_dims = replacement_dims.clone();
7399        if matches!(order, LinearizationOrder::RowMajor) {
7400            grouped_indices.reverse();
7401            grouped_dims.reverse();
7402        }
7403        let mut packed_indices =
7404            Vec::with_capacity(self.indices.len() - 1usize + grouped_indices.len());
7405        packed_indices.extend_from_slice(&self.indices[..axis]);
7406        packed_indices.extend(grouped_indices);
7407        packed_indices.extend_from_slice(&self.indices[axis + 1..]);
7408
7409        let mut packed_dims = Vec::with_capacity(old_dims.len() - 1usize + grouped_dims.len());
7410        packed_dims.extend_from_slice(&old_dims[..axis]);
7411        packed_dims.extend_from_slice(&grouped_dims);
7412        packed_dims.extend_from_slice(&old_dims[axis + 1..]);
7413
7414        let reshaped = self.try_materialized_inner()?.reshape(&packed_dims)?;
7415        let packed = Self::from_inner(packed_indices, reshaped)?;
7416        if matches!(order, LinearizationOrder::ColumnMajor) {
7417            Ok(packed)
7418        } else {
7419            packed.permute_indices(&result_indices)
7420        }
7421    }
7422
7423    /// Create a scalar (0-dimensional) tensor from a supported element value.
7424    ///
7425    /// # Errors
7426    /// Returns an error when the element type is not supported (an
7427    /// unsupported-dtype failure).
7428    /// # Example
7429    /// ```
7430    /// use tensor4all_core::IdxTensor;
7431    ///
7432    /// let scalar = IdxTensor::scalar(42.0).unwrap();
7433    /// assert_eq!(scalar.dims(), Vec::<usize>::new());
7434    /// assert_eq!(scalar.only().unwrap().real(), 42.0);
7435    /// ```
7436    pub fn scalar<T: TensorElement>(value: T) -> std::result::Result<Self, IdxTensorError> {
7437        Self::from_dense(vec![], vec![value])
7438    }
7439
7440    /// Create a tensor filled with zeros of a supported element type.
7441    ///
7442    /// # Errors
7443    /// Returns an error when the dimension product overflows (an overflow failure)
7444    /// or the element type is unsupported.
7445    /// # Example
7446    /// ```
7447    /// use tensor4all_core::IdxTensor;
7448    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7449    ///
7450    /// let i = Index::new_dyn(2);
7451    /// let j = Index::new_dyn(3);
7452    /// let tensor = IdxTensor::zeros::<f64>(vec![i, j]).unwrap();
7453    /// assert_eq!(tensor.dims(), vec![2, 3]);
7454    /// ```
7455    pub fn zeros<T: TensorElement + Zero + Clone>(
7456        indices: Vec<DynIndex>,
7457    ) -> std::result::Result<Self, IdxTensorError> {
7458        let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
7459        let size = checked_product(&dims)?;
7460        Self::from_dense(indices, vec![T::zero(); size])
7461    }
7462}
7463
7464// ============================================================================
7465// High-level API for data extraction (avoids direct .storage() access)
7466// ============================================================================
7467
7468impl IdxTensor {
7469    /// Extract tensor data as a column-major `Vec<T>`.
7470    ///
7471    /// # Type Parameters
7472    /// * `T` - The scalar element type (`f32`, `f64`, `Complex32`, or
7473    ///
7474    ///   `Complex64`).
7475    ///
7476    /// # Returns
7477    /// A vector of the tensor data in column-major order.
7478    ///
7479    /// # Errors
7480    /// Returns an error when the tensor dtype does not match the requested element
7481    /// type (a scalar-kind mismatch) or materialization fails.
7482    /// # Example
7483    /// ```
7484    /// use tensor4all_core::IdxTensor;
7485    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7486    ///
7487    /// let i = Index::new_dyn(2);
7488    /// let tensor = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
7489    /// let data = tensor.to_vec::<f64>().unwrap();
7490    /// assert_eq!(data, &[1.0, 2.0]);
7491    /// ```
7492    pub fn to_vec<T: TensorElement>(&self) -> std::result::Result<Vec<T>, IdxTensorError> {
7493        self.as_inner()?
7494            .duplicate_value()?
7495            .as_slice::<T>()
7496            .map(|values| values.to_vec())
7497            .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))
7498    }
7499
7500    /// Reads dense column-major tensor values without copying them into a new vector.
7501    ///
7502    /// The callback receives the same logical dense values and ordering as
7503    /// [`Self::to_vec`], with the first tensor index varying fastest. Use this
7504    /// for read-only kernels that can finish while the callback is active;
7505    /// use [`Self::to_vec`] when the values must outlive the call.
7506    ///
7507    /// # Arguments
7508    ///
7509    /// * `read` - A callback that consumes the temporary dense value slice and
7510    ///   returns the caller's result.
7511    ///
7512    /// # Returns
7513    ///
7514    /// Returns the callback's result without allocating a result vector.
7515    ///
7516    /// # Errors
7517    ///
7518    /// Returns an error when the tensor dtype does not match `T` or dense
7519    /// materialization fails.
7520    ///
7521    /// # Examples
7522    ///
7523    /// ```
7524    /// use tensor4all_core::{DynIndex, IdxTensor};
7525    ///
7526    /// let i = DynIndex::new_dyn(2);
7527    /// let j = DynIndex::new_dyn(2);
7528    /// let tensor = IdxTensor::from_dense(
7529    ///     vec![i, j],
7530    ///     vec![1.0_f64, 2.0, 3.0, 4.0],
7531    /// )?;
7532    /// let weighted_sum = tensor.with_dense_slice::<f64, _>(|values| {
7533    ///     values.iter().enumerate().map(|(i, value)| (i + 1) as f64 * value).sum::<f64>()
7534    /// })?;
7535    /// assert_eq!(weighted_sum, 30.0);
7536    /// # Ok::<(), anyhow::Error>(())
7537    /// ```
7538    pub fn with_dense_slice<T: TensorElement, R>(
7539        &self,
7540        read: impl FnOnce(&[T]) -> R,
7541    ) -> std::result::Result<R, IdxTensorError> {
7542        let inner = self.as_inner()?;
7543        let value = inner.value()?;
7544        if let Ok(values) = value.as_slice::<T>() {
7545            return Ok(read(values));
7546        }
7547
7548        // Backend-resident and non-contiguous values cannot be borrowed as a
7549        // host slice. Preserve the previous materializing behavior for those
7550        // cases while avoiding a full payload copy for the ordinary
7551        // host-contiguous path above.
7552        drop(value);
7553        let values = inner.duplicate_value()?;
7554        let values = values
7555            .as_slice::<T>()
7556            .map_err(|source| IdxTensorError::materialization(anyhow::Error::new(source)))?;
7557        Ok(read(values))
7558    }
7559
7560    /// Consume the tensor and return its indices with dense column-major values.
7561    ///
7562    /// Use this when a caller needs to move index metadata and dense payload
7563    /// values across an API boundary. The returned values are ordered with the
7564    /// first tensor index varying fastest. Compact diagonal or structured
7565    /// storage is materialized into dense logical values.
7566    ///
7567    /// # Type Parameters
7568    /// * `T` - The scalar element type to extract: `f32`, `f64`, `Complex32`,
7569    ///
7570    ///   or `Complex64`.
7571    ///
7572    /// # Returns
7573    /// The tensor's original indices and dense column-major flat data.
7574    ///
7575    /// # Errors
7576    /// Returns an error when the tensor dtype does not match the requested element
7577    /// type (a scalar-kind mismatch) or materialization fails.
7578    /// # Examples
7579    /// ```
7580    /// use tensor4all_core::{DynIndex, IdxTensor};
7581    ///
7582    /// let i = DynIndex::new_dyn(2);
7583    /// let j = DynIndex::new_dyn(2);
7584    /// let tensor = IdxTensor::from_dense(
7585    ///     vec![i.clone(), j.clone()],
7586    ///     vec![1.0_f64, 2.0, 3.0, 4.0],
7587    /// ).unwrap();
7588    ///
7589    /// let (indices, data) = tensor.into_dense_col_major_parts::<f64>().unwrap();
7590    ///
7591    /// assert_eq!(indices, vec![i, j]);
7592    /// assert_eq!(data, vec![1.0, 2.0, 3.0, 4.0]);
7593    /// ```
7594    pub fn into_dense_col_major_parts<T: TensorElement>(
7595        self,
7596    ) -> std::result::Result<(Vec<DynIndex>, Vec<T>), IdxTensorError> {
7597        if !(!self.tracks_grad()) {
7598            return Err(anyhow::anyhow!("IdxTensor::into_dense_col_major_parts cannot consume tensors with tracked autodiff state").into());
7599        };
7600        let data = self.to_vec::<T>()?;
7601        Ok((self.indices, data))
7602    }
7603
7604    /// Check if the tensor has `f64` storage.
7605    ///
7606    /// # Example
7607    /// ```
7608    /// use tensor4all_core::IdxTensor;
7609    /// use tensor4all_core::index::{DefaultIndex as Index, DynId};
7610    ///
7611    /// let i = Index::new_dyn(2);
7612    /// let tensor = IdxTensor::from_dense(vec![i], vec![1.0, 2.0]).unwrap();
7613    /// assert!(tensor.is_f64());
7614    /// assert!(!tensor.is_complex());
7615    /// ```
7616    pub fn is_f64(&self) -> bool {
7617        self.storage.dtype() == Some(DType::F64)
7618    }
7619
7620    /// Check if the tensor has `f32` storage.
7621    ///
7622    /// # Examples
7623    ///
7624    /// ```
7625    /// use tensor4all_core::{DynIndex, IdxTensor};
7626    ///
7627    /// let tensor = IdxTensor::from_dense(
7628    ///     vec![DynIndex::new_dyn(2)],
7629    ///     vec![1.0_f32, 2.0],
7630    /// )
7631    /// .unwrap();
7632    /// assert!(tensor.is_f32());
7633    /// ```
7634    pub fn is_f32(&self) -> bool {
7635        self.storage.dtype() == Some(DType::F32)
7636    }
7637
7638    /// Check if the tensor has complex-32 storage.
7639    ///
7640    /// # Examples
7641    ///
7642    /// ```
7643    /// use num_complex::Complex32;
7644    /// use tensor4all_core::{DynIndex, IdxTensor};
7645    ///
7646    /// let tensor = IdxTensor::from_dense(
7647    ///     vec![DynIndex::new_dyn(2)],
7648    ///     vec![Complex32::new(1.0, 0.0), Complex32::new(0.0, 1.0)],
7649    /// )
7650    /// .unwrap();
7651    /// assert!(tensor.is_c32());
7652    /// ```
7653    pub fn is_c32(&self) -> bool {
7654        self.storage.dtype() == Some(DType::C32)
7655    }
7656
7657    /// Check if the tensor has complex-64 storage.
7658    ///
7659    /// # Example
7660    /// ```
7661    /// use num_complex::Complex64;
7662    /// use tensor4all_core::{DynIndex, IdxTensor};
7663    ///
7664    /// let i = DynIndex::new_dyn(2);
7665    /// let tensor = IdxTensor::from_dense(
7666    ///     vec![i],
7667    ///     vec![Complex64::new(1.0, 0.0), Complex64::new(0.0, 1.0)],
7668    /// )
7669    /// .unwrap();
7670    /// assert!(tensor.is_c64());
7671    /// ```
7672    pub fn is_c64(&self) -> bool {
7673        self.storage.is_c64()
7674    }
7675
7676    /// Check whether the tensor carries diagonal logical axis metadata.
7677    ///
7678    /// # Examples
7679    ///
7680    /// ```
7681    /// use tensor4all_core::{DynIndex, IdxTensor};
7682    /// use tensor4all_tensorbackend::Storage;
7683    ///
7684    /// // Tensors from `from_dense` use dense storage
7685    /// let i = DynIndex::new_dyn(2);
7686    /// let j = DynIndex::new_dyn(2);
7687    /// let dense = IdxTensor::from_dense(vec![i, j], vec![1.0, 0.0, 0.0, 1.0]).unwrap();
7688    /// assert!(!dense.is_diag());
7689    ///
7690    /// // Diagonal metadata is preserved when constructing from diagonal storage.
7691    /// let k = DynIndex::new_dyn(2);
7692    /// let l = DynIndex::new_dyn(2);
7693    /// let diag = IdxTensor::from_storage(
7694    ///     vec![k, l],
7695    ///     Storage::from_diag_col_major(vec![1.0, 2.0], 2)
7696    ///         .map(std::sync::Arc::new)
7697    ///         .unwrap(),
7698    /// )
7699    /// .unwrap();
7700    /// assert!(diag.is_diag());
7701    /// ```
7702    pub fn is_diag(&self) -> bool {
7703        self.storage.is_diag()
7704    }
7705
7706    /// Check if the tensor has complex storage (C64).
7707    ///
7708    /// # Examples
7709    ///
7710    /// ```
7711    /// use tensor4all_core::{DynIndex, IdxTensor};
7712    /// use num_complex::Complex64;
7713    ///
7714    /// let i = DynIndex::new_dyn(2);
7715    /// let real_t = IdxTensor::from_dense(vec![i.clone()], vec![1.0, 2.0]).unwrap();
7716    /// assert!(!real_t.is_complex());
7717    ///
7718    /// let complex_t = IdxTensor::from_dense(
7719    ///     vec![i],
7720    ///     vec![Complex64::new(1.0, 0.0), Complex64::new(0.0, 1.0)],
7721    /// ).unwrap();
7722    /// assert!(complex_t.is_complex());
7723    /// ```
7724    pub fn is_complex(&self) -> bool {
7725        self.storage.is_complex()
7726    }
7727    /// Create a dense tensor in a caller-owned execution context.
7728    ///
7729    /// Host data is explicitly uploaded when `context` is CUDA-resident; no
7730    /// global default context is consulted.
7731    ///
7732    /// # Examples
7733    ///
7734    /// ```
7735    /// use std::sync::Arc;
7736    /// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor};
7737    /// use tensor4all_tensorbackend::CpuExecutionContext;
7738    /// use tenferro_cpu::CpuBackend;
7739    ///
7740    /// let context = ExecutionContext::Cpu(Arc::new(
7741    ///     CpuExecutionContext::from_backend(CpuBackend::new()),
7742    /// ));
7743    /// let index = DynIndex::new_dyn(2);
7744    /// let tensor = IdxTensor::from_dense_in(&context, vec![index], vec![1.0_f64, 2.0])?;
7745    /// assert_eq!(tensor.to_vec::<f64>()?, vec![1.0, 2.0]);
7746    /// # Ok::<(), Box<dyn std::error::Error>>(())
7747    /// ```
7748    ///
7749    /// # Errors
7750    ///
7751    /// Returns [`IdxTensorError`] when the indices, payload, context, or
7752    /// explicit transfer is invalid.
7753    pub fn from_dense_in<T: TensorElement>(
7754        context: &ExecutionContext,
7755        indices: Vec<DynIndex>,
7756        data: Vec<T>,
7757    ) -> std::result::Result<Self, IdxTensorError> {
7758        let dims = Self::expected_dims_from_indices(&indices);
7759        Self::validate_indices(&indices)?;
7760        Self::validate_dense_payload_len(data.len(), &dims)?;
7761        let native = dense_native_tensor_from_col_major(&data, &dims)
7762            .map_err(|error| IdxTensorError::operation("context-scoped construction", error))?;
7763        let inner = match context {
7764            ExecutionContext::Cpu(context) => {
7765                let runtime = context.eager_runtime().map_err(|error| {
7766                    IdxTensorError::operation("CPU context construction", anyhow::Error::new(error))
7767                })?;
7768                EagerTensor::from_tensor_in(native, runtime).map_err(|error| {
7769                    IdxTensorError::operation(
7770                        "CPU context eager wrapping",
7771                        anyhow::Error::new(error),
7772                    )
7773                })?
7774            }
7775            #[cfg(feature = "tenferro-cuda")]
7776            ExecutionContext::Cuda(context) => {
7777                let uploaded = context.upload_cuda(&native).map_err(|error| {
7778                    IdxTensorError::operation(
7779                        "CUDA context construction",
7780                        anyhow::Error::new(error),
7781                    )
7782                })?;
7783                EagerTensor::from_tensor_in(
7784                    uploaded,
7785                    context.eager_runtime().map_err(|error| {
7786                        IdxTensorError::operation(
7787                            "CUDA context eager runtime",
7788                            anyhow::Error::new(error),
7789                        )
7790                    })?,
7791                )
7792                .map_err(|error| {
7793                    IdxTensorError::operation(
7794                        "CUDA context eager wrapping",
7795                        anyhow::Error::new(error),
7796                    )
7797                })?
7798            }
7799        };
7800        Self::from_inner(indices, inner).map_err(IdxTensorError::from)
7801    }
7802
7803    /// Create an all-ones tensor in a caller-owned execution context.
7804    ///
7805    /// # Examples
7806    ///
7807    /// ```
7808    /// use std::sync::Arc;
7809    /// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor};
7810    /// use tensor4all_tensorbackend::CpuExecutionContext;
7811    /// use tenferro_cpu::CpuBackend;
7812    ///
7813    /// let context = ExecutionContext::Cpu(Arc::new(
7814    ///     CpuExecutionContext::from_backend(CpuBackend::new()),
7815    /// ));
7816    /// let index = DynIndex::new_dyn(2);
7817    /// let tensor = IdxTensor::ones_in(&context, &[index])?;
7818    /// assert_eq!(tensor.to_vec::<f64>()?, vec![1.0, 1.0]);
7819    /// # Ok::<(), Box<dyn std::error::Error>>(())
7820    /// ```
7821    ///
7822    /// # Errors
7823    ///
7824    /// Returns [`IdxTensorError`] when the index dimensions or context
7825    /// construction is invalid.
7826    pub fn ones_in(
7827        context: &ExecutionContext,
7828        indices: &[DynIndex],
7829    ) -> std::result::Result<Self, IdxTensorError> {
7830        let dims = Self::expected_dims_from_indices(indices);
7831        let total_size = checked_total_size(&dims)?;
7832        Self::from_dense_in(context, indices.to_vec(), vec![1.0_f64; total_size])
7833    }
7834
7835    /// Validate that this tensor belongs to the supplied execution context.
7836    ///
7837    /// # Examples
7838    ///
7839    /// ```
7840    /// use std::sync::Arc;
7841    /// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor};
7842    /// use tensor4all_tensorbackend::CpuExecutionContext;
7843    /// use tenferro_cpu::CpuBackend;
7844    ///
7845    /// let context = ExecutionContext::Cpu(Arc::new(
7846    ///     CpuExecutionContext::from_backend(CpuBackend::new()),
7847    /// ));
7848    /// let index = DynIndex::new_dyn(2);
7849    /// let tensor = IdxTensor::from_dense_in(&context, vec![index], vec![1.0_f64, 2.0])?;
7850    /// tensor.validate_context(&context)?;
7851    /// assert_eq!(tensor.dims(), vec![2]);
7852    /// # Ok::<(), Box<dyn std::error::Error>>(())
7853    /// ```
7854    ///
7855    /// # Errors
7856    ///
7857    /// Returns [`IdxTensorError`] when the tensor does not belong to `context`.
7858    pub fn validate_context(
7859        &self,
7860        context: &ExecutionContext,
7861    ) -> std::result::Result<(), IdxTensorError> {
7862        let Some(inner) = self.storage.eager() else {
7863            if matches!(context, ExecutionContext::Cpu(_)) {
7864                return Ok(());
7865            }
7866            return Err(IdxTensorError::operation(
7867                "context validation",
7868                anyhow::anyhow!("tensor has no CUDA eager runtime"),
7869            ));
7870        };
7871        match context {
7872            ExecutionContext::Cpu(context) => {
7873                let expected = context.eager_runtime().map_err(|error| {
7874                    IdxTensorError::operation("CPU context validation", anyhow::Error::new(error))
7875                })?;
7876                if inner.ctx_id() != expected.id() {
7877                    return Err(IdxTensorError::operation(
7878                        "context validation",
7879                        anyhow::anyhow!("tensor belongs to a different CPU eager context"),
7880                    ));
7881                }
7882            }
7883            #[cfg(feature = "tenferro-cuda")]
7884            ExecutionContext::Cuda(context) => {
7885                self.validate_cuda_residency(context).map_err(|error| {
7886                    IdxTensorError::operation("CUDA context validation", anyhow::Error::new(error))
7887                })?;
7888            }
7889        }
7890        Ok(())
7891    }
7892
7893    /// Read back a small algorithm-decision tensor through its owning context.
7894    ///
7895    /// This is the explicit synchronization/readback boundary for rank
7896    /// selection and adaptive stopping payloads (singular values, norm
7897    /// vectors, scalar estimates). With a CUDA context the resident value is
7898    /// downloaded through that exact context; no arithmetic operand is ever
7899    /// reconstructed on the host from this payload.
7900    ///
7901    /// # Arguments
7902    ///
7903    /// * `context` - Caller-owned execution context the tensor must belong to.
7904    ///
7905    /// # Returns
7906    ///
7907    /// The decision values in column-major order (`f64`, or the real part for
7908    /// `Complex64`, matching the existing rank-selection semantics).
7909    ///
7910    /// # Examples
7911    ///
7912    /// ```
7913    /// use std::sync::Arc;
7914    /// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor};
7915    /// use tensor4all_tensorbackend::CpuExecutionContext;
7916    /// use tenferro_cpu::CpuBackend;
7917    ///
7918    /// let context = ExecutionContext::Cpu(Arc::new(
7919    ///     CpuExecutionContext::from_backend(CpuBackend::new()),
7920    /// ));
7921    /// let tensor = IdxTensor::from_dense_in(
7922    ///     &context,
7923    ///     vec![DynIndex::new_dyn(3)],
7924    ///     vec![0.5_f64, 2.0, 1.0],
7925    /// )?;
7926    /// assert_eq!(tensor.read_decision_data(&context)?, vec![0.5, 2.0, 1.0]);
7927    /// # Ok::<(), Box<dyn std::error::Error>>(())
7928    /// ```
7929    ///
7930    /// # Errors
7931    ///
7932    /// Returns [`IdxTensorError`] when the tensor does not belong to `context`,
7933    /// when the explicit transfer fails, or when the storage is neither `f64`
7934    /// nor `Complex64`.
7935    pub fn read_decision_data(
7936        &self,
7937        context: &ExecutionContext,
7938    ) -> std::result::Result<Vec<f64>, IdxTensorError> {
7939        self.validate_context(context)?;
7940        #[cfg(feature = "tenferro-cuda")]
7941        if let ExecutionContext::Cuda(cuda) = context {
7942            let host = self.download(cuda).map_err(|error| {
7943                IdxTensorError::operation("decision readback download", anyhow::Error::new(error))
7944            })?;
7945            return Self::decision_values(&host);
7946        }
7947        Self::decision_values(self)
7948    }
7949
7950    fn decision_values(host: &Self) -> std::result::Result<Vec<f64>, IdxTensorError> {
7951        if host.is_f64() {
7952            host.to_vec::<f64>().map_err(|error| {
7953                IdxTensorError::operation("decision readback", anyhow::Error::new(error))
7954            })
7955        } else if host.is_c64() {
7956            host.to_vec::<Complex64>()
7957                .map(|values| values.into_iter().map(|value| value.re).collect())
7958                .map_err(|error| {
7959                    IdxTensorError::operation("decision readback", anyhow::Error::new(error))
7960                })
7961        } else {
7962            Err(IdxTensorError::operation(
7963                "decision readback",
7964                anyhow::anyhow!("decision payloads support f64 and Complex64 storage only"),
7965            ))
7966        }
7967    }
7968
7969    /// Scale every element by a real factor in a caller-owned context.
7970    ///
7971    /// The scalar is built in the tensor's owning runtime (explicit upload
7972    /// for CUDA) and multiplied there; host-storage CPU tensors are adopted
7973    /// into `context` first via explicit construction. With a CPU context
7974    /// this matches [`IdxTensor::scale`] bit-for-bit.
7975    ///
7976    /// # Arguments
7977    ///
7978    /// * `factor` - Real scaling factor.
7979    /// * `context` - Caller-owned execution context the tensor belongs to.
7980    ///
7981    /// # Examples
7982    ///
7983    /// ```
7984    /// use std::sync::Arc;
7985    /// use tensor4all_core::{DynIndex, ExecutionContext, IdxTensor};
7986    /// use tensor4all_tensorbackend::CpuExecutionContext;
7987    /// use tenferro_cpu::CpuBackend;
7988    ///
7989    /// let context = ExecutionContext::Cpu(Arc::new(
7990    ///     CpuExecutionContext::from_backend(CpuBackend::new()),
7991    /// ));
7992    /// let tensor = IdxTensor::from_dense_in(
7993    ///     &context,
7994    ///     vec![DynIndex::new_dyn(2)],
7995    ///     vec![1.0_f64, 2.0],
7996    /// )?;
7997    /// let scaled = tensor.scale_in(3.0, &context)?;
7998    /// assert_eq!(scaled.to_vec::<f64>()?, vec![3.0, 6.0]);
7999    /// # Ok::<(), Box<dyn std::error::Error>>(())
8000    /// ```
8001    ///
8002    /// # Errors
8003    ///
8004    /// Returns [`IdxTensorError`] when the tensor does not belong to
8005    /// `context`, when the storage is neither `f64` nor `Complex64`, or when
8006    /// the scalar construction, upload, or multiplication fails.
8007    pub fn scale_in(
8008        &self,
8009        factor: f64,
8010        context: &ExecutionContext,
8011    ) -> std::result::Result<Self, IdxTensorError> {
8012        self.validate_context(context)?;
8013        let adopted_storage;
8014        let operand: &EagerTensor = match self.eager_runtime() {
8015            Some(_) => self.as_inner()?,
8016            // Host storage (CPU only; CUDA validated above): adopt into the
8017            // context through explicit construction.
8018            None => {
8019                adopted_storage = Self::adopt_host_into_context(self, context)?;
8020                adopted_storage.as_inner()?
8021            }
8022        };
8023        let scalar = Self::context_scalar_in(operand, factor, context)?;
8024        let scaled = operand.mul(&scalar).map_err(|error| {
8025            IdxTensorError::operation(
8026                "context-scoped scaling",
8027                anyhow::Error::new(error).context("eager multiplication failed"),
8028            )
8029        })?;
8030        Self::from_inner(self.indices.clone(), scaled).map_err(IdxTensorError::from)
8031    }
8032
8033    /// Adopt host storage into a context through explicit construction.
8034    fn adopt_host_into_context(
8035        tensor: &Self,
8036        context: &ExecutionContext,
8037    ) -> std::result::Result<Self, IdxTensorError> {
8038        if tensor.is_f64() {
8039            let values = tensor.to_vec::<f64>().map_err(|error| {
8040                IdxTensorError::operation("context adoption", anyhow::Error::new(error))
8041            })?;
8042            Self::from_dense_in(context, tensor.indices.clone(), values)
8043        } else if tensor.is_c64() {
8044            let values = tensor.to_vec::<Complex64>().map_err(|error| {
8045                IdxTensorError::operation("context adoption", anyhow::Error::new(error))
8046            })?;
8047            Self::from_dense_in(context, tensor.indices.clone(), values)
8048        } else {
8049            Err(IdxTensorError::operation(
8050                "context adoption",
8051                anyhow::anyhow!("adoption supports f64 and Complex64 storage only"),
8052            ))
8053        }
8054    }
8055
8056    /// Build a scalar in the operand's owning runtime (explicit upload for CUDA).
8057    fn context_scalar_in(
8058        operand: &EagerTensor,
8059        factor: f64,
8060        context: &ExecutionContext,
8061    ) -> std::result::Result<EagerTensor, IdxTensorError> {
8062        let native: NativeTensor = match operand.dtype() {
8063            DType::F64 => NativeTensor::from_vec_col_major(vec![], vec![factor]),
8064            DType::C64 => {
8065                NativeTensor::from_vec_col_major(vec![], vec![Complex64::new(factor, 0.0)])
8066            }
8067            other => {
8068                return Err(IdxTensorError::operation(
8069                    "context-scoped scaling",
8070                    anyhow::anyhow!("unsupported scalar dtype {other:?}"),
8071                ));
8072            }
8073        }
8074        .map_err(|error| {
8075            IdxTensorError::operation("context-scoped scaling", anyhow::anyhow!("{error}"))
8076        })?;
8077        let runtime = Self::context_eager_runtime(context)?;
8078        let resident = match context {
8079            ExecutionContext::Cpu(_) => native,
8080            #[cfg(feature = "tenferro-cuda")]
8081            ExecutionContext::Cuda(cuda) => cuda.upload_cuda(&native).map_err(|error| {
8082                IdxTensorError::operation(
8083                    "context-scoped scaling",
8084                    anyhow::Error::new(error).context("scalar upload failed"),
8085                )
8086            })?,
8087        };
8088        EagerTensor::from_tensor_in(resident, runtime).map_err(|error| {
8089            IdxTensorError::operation(
8090                "context-scoped scaling",
8091                anyhow::Error::new(error).context("scalar wrapping failed"),
8092            )
8093        })
8094    }
8095
8096    /// Resolve the eager runtime owned by a context.
8097    fn context_eager_runtime(
8098        context: &ExecutionContext,
8099    ) -> std::result::Result<Arc<EagerRuntime>, IdxTensorError> {
8100        match context {
8101            ExecutionContext::Cpu(context) => context.eager_runtime().map_err(|error| {
8102                IdxTensorError::operation("context-scoped scaling", anyhow::Error::new(error))
8103            }),
8104            #[cfg(feature = "tenferro-cuda")]
8105            ExecutionContext::Cuda(context) => context.eager_runtime().map_err(|error| {
8106                IdxTensorError::operation("context-scoped scaling", anyhow::Error::new(error))
8107            }),
8108        }
8109    }
8110
8111    /// Frobenius norm in a caller-owned execution context.
8112    ///
8113    /// Reductions run in the tensor's owning runtime; only the scalar result
8114    /// crosses the explicit readback boundary on CUDA. With a CPU context
8115    /// this delegates to [`IdxTensor::norm`] unchanged.
8116    ///
8117    /// # Errors
8118    ///
8119    /// Returns [`IdxTensorError`] when the tensor does not belong to
8120    /// `context`, when the storage is neither `f64` nor `Complex64`, or when
8121    /// the reductions or explicit readback fail.
8122    pub fn norm_in(&self, context: &ExecutionContext) -> std::result::Result<f64, IdxTensorError> {
8123        self.validate_context(context)?;
8124        #[cfg(feature = "tenferro-cuda")]
8125        if matches!(context, ExecutionContext::Cuda(_)) {
8126            let inner = self.cuda_eager_inner().ok_or_else(|| {
8127                IdxTensorError::operation(
8128                    "context-scoped norm",
8129                    anyhow::anyhow!("tensor has no resident eager value"),
8130                )
8131            })?;
8132            let rank = inner.shape().len();
8133            let axes: Vec<usize> = (0..rank).collect();
8134            let squared = inner
8135                .abs()
8136                .and_then(|magnitudes| magnitudes.reduce_sum_squares(&axes))
8137                .map_err(|error| {
8138                    IdxTensorError::operation("context-scoped norm", anyhow::anyhow!("{error}"))
8139                })?;
8140            let scalar = squared.reshape(&[1]).map_err(|error| {
8141                IdxTensorError::operation("context-scoped norm", anyhow::anyhow!("{error}"))
8142            })?;
8143            let values = Self::read_resident_vector(&scalar, 1, context).map_err(|error| {
8144                IdxTensorError::operation("context-scoped norm", anyhow::Error::new(error))
8145            })?;
8146            return Ok(values[0].sqrt());
8147        }
8148        #[cfg(not(feature = "tenferro-cuda"))]
8149        let _ = context;
8150        self.norm()
8151    }
8152}
8153
8154fn checked_product(dims: &[usize]) -> Result<usize> {
8155    dims.iter().try_fold(1usize, |acc, &dim| {
8156        acc.checked_mul(dim)
8157            .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
8158    })
8159}
8160
8161fn increment_col_major_coordinate(coords: &mut [usize], dims: &[usize]) {
8162    let mut carry = true;
8163    for (coordinate, &dim) in coords.iter_mut().zip(dims.iter()) {
8164        if !carry {
8165            break;
8166        }
8167        *coordinate += 1;
8168        if *coordinate == dim {
8169            *coordinate = 0;
8170        } else {
8171            carry = false;
8172        }
8173    }
8174}
8175
8176fn map_payload_support_coordinate(
8177    source_axis_classes: &[usize],
8178    target_axis_classes: &[usize],
8179    source_to_target_axes: &[usize],
8180    source_coords: &[usize],
8181    target_coords: &mut [usize],
8182    target_seen: &mut [bool],
8183) -> bool {
8184    if source_axis_classes.len() != source_to_target_axes.len()
8185        || target_coords.len() != target_seen.len()
8186    {
8187        return false;
8188    }
8189    target_seen.fill(false);
8190    for (source_axis, &target_axis) in source_to_target_axes.iter().enumerate() {
8191        let Some(&source_class) = source_axis_classes.get(source_axis) else {
8192            return false;
8193        };
8194        let Some(&target_class) = target_axis_classes.get(target_axis) else {
8195            return false;
8196        };
8197        let Some(&source_value) = source_coords.get(source_class) else {
8198            return false;
8199        };
8200        let Some(target_value) = target_coords.get_mut(target_class) else {
8201            return false;
8202        };
8203        if target_seen[target_class] {
8204            if *target_value != source_value {
8205                return false;
8206            }
8207        } else {
8208            *target_value = source_value;
8209            target_seen[target_class] = true;
8210        }
8211    }
8212    target_seen.iter().all(|&seen| seen)
8213}
8214
8215fn decode_col_major_linear(linear: usize, dims: &[usize]) -> Result<Vec<usize>> {
8216    let total = checked_product(dims)?;
8217    if !(linear < total) {
8218        return Err(anyhow::anyhow!(
8219            "linear offset {} out of bounds for dims {:?}",
8220            linear,
8221            dims
8222        ));
8223    };
8224    let mut remaining = linear;
8225    let mut out = Vec::with_capacity(dims.len());
8226    for &dim in dims {
8227        out.push(remaining % dim);
8228        remaining /= dim;
8229    }
8230    Ok(out)
8231}
8232
8233fn encode_col_major_linear(indices: &[usize], dims: &[usize]) -> Result<usize> {
8234    if !(indices.len() == dims.len()) {
8235        return Err(anyhow::anyhow!(
8236            "index rank {} does not match dims {:?}",
8237            indices.len(),
8238            dims
8239        ));
8240    };
8241    let mut linear = 0usize;
8242    let mut stride = 1usize;
8243    for (&index, &dim) in indices.iter().zip(dims.iter()) {
8244        if !(index < dim) {
8245            return Err(anyhow::anyhow!(
8246                "index {} out of bounds for dimension {}",
8247                index,
8248                dim
8249            ));
8250        };
8251        let term = index
8252            .checked_mul(stride)
8253            .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
8254        linear = linear
8255            .checked_add(term)
8256            .ok_or_else(|| anyhow::anyhow!("linear offset overflow"))?;
8257        stride = stride
8258            .checked_mul(dim)
8259            .ok_or_else(|| anyhow::anyhow!("stride overflow"))?;
8260    }
8261    Ok(linear)
8262}
8263
8264#[cfg(test)]
8265mod tests {
8266    use super::*;
8267    use num_complex::{Complex32, Complex64};
8268    use std::cell::Cell;
8269    use tensor4all_tensorbackend::StorageError;
8270
8271    #[test]
8272    fn structured_contraction_does_not_install_logical_dense_cache() {
8273        let n = 8;
8274        let left = DynIndex::new_dyn(n);
8275        let site = DynIndex::new_dyn(3);
8276        let right = DynIndex::new_dyn(n);
8277        let far = DynIndex::new_dyn(n);
8278        let end = DynIndex::new_dyn(n);
8279        let a =
8280            IdxTensor::from_copy_selector(left, site.clone(), right.clone(), 1, 1.0_f64).unwrap();
8281        let b =
8282            IdxTensor::from_copy_selector(right, site.clone(), far.clone(), 1, 2.0_f64).unwrap();
8283        let c = IdxTensor::from_copy_selector(far, site.clone(), end, 1, 3.0_f64).unwrap();
8284        let result = crate::defaults::contract::contract_with_options(
8285            &[&a, &b, &c],
8286            crate::defaults::contract::ContractionOptions::new()
8287                .with_retain_indices(std::slice::from_ref(&site)),
8288        )
8289        .unwrap();
8290
8291        assert_eq!(result.storage_kind(), StorageKind::Structured);
8292        assert!(result.eager_cache.get().is_none());
8293        assert_eq!(result.storage().unwrap().payload_len(), n * 3);
8294    }
8295
8296    #[test]
8297    fn binary_contraction_axis_classes_preserve_uncontracted_order() {
8298        let axis_classes =
8299            IdxTensor::binary_contraction_axis_classes(&[0, 1, 0], &[1], &[0, 1], &[0]).unwrap();
8300
8301        assert_eq!(axis_classes, vec![0, 0, 1]);
8302    }
8303
8304    #[test]
8305    fn structured_metrics_use_authoritative_compact_payload_for_all_dtypes() {
8306        fn check(tensor: IdxTensor, expected_sum: f64, expected_norm_squared: f64) {
8307            assert!(matches!(tensor.storage, IdxTensorStorage::Compact(_)));
8308            assert!(tensor.eager_cache.get().is_none());
8309            assert!((tensor.sum().unwrap().real() - expected_sum).abs() < 1.0e-6);
8310            assert!((tensor.norm_squared().unwrap() - expected_norm_squared).abs() < 1.0e-6);
8311            assert!((tensor.maxabs().unwrap() - 2.0).abs() < 1.0e-6);
8312            assert!(tensor.isapprox(&tensor, 0.0, 0.0).unwrap());
8313            assert!(tensor.eager_cache.get().is_none());
8314        }
8315
8316        let indices = || vec![DynIndex::new_dyn(2), DynIndex::new_dyn(2)];
8317        check(
8318            IdxTensor::from_diag(indices(), vec![1.0_f32, 2.0]).unwrap(),
8319            3.0,
8320            5.0,
8321        );
8322        check(
8323            IdxTensor::from_diag(indices(), vec![1.0_f64, 2.0]).unwrap(),
8324            3.0,
8325            5.0,
8326        );
8327        check(
8328            IdxTensor::from_diag(
8329                indices(),
8330                vec![Complex32::new(1.0, 0.0), Complex32::new(2.0, 0.0)],
8331            )
8332            .unwrap(),
8333            3.0,
8334            5.0,
8335        );
8336        check(
8337            IdxTensor::from_diag(
8338                indices(),
8339                vec![Complex64::new(1.0, 0.0), Complex64::new(2.0, 0.0)],
8340            )
8341            .unwrap(),
8342            3.0,
8343            5.0,
8344        );
8345    }
8346
8347    #[test]
8348    fn payload_storage_error_retains_typed_source() {
8349        let storage = Storage::from_dense_col_major(vec![1.0_f64, 2.0], &[2]).unwrap();
8350        let error = storage.scalar_at(&[2]).unwrap_err();
8351        assert!(matches!(error, StorageError::InvalidStructuredStorage(_)));
8352    }
8353
8354    #[test]
8355    fn materialization_error_retains_backend_source() {
8356        let native = NativeTensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
8357        let inner = EagerTensor::from_tensor_in(native, default_eager_ctx().unwrap()).unwrap();
8358        let storage = IdxTensorStorage::Eager {
8359            inner: Arc::new(inner),
8360            axis_classes: vec![0, 0],
8361        };
8362
8363        let error = storage.materialize(2).unwrap_err();
8364        assert!(matches!(error, TensorStorageError::Materialization { .. }));
8365        assert!(std::error::Error::source(&error).is_some());
8366    }
8367
8368    fn conjugate_with_injected_failure(
8369        tensor: &IdxTensor,
8370        target: *const EagerTensor,
8371        message: &'static str,
8372    ) -> IdxTensor {
8373        let calls = Cell::new(0usize);
8374        let conjugated = tensor.conj_with(&|inner| {
8375            calls.set(calls.get() + 1);
8376            if std::ptr::eq(inner, target) {
8377                Err(Arc::new(std::io::Error::other(message)) as _)
8378            } else {
8379                conjugate_eager(inner)
8380            }
8381        });
8382        assert!(calls.get() > 0, "injected closure was not reached");
8383        conjugated
8384    }
8385
8386    fn assert_unwrapped_conjugation_error(
8387        tensor: IdxTensor,
8388        target: *const EagerTensor,
8389        message: &'static str,
8390    ) -> IdxTensor {
8391        let conjugated = conjugate_with_injected_failure(&tensor, target, message);
8392        let error = conjugated.to_storage().unwrap_err();
8393        assert!(matches!(error, TensorStorageError::Conjugation { .. }));
8394        let source = std::error::Error::source(&error).unwrap();
8395        assert_eq!(source.to_string(), message);
8396        assert!(
8397            source.source().is_none(),
8398            "source was wrapped more than once"
8399        );
8400        conjugated
8401    }
8402
8403    #[test]
8404    fn authoritative_storage_conjugation_failure_is_deferred_without_detaching() {
8405        let i = DynIndex::new_dyn(2);
8406        let native = NativeTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
8407        let inner = EagerTensor::requires_grad_in(native, default_eager_ctx().unwrap()).unwrap();
8408        let tensor = IdxTensor::from_inner(vec![i], inner).unwrap();
8409        let source = match &tensor.storage {
8410            IdxTensorStorage::Eager { inner, .. } => Arc::clone(inner),
8411            IdxTensorStorage::Compact(payload) => Arc::clone(&payload.payload),
8412            IdxTensorStorage::Materialized(_) | IdxTensorStorage::Deferred { .. } => {
8413                panic!("tracked eager source expected")
8414            }
8415        };
8416
8417        let conjugated = conjugate_with_injected_failure(
8418            &tensor,
8419            Arc::as_ptr(&source),
8420            "forced authoritative eager conjugation failure",
8421        );
8422        assert!(conjugated.tracks_grad());
8423        assert!(conjugated.is_f64());
8424        assert!(!conjugated.is_complex());
8425        assert!(!conjugated.is_diag());
8426        assert_eq!(conjugated.dims(), vec![2]);
8427
8428        let error = conjugated.to_storage().unwrap_err();
8429        let source = std::error::Error::source(&error).unwrap();
8430        assert_eq!(
8431            source.to_string(),
8432            "forced authoritative eager conjugation failure"
8433        );
8434        assert!(conjugated.detach().is_err());
8435    }
8436
8437    #[test]
8438    fn structured_payload_conjugation_failure_retains_graph_and_blocks_detached_primal() {
8439        let i = DynIndex::new_dyn(2);
8440        let j = DynIndex::new_dyn(2);
8441        let tensor = IdxTensor::from_diag(
8442            vec![i, j],
8443            vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, -4.0)],
8444        )
8445        .unwrap()
8446        .enable_grad()
8447        .unwrap();
8448
8449        let target = Arc::as_ptr(
8450            &tensor
8451                .storage
8452                .compact_payload()
8453                .expect("tracked compact payload")
8454                .payload,
8455        );
8456        let conjugated = assert_unwrapped_conjugation_error(
8457            tensor,
8458            target,
8459            "forced structured AD conjugation failure",
8460        );
8461        assert!(conjugated.tracks_grad());
8462        assert!(conjugated.detach().is_err());
8463        assert!(conjugated.clone().enable_grad().is_err());
8464        assert!(conjugated.sum().is_err());
8465        assert!(conjugated.grad().is_err());
8466        assert!(conjugated.clear_grad().is_err());
8467        assert!(conjugated.maxabs().is_err());
8468        assert!(conjugated.norm_squared().is_err());
8469
8470        let twice_conjugated = conjugated.conj();
8471        let error = twice_conjugated.to_storage().unwrap_err();
8472        assert_eq!(
8473            std::error::Error::source(&error).unwrap().to_string(),
8474            "forced structured AD conjugation failure"
8475        );
8476    }
8477
8478    #[test]
8479    fn eager_cache_conjugation_failure_is_deferred_with_original_diagnostic() {
8480        let i = DynIndex::new_dyn(2);
8481        let j = DynIndex::new_dyn(2);
8482        let tensor = IdxTensor::from_diag(vec![i, j], vec![1.0_f64, 2.0]).unwrap();
8483        tensor.as_inner().unwrap();
8484        let target = Arc::as_ptr(tensor.eager_cache.get().unwrap());
8485
8486        let conjugated = assert_unwrapped_conjugation_error(
8487            tensor,
8488            target,
8489            "forced eager cache conjugation failure",
8490        );
8491        assert!(!conjugated.tracks_grad());
8492        assert!(conjugated.detach().is_err());
8493    }
8494
8495    #[test]
8496    fn encode_col_major_linear_rejects_offset_overflow() {
8497        let error =
8498            encode_col_major_linear(&[usize::MAX - 1, usize::MAX - 1], &[usize::MAX, usize::MAX])
8499                .unwrap_err();
8500        assert!(error.to_string().contains("linear offset overflow"));
8501    }
8502
8503    #[test]
8504    fn factorize_probe_batch_incremental_matches_the_column_based_path() {
8505        let row = DynIndex::new_dyn(6);
8506        let batch = DynIndex::new_dyn(4);
8507        let data: Vec<f64> = (0..24).map(|i| i as f64 * 0.37 - 1.5).collect();
8508        let batch_tensor =
8509            IdxTensor::from_dense(vec![row.clone(), batch.clone()], data.clone()).unwrap();
8510
8511        let columns: Vec<IdxTensor> = (0..4)
8512            .map(|position| {
8513                batch_tensor
8514                    .select_indices(std::slice::from_ref(&batch), &[position])
8515                    .unwrap()
8516            })
8517            .collect();
8518        let column_refs: Vec<&IdxTensor> = columns.iter().collect();
8519
8520        let from_batch = IdxTensor::factorize_probe_batch_incremental(
8521            None,
8522            &batch_tensor,
8523            &batch,
8524            std::slice::from_ref(&row),
8525        )
8526        .unwrap();
8527        let from_columns = IdxTensor::factorize_probe_columns_incremental(
8528            None,
8529            &column_refs,
8530            &column_refs,
8531            &[row],
8532        )
8533        .unwrap();
8534
8535        assert_eq!(from_batch.rank, from_columns.rank);
8536        let batch_data = from_batch.left.to_vec::<f64>().unwrap();
8537        let columns_data = from_columns.left.to_vec::<f64>().unwrap();
8538        assert_eq!(batch_data.len(), columns_data.len());
8539        let max_diff = batch_data
8540            .iter()
8541            .zip(columns_data.iter())
8542            .map(|(a, b)| (a - b).abs())
8543            .fold(0.0_f64, f64::max);
8544        assert!(
8545            max_diff < 1e-12,
8546            "batch-native and column-based factorizations disagree: max diff = {}",
8547            max_diff
8548        );
8549    }
8550}