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