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::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};
18use tenferro_ad::EagerTensor;
19use tenferro_einsum::eager_tensor::einsum_subscripts as eager_einsum_ad;
20use tenferro_einsum::EinsumSubscripts;
21use tensor4all_tensorbackend::{
22 axpby_native_tensor, contract_native_tensor, default_eager_ctx,
23 dense_native_tensor_from_col_major, diag_native_tensor_from_col_major,
24 native_tensor_primal_to_dense_col_major, native_tensor_primal_to_diag_c64,
25 native_tensor_primal_to_diag_f64, native_tensor_primal_to_storage, scale_native_tensor,
26 storage_payload_native_read_input, storage_to_native_tensor, AnyScalar as BackendScalar,
27 StorageScalar, TensorElement,
28};
29use tensor4all_tensorbackend::{Storage, StorageKind};
30
31use super::contract::PairwiseContractionOptions;
32use super::structured_contraction::{
33 normalize_payload_read_for_roots, storage_from_payload_native, storage_payload_native,
34 OperandLayout, StructuredContractionPlan, StructuredContractionSpec,
35};
36
37#[derive(Debug, Default, Clone)]
38struct PairwiseContractProfileEntry {
39 calls: usize,
40 total_time: Duration,
41 total_bytes: usize,
42}
43
44#[derive(Debug, Clone)]
72pub struct TensorHermitianEigendecomposition {
73 pub eigenvalues: Vec<f64>,
75 pub eigenvectors: TensorDynLen,
77 pub eigenvector_index: DynIndex,
79}
80
81thread_local! {
82 static PAIRWISE_CONTRACT_PROFILE_STATE: RefCell<HashMap<&'static str, PairwiseContractProfileEntry>> =
83 RefCell::new(HashMap::new());
84}
85
86fn pairwise_contract_profile_enabled() -> bool {
87 static ENABLED: OnceLock<bool> = OnceLock::new();
88 *ENABLED.get_or_init(|| env::var("T4A_PROFILE_PAIRWISE_CONTRACT").is_ok())
89}
90
91fn record_pairwise_contract_profile(section: &'static str, elapsed: Duration) {
92 if !pairwise_contract_profile_enabled() {
93 return;
94 }
95 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
96 let mut state = state.borrow_mut();
97 let entry = state.entry(section).or_default();
98 entry.calls += 1;
99 entry.total_time += elapsed;
100 });
101}
102
103fn record_pairwise_contract_profile_bytes(section: &'static str, bytes: usize) {
104 if !pairwise_contract_profile_enabled() {
105 return;
106 }
107 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
108 let mut state = state.borrow_mut();
109 let entry = state.entry(section).or_default();
110 entry.total_bytes += bytes;
111 });
112}
113
114fn profile_pairwise_contract_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
115 if !pairwise_contract_profile_enabled() {
116 return f();
117 }
118 let started = Instant::now();
119 let result = f();
120 record_pairwise_contract_profile(section, started.elapsed());
121 result
122}
123
124pub fn reset_pairwise_contract_profile() {
126 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| state.borrow_mut().clear());
127}
128
129pub fn print_and_reset_pairwise_contract_profile() {
131 if !pairwise_contract_profile_enabled() {
132 return;
133 }
134 PAIRWISE_CONTRACT_PROFILE_STATE.with(|state| {
135 let mut entries: Vec<_> = state
136 .borrow()
137 .iter()
138 .map(|(section, entry)| (*section, entry.clone()))
139 .collect();
140 state.borrow_mut().clear();
141 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
142
143 eprintln!("=== TensorDynLen pairwise contract profile ===");
144 for (section, entry) in entries {
145 let per_call_us = if entry.calls == 0 {
146 0.0
147 } else {
148 entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64
149 };
150 eprintln!(
151 "{section}: calls={} total={:.6}ms per_call={:.3}us bytes={}",
152 entry.calls,
153 entry.total_time.as_secs_f64() * 1.0e3,
154 per_call_us,
155 entry.total_bytes,
156 );
157 }
158 });
159}
160
161fn native_tensor_profile_bytes(native: &NativeTensor) -> usize {
162 let element_size = match native.dtype() {
163 DType::F32 => 4,
164 DType::F64 => 8,
165 DType::C32 => 8,
166 DType::C64 => 16,
167 DType::I32 => 4,
168 DType::I64 => 8,
169 DType::Bool => 1,
170 };
171 native.shape().iter().product::<usize>() * element_size
172}
173
174pub trait RandomScalar: TensorElement {
179 fn random_value<R: Rng>(rng: &mut R) -> Self;
181}
182
183impl RandomScalar for f64 {
184 fn random_value<R: Rng>(rng: &mut R) -> Self {
185 StandardNormal.sample(rng)
186 }
187}
188
189impl RandomScalar for Complex64 {
190 fn random_value<R: Rng>(rng: &mut R) -> Self {
191 Complex64::new(StandardNormal.sample(rng), StandardNormal.sample(rng))
192 }
193}
194
195pub fn compute_permutation_from_indices(
228 original_indices: &[DynIndex],
229 new_indices: &[DynIndex],
230) -> Result<Vec<usize>> {
231 anyhow::ensure!(
232 new_indices.len() == original_indices.len(),
233 "new_indices length must match original_indices length"
234 );
235
236 let mut perm = Vec::with_capacity(new_indices.len());
237 let mut used = std::collections::HashSet::new();
238
239 for new_idx in new_indices {
240 let pos = original_indices
243 .iter()
244 .position(|old_idx| old_idx == new_idx)
245 .ok_or_else(|| {
246 anyhow::anyhow!("new_indices must be a permutation of original_indices")
247 })?;
248
249 anyhow::ensure!(used.insert(pos), "duplicate index in new_indices");
250 perm.push(pos);
251 }
252
253 Ok(perm)
254}
255
256#[derive(Clone)]
257pub(crate) struct StructuredAdValue {
258 payload: Arc<EagerTensor>,
259 payload_dims: Vec<usize>,
260 axis_classes: Vec<usize>,
261}
262
263#[derive(Clone)]
264pub(crate) enum TensorDynLenStorage {
265 Materialized(Arc<Storage>),
266 Eager {
267 inner: Arc<EagerTensor>,
268 axis_classes: Vec<usize>,
269 },
270}
271
272impl TensorDynLenStorage {
273 fn from_storage(storage: Arc<Storage>) -> Self {
274 Self::Materialized(storage)
275 }
276
277 fn from_eager_dense(inner: EagerTensor, rank: usize) -> Self {
278 Self::Eager {
279 inner: Arc::new(inner),
280 axis_classes: TensorDynLen::dense_axis_classes(rank),
281 }
282 }
283
284 fn eager(&self) -> Option<&EagerTensor> {
285 match self {
286 Self::Materialized(_) => None,
287 Self::Eager { inner, .. } => Some(inner.as_ref()),
288 }
289 }
290
291 fn axis_classes(&self) -> &[usize] {
292 match self {
293 Self::Materialized(storage) => storage.axis_classes(),
294 Self::Eager { axis_classes, .. } => axis_classes,
295 }
296 }
297
298 fn payload_dims(&self) -> &[usize] {
299 match self {
300 Self::Materialized(storage) => storage.payload_dims(),
301 Self::Eager { inner, .. } => inner.data().shape(),
302 }
303 }
304
305 fn payload_strides_vec(&self) -> Vec<isize> {
306 match self {
307 Self::Materialized(storage) => storage.payload_strides().to_vec(),
308 Self::Eager { inner, .. } => {
309 let mut stride = 1isize;
310 inner
311 .data()
312 .shape()
313 .iter()
314 .map(|&dim| {
315 let current = stride;
316 stride *= isize::try_from(dim).unwrap_or(isize::MAX);
317 current
318 })
319 .collect()
320 }
321 }
322 }
323
324 fn is_f64(&self) -> bool {
325 match self {
326 Self::Materialized(storage) => storage.is_f64(),
327 Self::Eager { inner, .. } => inner.data().dtype() == DType::F64,
328 }
329 }
330
331 fn is_c64(&self) -> bool {
332 match self {
333 Self::Materialized(storage) => storage.is_c64(),
334 Self::Eager { inner, .. } => inner.data().dtype() == DType::C64,
335 }
336 }
337
338 fn is_complex(&self) -> bool {
339 match self {
340 Self::Materialized(storage) => storage.is_complex(),
341 Self::Eager { inner, .. } => matches!(inner.data().dtype(), DType::C32 | DType::C64),
342 }
343 }
344
345 fn is_diag(&self) -> bool {
346 match self {
347 Self::Materialized(storage) => storage.is_diag(),
348 Self::Eager { axis_classes, .. } => TensorDynLen::is_diag_axis_classes(axis_classes),
349 }
350 }
351
352 fn storage_kind(&self) -> StorageKind {
353 match self {
354 Self::Materialized(storage) => storage.storage_kind(),
355 Self::Eager { axis_classes, .. } => {
356 if axis_classes.iter().copied().eq(0..axis_classes.len()) {
357 StorageKind::Dense
358 } else if TensorDynLen::is_diag_axis_classes(axis_classes) {
359 StorageKind::Diagonal
360 } else {
361 StorageKind::Structured
362 }
363 }
364 }
365 }
366
367 fn materialize(&self, logical_rank: usize) -> Result<Arc<Storage>> {
368 match self {
369 Self::Materialized(storage) => Ok(Arc::clone(storage)),
370 Self::Eager {
371 inner,
372 axis_classes,
373 } => Ok(Arc::new(
374 TensorDynLen::storage_from_native_with_axis_classes(
375 inner.data(),
376 axis_classes,
377 logical_rank,
378 )?,
379 )),
380 }
381 }
382
383 fn scale(&self, scalar: &BackendScalar) -> Result<Storage> {
384 Ok(self.materialize(self.axis_classes().len())?.scale(scalar))
385 }
386
387 fn conj(&self) -> Result<Self> {
388 match self {
389 Self::Materialized(storage) => Ok(Self::Materialized(Arc::new(storage.conj()))),
390 Self::Eager {
391 inner,
392 axis_classes,
393 } => Ok(Self::Eager {
394 inner: Arc::new(inner.conj()?),
395 axis_classes: axis_classes.clone(),
396 }),
397 }
398 }
399
400 fn max_abs(&self) -> Result<f64> {
401 Ok(self.materialize(self.axis_classes().len())?.max_abs())
402 }
403}
404
405#[derive(Clone)]
457pub struct TensorDynLen {
458 pub indices: Vec<DynIndex>,
460 pub(crate) storage: TensorDynLenStorage,
462 pub(crate) structured_ad: Option<Arc<StructuredAdValue>>,
464 pub(crate) eager_cache: Arc<OnceLock<Arc<EagerTensor>>>,
466}
467
468impl TensorDynLen {
469 fn dense_axis_classes(rank: usize) -> Vec<usize> {
470 (0..rank).collect()
471 }
472
473 fn diag_axis_classes(rank: usize) -> Vec<usize> {
474 if rank == 0 {
475 vec![]
476 } else {
477 vec![0; rank]
478 }
479 }
480
481 fn canonicalize_axis_classes(axis_classes: &[usize]) -> Vec<usize> {
482 let mut map = std::collections::HashMap::new();
483 let mut next = 0usize;
484 axis_classes
485 .iter()
486 .map(|&class_id| {
487 *map.entry(class_id).or_insert_with(|| {
488 let canonical = next;
489 next += 1;
490 canonical
491 })
492 })
493 .collect()
494 }
495
496 fn permute_axis_classes(&self, perm: &[usize]) -> Vec<usize> {
497 let axis_classes = self.storage.axis_classes();
498 let permuted: Vec<usize> = perm.iter().map(|&index| axis_classes[index]).collect();
499 Self::canonicalize_axis_classes(&permuted)
500 }
501
502 fn normalize_insert_axis(op: &str, axis: isize, rank: usize) -> Result<usize> {
503 let normalized = if axis < 0 {
504 rank as isize + 1 + axis
505 } else {
506 axis
507 };
508 anyhow::ensure!(
509 normalized >= 0 && normalized <= rank as isize,
510 "{op}: axis {axis} is out of bounds for inserting into rank {rank}"
511 );
512 Ok(normalized as usize)
513 }
514
515 fn is_diag_axis_classes(axis_classes: &[usize]) -> bool {
516 axis_classes.len() >= 2 && axis_classes.iter().all(|&class_id| class_id == 0)
517 }
518
519 fn einsum_subscripts_from_usize_ids(
520 inputs: &[Vec<usize>],
521 output: &[usize],
522 ) -> Result<EinsumSubscripts> {
523 let input_labels = inputs
524 .iter()
525 .map(|ids| {
526 ids.iter()
527 .map(|&id| {
528 u32::try_from(id)
529 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
530 })
531 .collect::<Result<Vec<_>>>()
532 })
533 .collect::<Result<Vec<_>>>()?;
534 let output_labels = output
535 .iter()
536 .map(|&id| {
537 u32::try_from(id)
538 .map_err(|_| anyhow::anyhow!("einsum label {id} exceeds u32 range"))
539 })
540 .collect::<Result<Vec<_>>>()?;
541 let input_refs = input_labels.iter().map(Vec::as_slice).collect::<Vec<_>>();
542 Ok(EinsumSubscripts::new(&input_refs, &output_labels))
543 }
544
545 fn build_binary_einsum_subscripts(
546 lhs_rank: usize,
547 axes_a: &[usize],
548 rhs_rank: usize,
549 axes_b: &[usize],
550 ) -> Result<EinsumSubscripts> {
551 anyhow::ensure!(
552 axes_a.len() == axes_b.len(),
553 "contract axis length mismatch: lhs {:?}, rhs {:?}",
554 axes_a,
555 axes_b
556 );
557
558 let mut lhs_ids = vec![usize::MAX; lhs_rank];
559 let mut rhs_ids = vec![usize::MAX; rhs_rank];
560 let mut next_id = 0usize;
561
562 let mut seen_lhs = vec![false; lhs_rank];
563 let mut seen_rhs = vec![false; rhs_rank];
564
565 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
566 anyhow::ensure!(
567 lhs_axis < lhs_rank,
568 "lhs contract axis {lhs_axis} out of range"
569 );
570 anyhow::ensure!(
571 rhs_axis < rhs_rank,
572 "rhs contract axis {rhs_axis} out of range"
573 );
574 anyhow::ensure!(
575 !seen_lhs[lhs_axis],
576 "duplicate lhs contract axis {lhs_axis}"
577 );
578 anyhow::ensure!(
579 !seen_rhs[rhs_axis],
580 "duplicate rhs contract axis {rhs_axis}"
581 );
582 seen_lhs[lhs_axis] = true;
583 seen_rhs[rhs_axis] = true;
584 lhs_ids[lhs_axis] = next_id;
585 rhs_ids[rhs_axis] = next_id;
586 next_id += 1;
587 }
588
589 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
590 for id in &mut lhs_ids {
591 if *id == usize::MAX {
592 *id = next_id;
593 output_ids.push(next_id);
594 next_id += 1;
595 }
596 }
597 for id in &mut rhs_ids {
598 if *id == usize::MAX {
599 *id = next_id;
600 output_ids.push(next_id);
601 next_id += 1;
602 }
603 }
604
605 Self::einsum_subscripts_from_usize_ids(&[lhs_ids, rhs_ids], &output_ids)
606 }
607
608 fn binary_dot_general_config(axes_a: &[usize], axes_b: &[usize]) -> Result<DotGeneralConfig> {
609 anyhow::ensure!(
610 axes_a.len() == axes_b.len(),
611 "contract axis length mismatch: lhs {:?}, rhs {:?}",
612 axes_a,
613 axes_b
614 );
615 Ok(DotGeneralConfig {
616 lhs_contracting_dims: axes_a.to_vec(),
617 rhs_contracting_dims: axes_b.to_vec(),
618 lhs_batch_dims: vec![],
619 rhs_batch_dims: vec![],
620 })
621 }
622
623 fn binary_contraction_axis_classes(
624 lhs_axis_classes: &[usize],
625 axes_a: &[usize],
626 rhs_axis_classes: &[usize],
627 axes_b: &[usize],
628 ) -> Vec<usize> {
629 debug_assert_eq!(axes_a.len(), axes_b.len());
630
631 fn find(parent: &mut [usize], value: usize) -> usize {
632 if parent[value] != value {
633 parent[value] = find(parent, parent[value]);
634 }
635 parent[value]
636 }
637
638 fn union(parent: &mut [usize], lhs: usize, rhs: usize) {
639 let lhs_root = find(parent, lhs);
640 let rhs_root = find(parent, rhs);
641 if lhs_root != rhs_root {
642 parent[rhs_root] = lhs_root;
643 }
644 }
645
646 let lhs_payload_rank = lhs_axis_classes
647 .iter()
648 .copied()
649 .max()
650 .map(|value| value + 1)
651 .unwrap_or(0);
652 let rhs_payload_rank = rhs_axis_classes
653 .iter()
654 .copied()
655 .max()
656 .map(|value| value + 1)
657 .unwrap_or(0);
658 let rhs_offset = lhs_payload_rank;
659 let mut parent: Vec<usize> = (0..lhs_payload_rank + rhs_payload_rank).collect();
660
661 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
662 union(
663 &mut parent,
664 lhs_axis_classes[lhs_axis],
665 rhs_offset + rhs_axis_classes[rhs_axis],
666 );
667 }
668
669 let mut lhs_contracted = vec![false; lhs_axis_classes.len()];
670 for &axis in axes_a {
671 lhs_contracted[axis] = true;
672 }
673 let mut rhs_contracted = vec![false; rhs_axis_classes.len()];
674 for &axis in axes_b {
675 rhs_contracted[axis] = true;
676 }
677
678 let mut root_to_class = std::collections::HashMap::new();
679 let mut next_class = 0usize;
680 let mut axis_classes = Vec::new();
681
682 for (axis, &class_id) in lhs_axis_classes.iter().enumerate() {
683 if !lhs_contracted[axis] {
684 let root = find(&mut parent, class_id);
685 let class = *root_to_class.entry(root).or_insert_with(|| {
686 let value = next_class;
687 next_class += 1;
688 value
689 });
690 axis_classes.push(class);
691 }
692 }
693 for (axis, &class_id) in rhs_axis_classes.iter().enumerate() {
694 if !rhs_contracted[axis] {
695 let root = find(&mut parent, rhs_offset + class_id);
696 let class = *root_to_class.entry(root).or_insert_with(|| {
697 let value = next_class;
698 next_class += 1;
699 value
700 });
701 axis_classes.push(class);
702 }
703 }
704
705 axis_classes
706 }
707
708 fn scale_subscripts(rank: usize) -> Result<EinsumSubscripts> {
709 let ids: Vec<usize> = (0..rank).collect();
710 Self::einsum_subscripts_from_usize_ids(&[ids.clone(), Vec::new()], &ids)
711 }
712
713 fn validate_indices(indices: &[DynIndex]) -> Result<()> {
714 let mut seen = HashSet::new();
715 for idx in indices {
716 anyhow::ensure!(
717 seen.insert(idx.clone()),
718 "Tensor indices must all be unique"
719 );
720 }
721 Ok(())
722 }
723
724 fn validate_diag_dims(dims: &[usize]) -> Result<()> {
725 if !dims.is_empty() {
726 let first_dim = dims[0];
727 for (i, &dim) in dims.iter().enumerate() {
728 anyhow::ensure!(
729 dim == first_dim,
730 "DiagTensor requires all indices to have the same dimension, but dims[{i}] = {dim} != dims[0] = {first_dim}"
731 );
732 }
733 }
734 Ok(())
735 }
736
737 fn seed_native_payload(storage: &Storage, dims: &[usize]) -> Result<NativeTensor> {
738 storage_to_native_tensor(storage, dims)
739 }
740
741 fn empty_eager_cache() -> Arc<OnceLock<Arc<EagerTensor>>> {
742 Arc::new(OnceLock::new())
743 }
744
745 fn eager_cache_with(inner: EagerTensor) -> Arc<OnceLock<Arc<EagerTensor>>> {
746 let cache = Arc::new(OnceLock::new());
747 let _ = cache.set(Arc::new(inner));
748 cache
749 }
750
751 fn compact_payload_inner(&self) -> Result<EagerTensor> {
752 Ok(EagerTensor::from_tensor_in(
753 storage_payload_native(self.storage.materialize(self.indices.len())?.as_ref())?,
754 default_eager_ctx(),
755 ))
756 }
757
758 fn tracked_compact_payload_value(&self) -> Option<&StructuredAdValue> {
759 self.structured_ad.as_deref()
760 }
761
762 fn compact_payload_is_logical_dense(&self, payload_dims: &[usize]) -> bool {
763 self.storage.axis_classes() == Self::dense_axis_classes(self.indices.len())
764 && payload_dims == self.dims()
765 }
766
767 fn uses_tracked_compact_storage(&self) -> bool {
768 self.tracked_compact_payload_value()
769 .is_some_and(|value| !self.compact_payload_is_logical_dense(&value.payload_dims))
770 }
771
772 fn ensure_shape_packing_preserves_ad(&self, op_name: &str) -> Result<()> {
773 anyhow::ensure!(
774 !self.uses_tracked_compact_storage(),
775 "{op_name}: structured AD tensors with compact storage are not supported because materializing compact storage would detach gradients"
776 );
777 Ok(())
778 }
779
780 fn operand_indices_for_contraction(&self, conjugate: bool) -> Vec<DynIndex> {
781 if conjugate {
782 self.indices.iter().map(|index| index.conj()).collect()
783 } else {
784 self.indices.clone()
785 }
786 }
787
788 fn build_binary_contraction_labels(
789 lhs_rank: usize,
790 axes_a: &[usize],
791 rhs_rank: usize,
792 axes_b: &[usize],
793 ) -> Result<(Vec<usize>, Vec<usize>, Vec<usize>)> {
794 anyhow::ensure!(
795 axes_a.len() == axes_b.len(),
796 "contract axis length mismatch: lhs {:?}, rhs {:?}",
797 axes_a,
798 axes_b
799 );
800
801 let mut lhs_ids = vec![usize::MAX; lhs_rank];
802 let mut rhs_ids = vec![usize::MAX; rhs_rank];
803 let mut next_id = 0usize;
804
805 let mut seen_lhs = vec![false; lhs_rank];
806 let mut seen_rhs = vec![false; rhs_rank];
807
808 for (&lhs_axis, &rhs_axis) in axes_a.iter().zip(axes_b.iter()) {
809 anyhow::ensure!(
810 lhs_axis < lhs_rank,
811 "lhs contract axis {lhs_axis} out of range"
812 );
813 anyhow::ensure!(
814 rhs_axis < rhs_rank,
815 "rhs contract axis {rhs_axis} out of range"
816 );
817 anyhow::ensure!(
818 !seen_lhs[lhs_axis],
819 "duplicate lhs contract axis {lhs_axis}"
820 );
821 anyhow::ensure!(
822 !seen_rhs[rhs_axis],
823 "duplicate rhs contract axis {rhs_axis}"
824 );
825 seen_lhs[lhs_axis] = true;
826 seen_rhs[rhs_axis] = true;
827 lhs_ids[lhs_axis] = next_id;
828 rhs_ids[rhs_axis] = next_id;
829 next_id += 1;
830 }
831
832 let mut output_ids = Vec::with_capacity(lhs_rank + rhs_rank - 2 * axes_a.len());
833 for id in &mut lhs_ids {
834 if *id == usize::MAX {
835 *id = next_id;
836 output_ids.push(next_id);
837 next_id += 1;
838 }
839 }
840 for id in &mut rhs_ids {
841 if *id == usize::MAX {
842 *id = next_id;
843 output_ids.push(next_id);
844 next_id += 1;
845 }
846 }
847
848 Ok((lhs_ids, rhs_ids, output_ids))
849 }
850
851 fn build_payload_einsum_subscripts(
852 input_roots: &[Vec<usize>],
853 output_roots: &[usize],
854 ) -> Result<EinsumSubscripts> {
855 Self::einsum_subscripts_from_usize_ids(input_roots, output_roots)
856 }
857
858 fn normalize_eager_payload_for_roots(
859 payload: &EagerTensor,
860 roots: &[usize],
861 ) -> Result<(Option<EagerTensor>, Vec<usize>)> {
862 anyhow::ensure!(
863 payload.data().shape().len() == roots.len(),
864 "payload rank {} does not match root label count {}",
865 payload.data().shape().len(),
866 roots.len()
867 );
868
869 let mut current_payload = None;
870 let mut current_roots = roots.to_vec();
871 while let Some((axis_a, axis_b)) = Self::first_duplicate_pair(¤t_roots) {
872 let source = current_payload.as_ref().unwrap_or(payload);
873 current_payload = Some(source.extract_diag(axis_a, axis_b)?);
874 current_roots.remove(axis_b);
875 }
876
877 Ok((current_payload, current_roots))
878 }
879
880 fn first_duplicate_pair(values: &[usize]) -> Option<(usize, usize)> {
881 let mut first_axis_by_value = std::collections::HashMap::new();
882 for (axis, &value) in values.iter().enumerate() {
883 if let Some(&first_axis) = first_axis_by_value.get(&value) {
884 return Some((first_axis, axis));
885 }
886 first_axis_by_value.insert(value, axis);
887 }
888 None
889 }
890
891 fn binary_structured_contraction_plan(
892 &self,
893 other: &Self,
894 axes_a: &[usize],
895 axes_b: &[usize],
896 ) -> Result<(StructuredContractionPlan, Vec<Vec<usize>>, Vec<usize>)> {
897 let (lhs_labels, rhs_labels, output_labels) = Self::build_binary_contraction_labels(
898 self.indices.len(),
899 axes_a,
900 other.indices.len(),
901 axes_b,
902 )?;
903 let operands = vec![
904 OperandLayout::new(self.dims(), self.storage.axis_classes().to_vec())?,
905 OperandLayout::new(other.dims(), other.storage.axis_classes().to_vec())?,
906 ];
907 let spec = StructuredContractionSpec {
908 input_labels: vec![lhs_labels, rhs_labels],
909 output_labels,
910 retained_labels: Default::default(),
911 };
912 let plan = StructuredContractionPlan::new(&operands, &spec)?;
913 Ok((plan, spec.input_labels, spec.output_labels))
914 }
915
916 fn from_structured_payload_inner(
917 indices: Vec<DynIndex>,
918 payload_inner: EagerTensor,
919 payload_dims: Vec<usize>,
920 axis_classes: Vec<usize>,
921 ) -> Result<Self> {
922 Self::validate_indices(&indices)?;
923 if payload_inner.data().shape() != payload_dims {
924 return Err(anyhow::anyhow!(
925 "structured payload dims {:?} do not match planned payload dims {:?}",
926 payload_inner.data().shape(),
927 payload_dims
928 ));
929 }
930 let storage = storage_from_payload_native(
931 payload_inner.data().clone(),
932 &payload_dims,
933 axis_classes.clone(),
934 )?;
935 Self::validate_storage_matches_indices(&indices, &storage)?;
936 Ok(Self {
937 indices,
938 storage: TensorDynLenStorage::from_storage(Arc::new(storage)),
939 structured_ad: Some(Arc::new(StructuredAdValue {
940 payload: Arc::new(payload_inner),
941 payload_dims,
942 axis_classes,
943 })),
944 eager_cache: Self::empty_eager_cache(),
945 })
946 }
947
948 fn contract_structured_payloads(
949 &self,
950 other: &Self,
951 result_indices: Vec<DynIndex>,
952 axes_a: &[usize],
953 axes_b: &[usize],
954 ) -> Result<Self> {
955 let (plan, _, _) = self.binary_structured_contraction_plan(other, axes_a, axes_b)?;
956 let lhs_roots = plan.operand_plans[0].class_roots.clone();
957 let rhs_roots = plan.operand_plans[1].class_roots.clone();
958 let scalar_multiply =
959 lhs_roots.is_empty() && rhs_roots.is_empty() && plan.output_payload_roots.is_empty();
960
961 if let (Some(lhs_ad), Some(rhs_ad)) = (
962 self.tracked_compact_payload_value(),
963 other.tracked_compact_payload_value(),
964 ) {
965 if lhs_ad.payload.data().dtype() != rhs_ad.payload.data().dtype() {
966 return Err(anyhow::anyhow!(
967 "structured AD contraction requires matching payload dtypes"
968 ));
969 }
970 let (lhs_normalized, lhs_labels) =
971 Self::normalize_eager_payload_for_roots(lhs_ad.payload.as_ref(), &lhs_roots)?;
972 let (rhs_normalized, rhs_labels) =
973 Self::normalize_eager_payload_for_roots(rhs_ad.payload.as_ref(), &rhs_roots)?;
974 let lhs_payload = lhs_normalized
975 .as_ref()
976 .unwrap_or_else(|| lhs_ad.payload.as_ref());
977 let rhs_payload = rhs_normalized
978 .as_ref()
979 .unwrap_or_else(|| rhs_ad.payload.as_ref());
980 let payload = if scalar_multiply {
981 lhs_payload.mul(rhs_payload)?
982 } else {
983 let subscripts = Self::build_payload_einsum_subscripts(
984 &[lhs_labels, rhs_labels],
985 &plan.output_payload_roots,
986 )?;
987 eager_einsum_ad(&[lhs_payload, rhs_payload], &subscripts)?
988 };
989 return Self::from_structured_payload_inner(
990 result_indices,
991 payload,
992 plan.output_payload_dims,
993 plan.output_axis_classes,
994 );
995 }
996
997 if self.tracked_compact_payload_value().is_some()
998 || other.tracked_compact_payload_value().is_some()
999 {
1000 let lhs_owned = if self.tracked_compact_payload_value().is_some() {
1001 None
1002 } else {
1003 Some(self.compact_payload_inner()?)
1004 };
1005 let rhs_owned = if other.tracked_compact_payload_value().is_some() {
1006 None
1007 } else {
1008 Some(other.compact_payload_inner()?)
1009 };
1010 let lhs = if let Some(value) = self.tracked_compact_payload_value() {
1011 value.payload.as_ref()
1012 } else {
1013 lhs_owned
1014 .as_ref()
1015 .ok_or_else(|| anyhow::anyhow!("missing untracked left compact payload"))?
1016 };
1017 let rhs = if let Some(value) = other.tracked_compact_payload_value() {
1018 value.payload.as_ref()
1019 } else {
1020 rhs_owned
1021 .as_ref()
1022 .ok_or_else(|| anyhow::anyhow!("missing untracked right compact payload"))?
1023 };
1024 if lhs.data().dtype() != rhs.data().dtype() {
1025 return Err(anyhow::anyhow!(
1026 "structured AD contraction requires matching payload dtypes"
1027 ));
1028 }
1029 let (lhs_normalized, lhs_labels) =
1030 Self::normalize_eager_payload_for_roots(lhs, &lhs_roots)?;
1031 let (rhs_normalized, rhs_labels) =
1032 Self::normalize_eager_payload_for_roots(rhs, &rhs_roots)?;
1033 let lhs_payload = lhs_normalized.as_ref().unwrap_or(lhs);
1034 let rhs_payload = rhs_normalized.as_ref().unwrap_or(rhs);
1035 let payload = if scalar_multiply {
1036 lhs_payload.mul(rhs_payload)?
1037 } else {
1038 let subscripts = Self::build_payload_einsum_subscripts(
1039 &[lhs_labels, rhs_labels],
1040 &plan.output_payload_roots,
1041 )?;
1042 eager_einsum_ad(&[lhs_payload, rhs_payload], &subscripts)?
1043 };
1044 return Self::from_structured_payload_inner(
1045 result_indices,
1046 payload,
1047 plan.output_payload_dims,
1048 plan.output_axis_classes,
1049 );
1050 }
1051
1052 let lhs_storage = self.storage.materialize(self.indices.len())?;
1053 let rhs_storage = other.storage.materialize(other.indices.len())?;
1054 let lhs = storage_payload_native_read_input(lhs_storage.as_ref())?;
1055 let rhs = storage_payload_native_read_input(rhs_storage.as_ref())?;
1056 if lhs.dtype() != rhs.dtype() {
1057 return Err(anyhow::anyhow!(
1058 "structured payload contraction requires matching payload dtypes"
1059 ));
1060 }
1061 let (lhs, lhs_labels) = normalize_payload_read_for_roots(lhs, &lhs_roots)?;
1062 let (rhs, rhs_labels) = normalize_payload_read_for_roots(rhs, &rhs_roots)?;
1063 let payload = tensor4all_tensorbackend::einsum_native_tensor_reads(
1064 &[(&lhs, lhs_labels.as_slice()), (&rhs, rhs_labels.as_slice())],
1065 &plan.output_payload_roots,
1066 )?;
1067 let storage = storage_from_payload_native(
1068 payload,
1069 &plan.output_payload_dims,
1070 plan.output_axis_classes,
1071 )?;
1072 Self::from_storage(result_indices, Arc::new(storage))
1073 }
1074
1075 fn should_use_structured_payload_contract(&self, other: &Self) -> bool {
1076 let same_payload_dtype = self.storage.is_f64() == other.storage.is_f64()
1077 && self.storage.is_complex() == other.storage.is_complex();
1078 same_payload_dtype
1079 && (self.tracked_compact_payload_value().is_some()
1080 || other.tracked_compact_payload_value().is_some()
1081 || self.storage.axis_classes() != Self::dense_axis_classes(self.indices.len())
1082 || other.storage.axis_classes() != Self::dense_axis_classes(other.indices.len()))
1083 }
1084
1085 fn storage_from_native_with_axis_classes(
1086 native: &NativeTensor,
1087 axis_classes: &[usize],
1088 logical_rank: usize,
1089 ) -> Result<Storage> {
1090 if Self::is_diag_axis_classes(axis_classes) {
1091 match native.dtype() {
1092 DType::F32 | DType::F64 | DType::I32 | DType::I64 | DType::Bool => {
1093 Storage::from_diag_col_major(
1094 native_tensor_primal_to_diag_f64(native)?,
1095 logical_rank,
1096 )
1097 }
1098 DType::C32 | DType::C64 => Storage::from_diag_col_major(
1099 native_tensor_primal_to_diag_c64(native)?,
1100 logical_rank,
1101 ),
1102 }
1103 } else {
1104 native_tensor_primal_to_storage(native)
1105 }
1106 }
1107
1108 fn dense_selected_diag_payload<T: TensorElement + Copy + Zero>(
1109 payload: Vec<T>,
1110 kept_dims: &[usize],
1111 selected_positions: &[usize],
1112 ) -> Vec<T> {
1113 let output_len = kept_dims.iter().product::<usize>();
1114 let mut data = vec![T::zero(); output_len];
1115 if output_len == 0 {
1116 return data;
1117 }
1118
1119 let Some((&first_position, rest)) = selected_positions.split_first() else {
1120 return data;
1121 };
1122 if rest.iter().any(|&position| position != first_position) {
1123 return data;
1124 }
1125
1126 let value = payload[first_position];
1127 if kept_dims.is_empty() {
1128 data[0] = value;
1129 return data;
1130 }
1131
1132 let mut offset = 0usize;
1133 let mut stride = 1usize;
1134 for &dim in kept_dims {
1135 offset += first_position * stride;
1136 stride *= dim;
1137 }
1138 data[offset] = value;
1139 data
1140 }
1141
1142 fn select_diag_indices(
1143 &self,
1144 kept_indices: Vec<DynIndex>,
1145 kept_dims: Vec<usize>,
1146 positions: &[usize],
1147 ) -> Result<Self> {
1148 if self.storage.is_f64() {
1149 let storage = self.storage.materialize(self.indices.len())?;
1150 let payload = storage
1151 .payload_f64_col_major_vec()
1152 .map_err(anyhow::Error::msg)?;
1153 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions);
1154 Self::from_dense(kept_indices, data)
1155 } else if self.storage.is_c64() {
1156 let storage = self.storage.materialize(self.indices.len())?;
1157 let payload = storage
1158 .payload_c64_col_major_vec()
1159 .map_err(anyhow::Error::msg)?;
1160 let data = Self::dense_selected_diag_payload(payload, &kept_dims, positions);
1161 Self::from_dense(kept_indices, data)
1162 } else {
1163 Err(anyhow::anyhow!("unsupported diagonal storage scalar type"))
1164 }
1165 }
1166
1167 fn col_major_strides(dims: &[usize]) -> Result<Vec<isize>> {
1168 let mut strides = Vec::with_capacity(dims.len());
1169 let mut stride = 1isize;
1170 for &dim in dims {
1171 strides.push(stride);
1172 let dim = isize::try_from(dim)
1173 .map_err(|_| anyhow::anyhow!("dimension does not fit in isize"))?;
1174 stride = stride
1175 .checked_mul(dim)
1176 .ok_or_else(|| anyhow::anyhow!("column-major stride overflow"))?;
1177 }
1178 Ok(strides)
1179 }
1180
1181 fn zero_structured_selection<T>(
1182 kept_indices: Vec<DynIndex>,
1183 kept_dims: &[usize],
1184 ) -> Result<Self>
1185 where
1186 T: TensorElement + Zero,
1187 {
1188 let output_len = checked_product(kept_dims)?;
1189 Self::from_dense(kept_indices, vec![T::zero(); output_len])
1190 }
1191
1192 fn select_structured_indices_typed<T>(
1193 &self,
1194 payload: Vec<T>,
1195 kept_axes: &[usize],
1196 kept_indices: Vec<DynIndex>,
1197 kept_dims: Vec<usize>,
1198 selected_axes: &[usize],
1199 positions: &[usize],
1200 ) -> Result<Self>
1201 where
1202 T: TensorElement + StorageScalar + Zero,
1203 {
1204 let payload_dims = self.storage.payload_dims();
1205 let axis_classes = self.storage.axis_classes();
1206 let payload_rank = payload_dims.len();
1207 let mut selected_class_positions = vec![None; payload_rank];
1208
1209 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
1210 let class_id = axis_classes[axis];
1211 if let Some(existing) = selected_class_positions[class_id] {
1212 if existing != position {
1213 return Self::zero_structured_selection::<T>(kept_indices, &kept_dims);
1214 }
1215 } else {
1216 selected_class_positions[class_id] = Some(position);
1217 }
1218 }
1219
1220 let selected_class_kept = kept_axes
1221 .iter()
1222 .any(|&axis| selected_class_positions[axis_classes[axis]].is_some());
1223 if selected_class_kept {
1224 return self.select_structured_indices_dense(
1225 payload,
1226 kept_axes,
1227 kept_indices,
1228 kept_dims,
1229 &selected_class_positions,
1230 );
1231 }
1232
1233 let mut old_to_new_class = vec![None; payload_rank];
1234 let mut output_payload_dims = Vec::new();
1235 let mut output_axis_classes = Vec::with_capacity(kept_axes.len());
1236 for &axis in kept_axes {
1237 let class_id = axis_classes[axis];
1238 let new_class = match old_to_new_class[class_id] {
1239 Some(new_class) => new_class,
1240 None => {
1241 let new_class = output_payload_dims.len();
1242 old_to_new_class[class_id] = Some(new_class);
1243 output_payload_dims.push(payload_dims[class_id]);
1244 new_class
1245 }
1246 };
1247 output_axis_classes.push(new_class);
1248 }
1249
1250 let output_len = checked_product(&output_payload_dims)?;
1251 let mut output_payload = Vec::with_capacity(output_len);
1252 for linear in 0..output_len {
1253 let output_payload_index = decode_col_major_linear(linear, &output_payload_dims)?;
1254 let mut input_payload_index = vec![0usize; payload_rank];
1255 for class_id in 0..payload_rank {
1256 input_payload_index[class_id] =
1257 if let Some(position) = selected_class_positions[class_id] {
1258 position
1259 } else if let Some(new_class) = old_to_new_class[class_id] {
1260 output_payload_index[new_class]
1261 } else {
1262 return Err(anyhow::anyhow!(
1263 "structured payload class {class_id} is neither selected nor kept"
1264 ));
1265 };
1266 }
1267 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
1268 output_payload.push(payload[input_linear]);
1269 }
1270
1271 let output_strides = Self::col_major_strides(&output_payload_dims)?;
1272 let storage = Storage::new_structured(
1273 output_payload,
1274 output_payload_dims,
1275 output_strides,
1276 output_axis_classes,
1277 )?;
1278 Self::from_storage(kept_indices, Arc::new(storage))
1279 }
1280
1281 fn select_structured_indices_dense<T>(
1282 &self,
1283 payload: Vec<T>,
1284 kept_axes: &[usize],
1285 kept_indices: Vec<DynIndex>,
1286 kept_dims: Vec<usize>,
1287 selected_class_positions: &[Option<usize>],
1288 ) -> Result<Self>
1289 where
1290 T: TensorElement + Zero,
1291 {
1292 let payload_dims = self.storage.payload_dims();
1293 let axis_classes = self.storage.axis_classes();
1294 let output_len = checked_product(&kept_dims)?;
1295 let mut output = Vec::with_capacity(output_len);
1296
1297 for linear in 0..output_len {
1298 let kept_position = decode_col_major_linear(linear, &kept_dims)?;
1299 let mut input_payload_index = selected_class_positions.to_vec();
1300 let mut is_structural_zero = false;
1301
1302 for (&axis, &position) in kept_axes.iter().zip(kept_position.iter()) {
1303 let class_id = axis_classes[axis];
1304 match input_payload_index[class_id] {
1305 Some(existing) if existing != position => {
1306 is_structural_zero = true;
1307 break;
1308 }
1309 Some(_) => {}
1310 None => input_payload_index[class_id] = Some(position),
1311 }
1312 }
1313
1314 if is_structural_zero {
1315 output.push(T::zero());
1316 continue;
1317 }
1318
1319 let input_payload_index = input_payload_index
1320 .into_iter()
1321 .enumerate()
1322 .map(|(class_id, position)| {
1323 position.ok_or_else(|| {
1324 anyhow::anyhow!(
1325 "structured payload class {class_id} is neither selected nor kept"
1326 )
1327 })
1328 })
1329 .collect::<Result<Vec<_>>>()?;
1330 let input_linear = encode_col_major_linear(&input_payload_index, payload_dims)?;
1331 output.push(payload[input_linear]);
1332 }
1333
1334 Self::from_dense(kept_indices, output)
1335 }
1336
1337 fn select_structured_indices(
1338 &self,
1339 kept_axes: &[usize],
1340 kept_indices: Vec<DynIndex>,
1341 kept_dims: Vec<usize>,
1342 selected_axes: &[usize],
1343 positions: &[usize],
1344 ) -> Result<Self> {
1345 if self.storage.is_f64() {
1346 let storage = self.storage.materialize(self.indices.len())?;
1347 let payload = storage
1348 .payload_f64_col_major_vec()
1349 .map_err(anyhow::Error::msg)?;
1350 self.select_structured_indices_typed(
1351 payload,
1352 kept_axes,
1353 kept_indices,
1354 kept_dims,
1355 selected_axes,
1356 positions,
1357 )
1358 } else if self.storage.is_c64() {
1359 let storage = self.storage.materialize(self.indices.len())?;
1360 let payload = storage
1361 .payload_c64_col_major_vec()
1362 .map_err(anyhow::Error::msg)?;
1363 self.select_structured_indices_typed(
1364 payload,
1365 kept_axes,
1366 kept_indices,
1367 kept_dims,
1368 selected_axes,
1369 positions,
1370 )
1371 } else {
1372 Err(anyhow::anyhow!(
1373 "unsupported structured storage scalar type"
1374 ))
1375 }
1376 }
1377
1378 fn validate_storage_matches_indices(indices: &[DynIndex], storage: &Storage) -> Result<()> {
1379 let dims = Self::expected_dims_from_indices(indices);
1380 let storage_dims = storage.logical_dims();
1381 if storage_dims != dims {
1382 return Err(anyhow::anyhow!(
1383 "storage logical dims {:?} do not match indices dims {:?}",
1384 storage_dims,
1385 dims
1386 ));
1387 }
1388 if storage.is_diag() {
1389 Self::validate_diag_dims(&dims)?;
1390 }
1391 Ok(())
1392 }
1393
1394 fn try_materialized_inner(&self) -> Result<&EagerTensor> {
1395 if let Some(value) = self.tracked_compact_payload_value() {
1396 if self.compact_payload_is_logical_dense(&value.payload_dims) {
1397 return Ok(value.payload.as_ref());
1398 }
1399 }
1400 if let Some(inner) = self.storage.eager() {
1401 return Ok(inner);
1402 }
1403 if self.eager_cache.get().is_none() {
1404 let dims = self.dims();
1405 let native = profile_pairwise_contract_section("materialize_storage_to_native", || {
1406 let storage = self.storage.materialize(self.indices.len())?;
1407 Self::seed_native_payload(storage.as_ref(), &dims)
1408 })
1409 .context("TensorDynLen materialization failed")?;
1410 record_pairwise_contract_profile_bytes(
1411 "materialize_storage_to_native",
1412 native_tensor_profile_bytes(&native),
1413 );
1414 let _ = self.eager_cache.set(Arc::new(EagerTensor::from_tensor_in(
1415 native,
1416 default_eager_ctx(),
1417 )));
1418 }
1419 self.eager_cache
1420 .get()
1421 .map(|inner| inner.as_ref())
1422 .ok_or_else(|| {
1423 anyhow::anyhow!("TensorDynLen materialization cache was not initialized")
1424 })
1425 }
1426
1427 pub(crate) fn as_inner(&self) -> Result<&EagerTensor> {
1428 self.try_materialized_inner()
1429 }
1430
1431 #[inline]
1433 fn expected_dims_from_indices(indices: &[DynIndex]) -> Vec<usize> {
1434 indices.iter().map(|idx| idx.dim()).collect()
1435 }
1436
1437 pub fn dims(&self) -> Vec<usize> {
1456 Self::expected_dims_from_indices(&self.indices)
1457 }
1458
1459 pub fn select_indices(
1501 &self,
1502 selected_indices: &[DynIndex],
1503 positions: &[usize],
1504 ) -> Result<Self> {
1505 if selected_indices.len() != positions.len() {
1506 return Err(anyhow::anyhow!(
1507 "selected_indices length {} does not match positions length {}",
1508 selected_indices.len(),
1509 positions.len()
1510 ));
1511 }
1512 if selected_indices.is_empty() {
1513 return Ok(self.clone());
1514 }
1515
1516 let mut selected_axes = Vec::with_capacity(selected_indices.len());
1517 let mut seen_axes = HashSet::with_capacity(selected_indices.len());
1518 for (selected, &position) in selected_indices.iter().zip(positions.iter()) {
1519 let axis = self
1520 .indices
1521 .iter()
1522 .position(|index| index == selected)
1523 .ok_or_else(|| anyhow::anyhow!("selected index is not present in tensor"))?;
1524 if !seen_axes.insert(axis) {
1525 return Err(anyhow::anyhow!("selected index appears more than once"));
1526 }
1527 let dim = self.indices[axis].dim();
1528 if position >= dim {
1529 return Err(anyhow::anyhow!(
1530 "selected coordinate {position} is out of range for axis {axis} with dim {dim}"
1531 ));
1532 }
1533 selected_axes.push(axis);
1534 }
1535
1536 let kept_axes = self
1537 .indices
1538 .iter()
1539 .enumerate()
1540 .filter(|(axis, _)| !seen_axes.contains(axis))
1541 .map(|(axis, _)| axis)
1542 .collect::<Vec<_>>();
1543 let kept_indices = kept_axes
1544 .iter()
1545 .map(|&axis| self.indices[axis].clone())
1546 .collect::<Vec<_>>();
1547 let kept_dims = kept_axes
1548 .iter()
1549 .map(|&axis| self.indices[axis].dim())
1550 .collect::<Vec<_>>();
1551
1552 if self.storage.storage_kind() == StorageKind::Diagonal {
1553 return self.select_diag_indices(kept_indices, kept_dims, positions);
1554 }
1555 if self.storage.storage_kind() == StorageKind::Structured {
1556 return self.select_structured_indices(
1557 &kept_axes,
1558 kept_indices,
1559 kept_dims,
1560 &selected_axes,
1561 positions,
1562 );
1563 }
1564 if self.storage.storage_kind() != StorageKind::Dense {
1565 return Err(anyhow::anyhow!(
1566 "select_indices got unsupported storage kind {:?}",
1567 self.storage.storage_kind()
1568 ));
1569 }
1570
1571 let rank = self.indices.len();
1572 let mut starts = vec![0_i64; rank];
1573 let mut slice_sizes = self.dims();
1574 for (&axis, &position) in selected_axes.iter().zip(positions.iter()) {
1575 starts[axis] = i64::try_from(position)
1576 .map_err(|_| anyhow::anyhow!("selected coordinate does not fit in i64"))?;
1577 slice_sizes[axis] = 1;
1578 }
1579
1580 let starts_tensor = EagerTensor::from_tensor_in(
1581 NativeTensor::from_vec_col_major(vec![rank], starts),
1582 default_eager_ctx(),
1583 );
1584 let sliced = self
1585 .try_materialized_inner()?
1586 .dynamic_slice(&starts_tensor, &slice_sizes)?;
1587 Self::from_inner(kept_indices, sliced.reshape(&kept_dims)?)
1588 }
1589
1590 pub fn stack_along_new_index(
1623 tensors: &[&Self],
1624 new_index: DynIndex,
1625 axis: isize,
1626 ) -> Result<Self> {
1627 let first = tensors
1628 .first()
1629 .copied()
1630 .ok_or_else(|| anyhow::anyhow!("stack_along_new_index requires at least one tensor"))?;
1631 anyhow::ensure!(
1632 new_index.dim() == tensors.len(),
1633 "stack_along_new_index: new index dim {} does not match tensor count {}",
1634 new_index.dim(),
1635 tensors.len()
1636 );
1637
1638 let base_indices = first.indices.clone();
1639 for tensor in tensors.iter().copied().skip(1) {
1640 anyhow::ensure!(
1641 tensor.indices == base_indices,
1642 "stack_along_new_index: input tensors must have identical index order"
1643 );
1644 }
1645 for &tensor in tensors {
1646 tensor.ensure_shape_packing_preserves_ad("stack_along_new_index")?;
1647 }
1648
1649 let insert_axis =
1650 Self::normalize_insert_axis("stack_along_new_index", axis, base_indices.len())?;
1651 let mut result_indices = base_indices;
1652 result_indices.insert(insert_axis, new_index);
1653
1654 let inner_refs = tensors
1655 .iter()
1656 .map(|tensor| tensor.try_materialized_inner())
1657 .collect::<Result<Vec<_>>>()?;
1658 let stacked = EagerTensor::stack(&inner_refs, axis)?;
1659 Self::from_inner(result_indices, stacked)
1660 }
1661
1662 pub fn index_select(
1695 &self,
1696 source_index: &DynIndex,
1697 target_index: DynIndex,
1698 positions: &[usize],
1699 ) -> Result<Self> {
1700 anyhow::ensure!(
1701 target_index.dim() == positions.len(),
1702 "index_select: target index dim {} does not match position count {}",
1703 target_index.dim(),
1704 positions.len()
1705 );
1706 let axis = self
1707 .indices
1708 .iter()
1709 .position(|index| index == source_index)
1710 .ok_or_else(|| anyhow::anyhow!("index_select: source index is not present"))?;
1711 let source_dim = self.indices[axis].dim();
1712 for &position in positions {
1713 anyhow::ensure!(
1714 position < source_dim,
1715 "index_select: position {position} is out of range for source dim {source_dim}"
1716 );
1717 }
1718 self.ensure_shape_packing_preserves_ad("index_select")?;
1719
1720 let axis = isize::try_from(axis)
1721 .map_err(|_| anyhow::anyhow!("index_select: axis does not fit in isize"))?;
1722 let selected = self
1723 .try_materialized_inner()?
1724 .index_select(axis, positions)?;
1725 let mut result_indices = self.indices.clone();
1726 result_indices[axis as usize] = target_index;
1727 Self::from_inner(result_indices, selected)
1728 }
1729
1730 pub fn new(indices: Vec<DynIndex>, storage: Arc<Storage>) -> Result<Self> {
1750 Self::from_storage(indices, storage)
1751 }
1752
1753 pub fn from_indices(indices: Vec<DynIndex>, storage: Arc<Storage>) -> Result<Self> {
1775 Self::new(indices, storage)
1776 }
1777
1778 pub fn from_storage(indices: Vec<DynIndex>, storage: Arc<Storage>) -> Result<Self> {
1794 Self::validate_indices(&indices)?;
1795 Self::validate_storage_matches_indices(&indices, storage.as_ref())?;
1796 Ok(Self {
1797 indices,
1798 storage: TensorDynLenStorage::from_storage(storage),
1799 structured_ad: None,
1800 eager_cache: Self::empty_eager_cache(),
1801 })
1802 }
1803
1804 pub fn from_structured_storage(indices: Vec<DynIndex>, storage: Arc<Storage>) -> Result<Self> {
1828 Self::from_storage(indices, storage)
1829 }
1830
1831 pub(crate) fn from_native(indices: Vec<DynIndex>, native: NativeTensor) -> Result<Self> {
1833 let axis_classes = Self::dense_axis_classes(indices.len());
1834 Self::from_native_with_axis_classes(indices, native, axis_classes)
1835 }
1836
1837 pub(crate) fn from_native_with_axis_classes(
1838 indices: Vec<DynIndex>,
1839 native: NativeTensor,
1840 axis_classes: Vec<usize>,
1841 ) -> Result<Self> {
1842 Self::from_inner_with_axis_classes(
1843 indices,
1844 EagerTensor::from_tensor_in(native, default_eager_ctx()),
1845 axis_classes,
1846 )
1847 }
1848
1849 pub(crate) fn from_inner(indices: Vec<DynIndex>, inner: EagerTensor) -> Result<Self> {
1850 let axis_classes = Self::dense_axis_classes(indices.len());
1851 Self::from_inner_with_axis_classes(indices, inner, axis_classes)
1852 }
1853
1854 pub fn hermitian_eigendecomposition(
1896 &self,
1897 hermitian_tol: f64,
1898 ) -> Result<TensorHermitianEigendecomposition> {
1899 anyhow::ensure!(
1900 self.indices.len() == 2,
1901 "TensorDynLen::hermitian_eigendecomposition requires a rank-2 tensor, got rank {}",
1902 self.indices.len()
1903 );
1904 let dims = self.dims();
1905 anyhow::ensure!(
1906 dims[0] == dims[1],
1907 "TensorDynLen::hermitian_eigendecomposition requires a square matrix, got {}x{}",
1908 dims[0],
1909 dims[1]
1910 );
1911 anyhow::ensure!(
1912 dims[0] > 0,
1913 "TensorDynLen::hermitian_eigendecomposition requires a non-empty matrix"
1914 );
1915 anyhow::ensure!(
1916 hermitian_tol.is_finite() && hermitian_tol >= 0.0,
1917 "TensorDynLen::hermitian_eigendecomposition requires a finite non-negative tolerance"
1918 );
1919
1920 let input = self.try_materialized_inner()?;
1921 let (values, vectors) = tenferro_linalg::eager_tensor::eigh(input)
1922 .map_err(|source| anyhow::anyhow!("Hermitian eigendecomposition failed: {source}"))?;
1923
1924 let eigenvalue_index = DynIndex::new_dyn(dims[0]);
1925 let eigenvector_index = DynIndex::new_dyn(dims[0]);
1926 let eigenvalue_tensor = Self::from_inner(vec![eigenvalue_index], values)?;
1927 let eigenvalues = Self::read_real_eigenvalues(&eigenvalue_tensor, hermitian_tol)
1928 .with_context(|| {
1929 "TensorDynLen::hermitian_eigendecomposition failed to read eigenvalues"
1930 })?;
1931 let eigenvectors = Self::from_inner(
1932 vec![self.indices[0].clone(), eigenvector_index.clone()],
1933 vectors,
1934 )?;
1935
1936 Ok(TensorHermitianEigendecomposition {
1937 eigenvalues,
1938 eigenvectors,
1939 eigenvector_index,
1940 })
1941 }
1942
1943 fn read_real_eigenvalues(values: &Self, hermitian_tol: f64) -> Result<Vec<f64>> {
1944 if values.is_complex() {
1945 values
1946 .to_vec::<Complex64>()?
1947 .into_iter()
1948 .enumerate()
1949 .map(|(index, value)| {
1950 let imaginary = value.im.abs();
1951 let allowed = hermitian_tol * value.norm().max(1.0);
1952 anyhow::ensure!(
1953 imaginary <= allowed,
1954 "Hermitian eigenvalue {index} has imaginary part {imaginary}, exceeding tolerance {allowed}"
1955 );
1956 Ok(value.re)
1957 })
1958 .collect()
1959 } else {
1960 values.to_vec::<f64>()
1961 }
1962 }
1963
1964 pub(crate) fn from_diag_inner(
1965 indices: Vec<DynIndex>,
1966 payload_inner: EagerTensor,
1967 ) -> Result<Self> {
1968 let dims = Self::expected_dims_from_indices(&indices);
1969 Self::validate_indices(&indices)?;
1970 Self::validate_diag_dims(&dims)?;
1971 Self::validate_diag_payload_len(payload_inner.data().shape().iter().product(), &dims)?;
1972 let axis_classes = Self::diag_axis_classes(dims.len());
1973 let diag_inner = payload_inner.embed_diag(0, 1)?;
1974 Self::from_inner_with_axis_classes(indices, diag_inner, axis_classes)
1975 }
1976
1977 pub(crate) fn from_inner_with_axis_classes(
1978 indices: Vec<DynIndex>,
1979 inner: EagerTensor,
1980 axis_classes: Vec<usize>,
1981 ) -> Result<Self> {
1982 let dims = profile_pairwise_contract_section("from_inner_expected_dims", || {
1983 Self::expected_dims_from_indices(&indices)
1984 });
1985 profile_pairwise_contract_section("from_inner_validate_indices", || {
1986 Self::validate_indices(&indices)
1987 })?;
1988 if dims != inner.data().shape() {
1989 return Err(anyhow::anyhow!(
1990 "native payload dims {:?} do not match indices dims {:?}",
1991 inner.data().shape(),
1992 dims
1993 ));
1994 }
1995 if Self::is_diag_axis_classes(&axis_classes) {
1996 profile_pairwise_contract_section("from_inner_validate_diag_dims", || {
1997 Self::validate_diag_dims(&dims)
1998 })?;
1999 }
2000 let (storage, eager_cache) = if axis_classes == Self::dense_axis_classes(indices.len()) {
2001 (
2002 TensorDynLenStorage::from_eager_dense(inner, indices.len()),
2003 Self::empty_eager_cache(),
2004 )
2005 } else {
2006 let storage = profile_pairwise_contract_section("from_inner_storage_snapshot", || {
2007 Self::storage_from_native_with_axis_classes(
2008 inner.data(),
2009 &axis_classes,
2010 indices.len(),
2011 )
2012 })?;
2013 record_pairwise_contract_profile_bytes(
2014 "from_inner_storage_snapshot",
2015 native_tensor_profile_bytes(inner.data()),
2016 );
2017 (
2018 TensorDynLenStorage::from_storage(Arc::new(storage)),
2019 profile_pairwise_contract_section("from_inner_eager_cache", || {
2020 Self::eager_cache_with(inner)
2021 }),
2022 )
2023 };
2024 Ok(Self {
2025 indices,
2026 storage,
2027 structured_ad: None,
2028 eager_cache,
2029 })
2030 }
2031
2032 pub fn indices(&self) -> &[DynIndex] {
2034 &self.indices
2035 }
2036
2037 pub(crate) fn as_native(&self) -> Result<&NativeTensor> {
2039 Ok(self.try_materialized_inner()?.data())
2040 }
2041
2042 pub fn enable_grad(self) -> Result<Self> {
2044 let materialized = self.storage.materialize(self.indices.len())?;
2045 let payload = storage_payload_native(materialized.as_ref())
2046 .context("TensorDynLen::enable_grad failed")?;
2047 let payload_dims = self.storage.payload_dims().to_vec();
2048 let axis_classes = self.storage.axis_classes().to_vec();
2049 Ok(Self {
2050 indices: self.indices,
2051 storage: self.storage,
2052 structured_ad: Some(Arc::new(StructuredAdValue {
2053 payload: Arc::new(EagerTensor::requires_grad_in(payload, default_eager_ctx())),
2054 payload_dims,
2055 axis_classes,
2056 })),
2057 eager_cache: Self::empty_eager_cache(),
2058 })
2059 }
2060
2061 pub fn tracks_grad(&self) -> bool {
2063 self.structured_ad
2064 .as_ref()
2065 .is_some_and(|value| value.payload.tracks_grad())
2066 || self.storage.eager().is_some_and(EagerTensor::tracks_grad)
2067 || self
2068 .eager_cache
2069 .get()
2070 .is_some_and(|inner| inner.tracks_grad())
2071 }
2072
2073 pub fn grad(&self) -> Result<Option<Self>> {
2075 if let Some(value) = self.tracked_compact_payload_value() {
2076 return value
2077 .payload
2078 .grad()
2079 .map(|grad| {
2080 let storage = storage_from_payload_native(
2081 grad.as_ref().clone(),
2082 &value.payload_dims,
2083 value.axis_classes.clone(),
2084 )?;
2085 Self::from_storage(self.indices.clone(), Arc::new(storage))
2086 })
2087 .transpose();
2088 }
2089 self.try_materialized_inner()?
2090 .grad()
2091 .map(|grad| {
2092 Self::from_native_with_axis_classes(
2093 self.indices.clone(),
2094 grad.as_ref().clone(),
2095 self.storage.axis_classes().to_vec(),
2096 )
2097 })
2098 .transpose()
2099 }
2100
2101 pub fn clear_grad(&self) -> Result<()> {
2103 if let Some(value) = self.tracked_compact_payload_value() {
2104 value.payload.clear_grad();
2105 }
2106 if let Some(inner) = self.storage.eager() {
2107 inner.clear_grad();
2108 }
2109 if let Some(inner) = self.eager_cache.get() {
2110 inner.clear_grad();
2111 }
2112 Ok(())
2113 }
2114
2115 pub fn backward(&self) -> Result<()> {
2117 if let Some(value) = self.tracked_compact_payload_value() {
2118 return value
2119 .payload
2120 .backward()
2121 .map(|_| ())
2122 .map_err(|e| anyhow::anyhow!("TensorDynLen::backward failed: {e}"));
2123 }
2124 self.try_materialized_inner()?
2125 .backward()
2126 .map(|_| ())
2127 .map_err(|e| anyhow::anyhow!("TensorDynLen::backward failed: {e}"))
2128 }
2129
2130 pub fn detach(&self) -> Result<Self> {
2132 if self.tracked_compact_payload_value().is_some() {
2133 return Self::from_storage(
2134 self.indices.clone(),
2135 self.storage.materialize(self.indices.len())?,
2136 );
2137 }
2138 Self::from_inner_with_axis_classes(
2139 self.indices.clone(),
2140 self.try_materialized_inner()?.detach(),
2141 self.storage.axis_classes().to_vec(),
2142 )
2143 }
2144
2145 pub fn is_simple(&self) -> bool {
2147 true
2148 }
2149
2150 pub fn to_storage(&self) -> Result<Arc<Storage>> {
2152 self.storage.materialize(self.indices.len())
2153 }
2154
2155 pub fn storage(&self) -> Arc<Storage> {
2157 self.storage
2158 .materialize(self.indices.len())
2159 .expect("TensorDynLen storage materialization failed")
2160 }
2161
2162 pub fn sum(&self) -> Result<AnyScalar> {
2175 if self.indices.is_empty() {
2176 return AnyScalar::from_tensor(self.clone());
2177 }
2178 let axes: Vec<usize> = (0..self.indices.len()).collect();
2179 let reduced = self.try_materialized_inner()?.reduce_sum(&axes)?;
2180 AnyScalar::from_tensor(Self::from_inner(Vec::new(), reduced)?)
2181 }
2182
2183 pub fn only(&self) -> Result<AnyScalar> {
2204 let dims = self.dims();
2205 let total_size = checked_product(&dims)?;
2206 anyhow::ensure!(
2207 total_size == 1 || dims.is_empty(),
2208 "only() requires a scalar tensor (1 element), got {} elements with dims {:?}",
2209 if dims.is_empty() { 1 } else { total_size },
2210 dims
2211 );
2212 self.sum()
2213 }
2214
2215 pub fn permute_indices(&self, new_indices: &[DynIndex]) -> Result<Self> {
2246 let perm = compute_permutation_from_indices(&self.indices, new_indices)?;
2248 if perm.iter().copied().eq(0..perm.len()) {
2249 return Ok(Self {
2250 indices: new_indices.to_vec(),
2251 storage: self.storage.clone(),
2252 structured_ad: self.structured_ad.clone(),
2253 eager_cache: Arc::clone(&self.eager_cache),
2254 });
2255 }
2256
2257 let permuted = self.try_materialized_inner()?.transpose(&perm)?;
2258 let axis_classes = self.permute_axis_classes(&perm);
2259 Self::from_inner_with_axis_classes(new_indices.to_vec(), permuted, axis_classes)
2260 }
2261
2262 pub fn permute(&self, perm: &[usize]) -> Result<Self> {
2291 anyhow::ensure!(
2292 perm.len() == self.indices.len(),
2293 "permutation length must match tensor rank"
2294 );
2295 let mut seen = HashSet::new();
2296 for &axis in perm {
2297 anyhow::ensure!(
2298 axis < self.indices.len(),
2299 "permutation axis {axis} out of range"
2300 );
2301 anyhow::ensure!(seen.insert(axis), "duplicate axis {axis} in permutation");
2302 }
2303 if perm.iter().copied().eq(0..perm.len()) {
2304 return Ok(self.clone());
2305 }
2306
2307 let new_indices: Vec<DynIndex> = perm.iter().map(|&i| self.indices[i].clone()).collect();
2309 let permuted = self.try_materialized_inner()?.transpose(perm)?;
2310 let axis_classes = self.permute_axis_classes(perm);
2311 Self::from_inner_with_axis_classes(new_indices, permuted, axis_classes)
2312 }
2313
2314 pub(crate) fn try_contract_pairwise_default(&self, other: &Self) -> Result<Self> {
2315 self.try_contract_pairwise_default_with_options(other, PairwiseContractionOptions::new())
2316 }
2317
2318 pub(crate) fn try_contract_pairwise_default_with_options(
2319 &self,
2320 other: &Self,
2321 options: PairwiseContractionOptions,
2322 ) -> Result<Self> {
2323 let self_indices = profile_pairwise_contract_section("operand_indices", || {
2324 self.operand_indices_for_contraction(options.lhs_conj)
2325 });
2326 let other_indices = profile_pairwise_contract_section("operand_indices", || {
2327 other.operand_indices_for_contraction(options.rhs_conj)
2328 });
2329 let self_dims = profile_pairwise_contract_section("expected_dims", || {
2330 Self::expected_dims_from_indices(&self_indices)
2331 });
2332 let other_dims = profile_pairwise_contract_section("expected_dims", || {
2333 Self::expected_dims_from_indices(&other_indices)
2334 });
2335 let spec = profile_pairwise_contract_section("prepare_contraction", || {
2336 prepare_contraction(&self_indices, &self_dims, &other_indices, &other_dims)
2337 })
2338 .context("contraction preparation failed")?;
2339 let result_axis_classes = profile_pairwise_contract_section("result_axis_classes", || {
2340 Self::binary_contraction_axis_classes(
2341 self.storage.axis_classes(),
2342 &spec.axes_a,
2343 other.storage.axis_classes(),
2344 &spec.axes_b,
2345 )
2346 });
2347
2348 if profile_pairwise_contract_section("structured_check", || {
2349 self.should_use_structured_payload_contract(other)
2350 }) {
2351 if options.has_conj() {
2352 let lhs = if options.lhs_conj {
2353 self.conj()
2354 } else {
2355 self.clone()
2356 };
2357 let rhs = if options.rhs_conj {
2358 other.conj()
2359 } else {
2360 other.clone()
2361 };
2362 return profile_pairwise_contract_section("structured_conj_fallback", || {
2363 lhs.try_contract_pairwise_default(&rhs)
2364 });
2365 }
2366 return profile_pairwise_contract_section("structured_payload_contract", || {
2367 self.contract_structured_payloads(
2368 other,
2369 spec.result_indices.into_vec(),
2370 &spec.axes_a,
2371 &spec.axes_b,
2372 )
2373 });
2374 }
2375
2376 if self.indices.is_empty() && other.indices.is_empty() {
2377 if options.has_conj() {
2378 let lhs = if options.lhs_conj {
2379 self.conj()
2380 } else {
2381 self.clone()
2382 };
2383 let rhs = if options.rhs_conj {
2384 other.conj()
2385 } else {
2386 other.clone()
2387 };
2388 return lhs.try_contract_pairwise_default(&rhs);
2389 }
2390 let result = profile_pairwise_contract_section("scalar_mul", || {
2391 Ok::<_, anyhow::Error>(
2392 self.try_materialized_inner()?
2393 .mul(other.try_materialized_inner()?)?,
2394 )
2395 })?;
2396 return profile_pairwise_contract_section("from_inner", || {
2397 Self::from_inner(spec.result_indices.into_vec(), result)
2398 });
2399 }
2400
2401 let self_native = profile_pairwise_contract_section("as_native", || self.as_native())?;
2402 let other_native = profile_pairwise_contract_section("as_native", || other.as_native())?;
2403 if self_native.dtype() != other_native.dtype() {
2404 if options.has_conj() {
2405 let lhs = if options.lhs_conj {
2406 self.conj()
2407 } else {
2408 self.clone()
2409 };
2410 let rhs = if options.rhs_conj {
2411 other.conj()
2412 } else {
2413 other.clone()
2414 };
2415 return lhs.try_contract_pairwise_default(&rhs);
2416 }
2417 let result_native = profile_pairwise_contract_section("native_contract", || {
2418 contract_native_tensor(self_native, &spec.axes_a, other_native, &spec.axes_b)
2419 })?;
2420 return profile_pairwise_contract_section("from_native", || {
2421 Self::from_native_with_axis_classes(
2422 spec.result_indices.into_vec(),
2423 result_native,
2424 result_axis_classes,
2425 )
2426 });
2427 }
2428
2429 let config = profile_pairwise_contract_section("build_dot_general_config", || {
2430 Self::binary_dot_general_config(&spec.axes_a, &spec.axes_b)
2431 })?;
2432 let result = profile_pairwise_contract_section("dot_general_with_conj", || {
2433 let lhs = profile_pairwise_contract_section("lhs_try_materialized_inner", || {
2434 self.try_materialized_inner()
2435 })?;
2436 let rhs = profile_pairwise_contract_section("rhs_try_materialized_inner", || {
2437 other.try_materialized_inner()
2438 })?;
2439 profile_pairwise_contract_section("dot_general_execute", || {
2440 lhs.dot_general_with_conj(rhs, &config, options.lhs_conj, options.rhs_conj)
2441 })
2442 .map_err(anyhow::Error::from)
2443 })?;
2444 record_pairwise_contract_profile_bytes(
2445 "dot_general_output",
2446 native_tensor_profile_bytes(result.data()),
2447 );
2448 profile_pairwise_contract_section("from_inner_axis_classes", || {
2449 Self::from_inner_with_axis_classes(
2450 spec.result_indices.into_vec(),
2451 result,
2452 result_axis_classes,
2453 )
2454 })
2455 }
2456
2457 pub(crate) fn try_tensordot_pairwise_explicit(
2458 &self,
2459 other: &Self,
2460 pairs: &[(DynIndex, DynIndex)],
2461 ) -> Result<Self> {
2462 use crate::index_ops::ContractionError;
2463
2464 let self_dims = Self::expected_dims_from_indices(&self.indices);
2465 let other_dims = Self::expected_dims_from_indices(&other.indices);
2466 let spec = prepare_contraction_pairs(
2467 &self.indices,
2468 &self_dims,
2469 &other.indices,
2470 &other_dims,
2471 pairs,
2472 )
2473 .map_err(|e| match e {
2474 ContractionError::NoCommonIndices => {
2475 anyhow::anyhow!("tensordot: No pairs specified for contraction")
2476 }
2477 ContractionError::BatchContractionNotImplemented => anyhow::anyhow!(
2478 "tensordot: Common index found but not in contraction pairs. \
2479 Batch contraction is not yet implemented."
2480 ),
2481 ContractionError::IndexNotFound { tensor } => {
2482 anyhow::anyhow!("tensordot: Index not found in {} tensor", tensor)
2483 }
2484 ContractionError::DimensionMismatch {
2485 pos_a,
2486 pos_b,
2487 dim_a,
2488 dim_b,
2489 } => anyhow::anyhow!(
2490 "tensordot: Dimension mismatch: self[{}]={} != other[{}]={}",
2491 pos_a,
2492 dim_a,
2493 pos_b,
2494 dim_b
2495 ),
2496 ContractionError::DuplicateAxis { tensor, pos } => {
2497 anyhow::anyhow!("tensordot: Duplicate axis {} in {} tensor", pos, tensor)
2498 }
2499 })?;
2500 let result_axis_classes = Self::binary_contraction_axis_classes(
2501 self.storage.axis_classes(),
2502 &spec.axes_a,
2503 other.storage.axis_classes(),
2504 &spec.axes_b,
2505 );
2506
2507 if self.should_use_structured_payload_contract(other) {
2508 return self.contract_structured_payloads(
2509 other,
2510 spec.result_indices.into_vec(),
2511 &spec.axes_a,
2512 &spec.axes_b,
2513 );
2514 }
2515
2516 if self.indices.is_empty() && other.indices.is_empty() {
2517 let result = self
2518 .try_materialized_inner()?
2519 .mul(other.try_materialized_inner()?)
2520 .map_err(|e| anyhow::anyhow!("tensordot scalar multiply failed: {e}"))?;
2521 return Self::from_inner(spec.result_indices.into_vec(), result);
2522 }
2523
2524 let self_native = self.as_native()?;
2525 let other_native = other.as_native()?;
2526 if self_native.dtype() != other_native.dtype() {
2527 let result_native =
2528 contract_native_tensor(self_native, &spec.axes_a, other_native, &spec.axes_b)?;
2529 return Self::from_native_with_axis_classes(
2530 spec.result_indices.into_vec(),
2531 result_native,
2532 result_axis_classes,
2533 );
2534 }
2535
2536 let subscripts = Self::build_binary_einsum_subscripts(
2537 self.indices.len(),
2538 &spec.axes_a,
2539 other.indices.len(),
2540 &spec.axes_b,
2541 )?;
2542 let result = eager_einsum_ad(
2543 &[
2544 self.try_materialized_inner()?,
2545 other.try_materialized_inner()?,
2546 ],
2547 &subscripts,
2548 )
2549 .map_err(|e| anyhow::anyhow!("tensordot failed: {e}"))?;
2550 Self::from_inner_with_axis_classes(
2551 spec.result_indices.into_vec(),
2552 result,
2553 result_axis_classes,
2554 )
2555 }
2556
2557 pub(crate) fn try_outer_product_pairwise(&self, other: &Self) -> Result<Self> {
2558 use anyhow::Context;
2559
2560 let common_positions = common_ind_positions(&self.indices, &other.indices);
2562 if !common_positions.is_empty() {
2563 let common_ids: Vec<_> = common_positions
2564 .iter()
2565 .map(|(pos_a, _)| self.indices[*pos_a].id())
2566 .collect();
2567 return Err(anyhow::anyhow!(
2568 "outer_product: tensors have common indices {:?}. \
2569 Use tensordot to contract common indices, or use sim() to replace \
2570 indices with fresh IDs before computing outer product.",
2571 common_ids
2572 ))
2573 .context("outer_product: common indices found");
2574 }
2575
2576 let mut result_indices = self.indices.clone();
2578 result_indices.extend(other.indices.iter().cloned());
2579 let result_axis_classes = Self::binary_contraction_axis_classes(
2580 self.storage.axis_classes(),
2581 &[],
2582 other.storage.axis_classes(),
2583 &[],
2584 );
2585 if self.should_use_structured_payload_contract(other) {
2586 return self.contract_structured_payloads(other, result_indices, &[], &[]);
2587 }
2588 let self_native = self.as_native()?;
2589 let other_native = other.as_native()?;
2590 if self_native.dtype() != other_native.dtype() {
2591 let result_native = contract_native_tensor(self_native, &[], other_native, &[])?;
2592 return Self::from_native_with_axis_classes(
2593 result_indices,
2594 result_native,
2595 result_axis_classes,
2596 );
2597 }
2598
2599 let subscripts = Self::build_binary_einsum_subscripts(
2600 self.indices.len(),
2601 &[],
2602 other.indices.len(),
2603 &[],
2604 )?;
2605 let result = eager_einsum_ad(
2606 &[
2607 self.try_materialized_inner()?,
2608 other.try_materialized_inner()?,
2609 ],
2610 &subscripts,
2611 )
2612 .map_err(|e| anyhow::anyhow!("outer_product failed: {e}"))?;
2613 Self::from_inner_with_axis_classes(result_indices, result, result_axis_classes)
2614 }
2615}
2616
2617impl TensorDynLen {
2622 pub fn random<T: RandomScalar, R: Rng>(rng: &mut R, indices: Vec<DynIndex>) -> Result<Self> {
2649 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
2650 let size = checked_product(&dims)?;
2651 let data: Vec<T> = (0..size).map(|_| T::random_value(rng)).collect();
2652 Self::from_dense(indices, data)
2653 }
2654}
2655
2656impl TensorDynLen {
2657 pub fn add(&self, other: &Self) -> Result<Self> {
2691 if self.indices.len() != other.indices.len() {
2693 return Err(anyhow::anyhow!(
2694 "Index count mismatch: self has {} indices, other has {}",
2695 self.indices.len(),
2696 other.indices.len()
2697 ));
2698 }
2699
2700 let self_set: HashSet<_> = self.indices.iter().collect();
2702 let other_set: HashSet<_> = other.indices.iter().collect();
2703
2704 if self_set != other_set {
2705 return Err(anyhow::anyhow!(
2706 "Index set mismatch: tensors must have the same indices"
2707 ));
2708 }
2709
2710 let other_aligned = other.permute_indices(&self.indices)?;
2712
2713 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
2715 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
2716 if self_expected_dims != other_expected_dims {
2717 use crate::TagSetLike;
2718 let fmt = |indices: &[DynIndex]| -> Vec<String> {
2719 indices
2720 .iter()
2721 .map(|idx| {
2722 let tags: Vec<String> = idx.tags().iter().collect();
2723 format!("{:?}(dim={},tags={:?})", idx.id(), idx.dim(), tags)
2724 })
2725 .collect()
2726 };
2727 return Err(anyhow::anyhow!(
2728 "Dimension mismatch after alignment.\n\
2729 self: dims={:?}, indices(order)={:?}\n\
2730 other_aligned: dims={:?}, indices(order)={:?}",
2731 self_expected_dims,
2732 fmt(&self.indices),
2733 other_expected_dims,
2734 fmt(&other_aligned.indices)
2735 ));
2736 }
2737
2738 self.axpby(
2739 AnyScalar::new_real(1.0),
2740 &other_aligned,
2741 AnyScalar::new_real(1.0),
2742 )
2743 }
2744
2745 pub fn axpby(&self, a: AnyScalar, other: &Self, b: AnyScalar) -> Result<Self> {
2767 if self.indices.len() != other.indices.len() {
2769 return Err(anyhow::anyhow!(
2770 "Index count mismatch: self has {} indices, other has {}",
2771 self.indices.len(),
2772 other.indices.len()
2773 ));
2774 }
2775
2776 let self_set: HashSet<_> = self.indices.iter().collect();
2778 let other_set: HashSet<_> = other.indices.iter().collect();
2779 if self_set != other_set {
2780 return Err(anyhow::anyhow!(
2781 "Index set mismatch: tensors must have the same indices"
2782 ));
2783 }
2784
2785 let other_aligned = other.permute_indices(&self.indices)?;
2787
2788 let self_expected_dims = Self::expected_dims_from_indices(&self.indices);
2790 let other_expected_dims = Self::expected_dims_from_indices(&other_aligned.indices);
2791 if self_expected_dims != other_expected_dims {
2792 return Err(anyhow::anyhow!(
2793 "Dimension mismatch after alignment: self={:?}, other_aligned={:?}",
2794 self_expected_dims,
2795 other_expected_dims
2796 ));
2797 }
2798
2799 let axis_classes = if self.storage.axis_classes() == other_aligned.storage.axis_classes() {
2800 self.storage.axis_classes().to_vec()
2801 } else {
2802 Self::dense_axis_classes(self.indices.len())
2803 };
2804
2805 let same_compact_layout = self.storage.payload_dims()
2806 == other_aligned.storage.payload_dims()
2807 && self.storage.payload_strides_vec() == other_aligned.storage.payload_strides_vec()
2808 && self.storage.axis_classes() == other_aligned.storage.axis_classes();
2809 if same_compact_layout
2810 && !self.tracks_grad()
2811 && !other_aligned.tracks_grad()
2812 && !a.tracks_grad()
2813 && !b.tracks_grad()
2814 {
2815 let lhs_storage = self.storage.materialize(self.indices.len())?;
2816 let rhs_storage = other_aligned
2817 .storage
2818 .materialize(other_aligned.indices.len())?;
2819 let combined = lhs_storage
2820 .axpby(
2821 &a.to_backend_scalar(),
2822 rhs_storage.as_ref(),
2823 &b.to_backend_scalar(),
2824 )
2825 .map_err(|e| anyhow::anyhow!("storage axpby failed: {e}"))?;
2826 return Self::from_storage(self.indices.clone(), Arc::new(combined));
2827 }
2828
2829 let self_native = self.as_native()?;
2830 let other_native = other_aligned.as_native()?;
2831 let a_native = a.as_tensor()?.as_native()?;
2832 let b_native = b.as_tensor()?.as_native()?;
2833 if self_native.dtype() != other_native.dtype()
2834 || self_native.dtype() != a_native.dtype()
2835 || other_native.dtype() != b_native.dtype()
2836 {
2837 let combined = axpby_native_tensor(
2838 self_native,
2839 &a.to_backend_scalar(),
2840 other_native,
2841 &b.to_backend_scalar(),
2842 )?;
2843 return Self::from_native_with_axis_classes(
2844 self.indices.clone(),
2845 combined,
2846 axis_classes,
2847 );
2848 }
2849
2850 let lhs = self.scale(a)?;
2851 let rhs = other_aligned.scale(b)?;
2852 let combined = lhs
2853 .try_materialized_inner()?
2854 .add(rhs.try_materialized_inner()?)
2855 .map_err(|e| anyhow::anyhow!("tensor addition failed: {e}"))?;
2856 Self::from_inner_with_axis_classes(self.indices.clone(), combined, axis_classes)
2857 }
2858
2859 pub fn scale(&self, scalar: AnyScalar) -> Result<Self> {
2874 if !self.tracks_grad() && !scalar.tracks_grad() {
2875 let scaled = self.storage.scale(&scalar.to_backend_scalar())?;
2876 return Self::from_storage(self.indices.clone(), Arc::new(scaled));
2877 }
2878
2879 let self_native = self.as_native()?;
2880 let scalar_native = scalar.as_tensor()?.as_native()?;
2881 if self_native.dtype() != scalar_native.dtype() {
2882 let scaled = scale_native_tensor(self_native, &scalar.to_backend_scalar())?;
2883 return Self::from_native_with_axis_classes(
2884 self.indices.clone(),
2885 scaled,
2886 self.storage.axis_classes().to_vec(),
2887 );
2888 }
2889
2890 let scaled = if self.indices.is_empty() {
2891 self.try_materialized_inner()?
2892 .mul(scalar.as_tensor()?.try_materialized_inner()?)
2893 .map_err(|e| anyhow::anyhow!("scalar multiplication failed: {e}"))?
2894 } else {
2895 let subscripts = Self::scale_subscripts(self.indices.len())?;
2896 eager_einsum_ad(
2897 &[
2898 self.try_materialized_inner()?,
2899 scalar.as_tensor()?.try_materialized_inner()?,
2900 ],
2901 &subscripts,
2902 )
2903 .map_err(|e| anyhow::anyhow!("tensor scaling failed: {e}"))?
2904 };
2905 Self::from_inner_with_axis_classes(
2906 self.indices.clone(),
2907 scaled,
2908 self.storage.axis_classes().to_vec(),
2909 )
2910 }
2911
2912 pub fn inner_product(&self, other: &Self) -> Result<AnyScalar> {
2930 if self.indices.len() == other.indices.len() {
2931 let self_set: HashSet<_> = self.indices.iter().collect();
2932 let other_set: HashSet<_> = other.indices.iter().collect();
2933 if self_set == other_set {
2934 let other_aligned = other.permute_indices(&self.indices)?;
2935 let result = super::contract::contract_pair_with_operand_options(
2936 self,
2937 &other_aligned,
2938 PairwiseContractionOptions::new().with_lhs_conj(true),
2939 )?;
2940 return result.sum();
2941 }
2942 }
2943
2944 let result = super::contract::contract_pair_with_operand_options(
2946 self,
2947 other,
2948 PairwiseContractionOptions::new().with_lhs_conj(true),
2949 )?;
2950 result.sum()
2952 }
2953}
2954
2955impl TensorDynLen {
2960 pub fn replaceind(&self, old_index: &DynIndex, new_index: &DynIndex) -> Result<Self> {
2994 if old_index.dim() != new_index.dim() {
2996 return Err(anyhow::anyhow!(
2997 "Index space mismatch: cannot replace index with dimension {} with index of dimension {}",
2998 old_index.dim(),
2999 new_index.dim()
3000 ));
3001 }
3002
3003 let new_indices: Vec<_> = self
3004 .indices
3005 .iter()
3006 .map(|idx| {
3007 if *idx == *old_index {
3008 new_index.clone()
3009 } else {
3010 idx.clone()
3011 }
3012 })
3013 .collect();
3014
3015 Ok(Self {
3016 indices: new_indices,
3017 storage: self.storage.clone(),
3018 structured_ad: self.structured_ad.clone(),
3019 eager_cache: Arc::clone(&self.eager_cache),
3020 })
3021 }
3022
3023 pub fn replaceinds(&self, old_indices: &[DynIndex], new_indices: &[DynIndex]) -> Result<Self> {
3061 anyhow::ensure!(
3062 old_indices.len() == new_indices.len(),
3063 "old_indices and new_indices must have the same length"
3064 );
3065
3066 for (old, new) in old_indices.iter().zip(new_indices.iter()) {
3068 if old.dim() != new.dim() {
3069 return Err(anyhow::anyhow!(
3070 "Index space mismatch: cannot replace index with dimension {} with index of dimension {}",
3071 old.dim(),
3072 new.dim()
3073 ));
3074 }
3075 }
3076
3077 let replacement_map: std::collections::HashMap<_, _> =
3079 old_indices.iter().zip(new_indices.iter()).collect();
3080
3081 let new_indices_vec: Vec<_> = self
3082 .indices
3083 .iter()
3084 .map(|idx| {
3085 if let Some(new_idx) = replacement_map.get(idx) {
3086 (*new_idx).clone()
3087 } else {
3088 idx.clone()
3089 }
3090 })
3091 .collect();
3092
3093 Ok(Self {
3094 indices: new_indices_vec,
3095 storage: self.storage.clone(),
3096 structured_ad: self.structured_ad.clone(),
3097 eager_cache: Arc::clone(&self.eager_cache),
3098 })
3099 }
3100}
3101
3102impl TensorDynLen {
3107 pub fn conj(&self) -> Self {
3130 let new_indices: Vec<DynIndex> = self.indices.iter().map(|idx| idx.conj()).collect();
3134 let structured_ad = self.tracked_compact_payload_value().and_then(|value| {
3135 value.payload.conj().ok().map(|payload| {
3136 Arc::new(StructuredAdValue {
3137 payload: Arc::new(payload),
3138 payload_dims: value.payload_dims.clone(),
3139 axis_classes: value.axis_classes.clone(),
3140 })
3141 })
3142 });
3143 let eager_cache = self
3144 .eager_cache
3145 .get()
3146 .and_then(|inner| inner.conj().ok())
3147 .map(Self::eager_cache_with)
3148 .unwrap_or_else(Self::empty_eager_cache);
3149 Self {
3150 indices: new_indices,
3151 storage: self.storage.conj().unwrap_or_else(|_| {
3152 TensorDynLenStorage::from_storage(Arc::new(self.storage().conj()))
3153 }),
3154 structured_ad,
3155 eager_cache,
3156 }
3157 }
3158}
3159
3160impl TensorDynLen {
3165 pub fn norm_squared(&self) -> f64 {
3183 self.try_norm_squared().unwrap_or(f64::NAN)
3184 }
3185
3186 pub fn try_norm_squared(&self) -> Result<f64> {
3191 if self.indices.is_empty() {
3193 let value = self.sum()?;
3195 let abs_val = value.abs();
3196 return Ok(abs_val * abs_val);
3197 }
3198
3199 let conj = self.conj();
3202 let scalar = super::contract::contract_pair(self, &conj)?;
3203 Ok(scalar.sum()?.real().max(0.0))
3206 }
3207
3208 pub fn norm(&self) -> f64 {
3222 self.norm_squared().sqrt()
3223 }
3224
3225 pub fn maxabs(&self) -> f64 {
3237 self.storage.max_abs().unwrap_or(0.0)
3238 }
3239
3240 pub fn sub(&self, other: &Self) -> Result<Self> {
3248 self.axpby(AnyScalar::new_real(1.0), other, AnyScalar::new_real(-1.0))
3249 }
3250
3251 pub fn neg(&self) -> Result<Self> {
3256 self.scale(AnyScalar::new_real(-1.0))
3257 }
3258
3259 pub fn isapprox(&self, other: &Self, atol: f64, rtol: f64) -> bool {
3264 let diff = match self.sub(other) {
3265 Ok(d) => d,
3266 Err(_) => return false,
3267 };
3268 let diff_norm = diff.norm();
3269 diff_norm <= atol.max(rtol * self.norm().max(other.norm()))
3270 }
3271
3272 pub fn diagonal(input_index: &DynIndex, output_index: &DynIndex) -> Result<Self> {
3277 <Self as TensorConstructionLike>::diagonal(input_index, output_index)
3278 }
3279
3280 pub fn delta(input_indices: &[DynIndex], output_indices: &[DynIndex]) -> Result<Self> {
3286 <Self as TensorConstructionLike>::delta(input_indices, output_indices)
3287 }
3288
3289 pub fn scalar_one() -> Result<Self> {
3294 <Self as TensorConstructionLike>::scalar_one()
3295 }
3296
3297 pub fn ones(indices: &[DynIndex]) -> Result<Self> {
3302 <Self as TensorConstructionLike>::ones(indices)
3303 }
3304
3305 pub fn onehot(index_vals: &[(DynIndex, usize)]) -> Result<Self> {
3310 <Self as TensorConstructionLike>::onehot(index_vals)
3311 }
3312
3313 pub fn distance(&self, other: &Self) -> Result<f64> {
3344 let norm_self = self.norm();
3345
3346 let neg_other = other.scale(AnyScalar::new_real(-1.0))?;
3348 let diff = self.add(&neg_other)?;
3349 let norm_diff = diff.norm();
3350
3351 if norm_self > 0.0 {
3352 Ok(norm_diff / norm_self)
3353 } else {
3354 Ok(norm_diff)
3355 }
3356 }
3357}
3358
3359impl std::fmt::Debug for TensorDynLen {
3360 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3361 f.debug_struct("TensorDynLen")
3362 .field("indices", &self.indices)
3363 .field("dims", &self.dims())
3364 .field("is_diag", &self.is_diag())
3365 .finish()
3366 }
3367}
3368
3369pub fn diag_tensor_dyn_len(indices: Vec<DynIndex>, diag_data: Vec<f64>) -> Result<TensorDynLen> {
3394 TensorDynLen::from_diag(indices, diag_data)
3395}
3396
3397#[allow(clippy::type_complexity)]
3398pub(crate) type UnfoldSplitInnerResult = (
3399 EagerTensor,
3400 usize,
3401 usize,
3402 usize,
3403 Vec<DynIndex>,
3404 Vec<DynIndex>,
3405);
3406
3407#[allow(clippy::type_complexity)]
3454pub fn unfold_split(
3455 t: &TensorDynLen,
3456 left_inds: &[DynIndex],
3457) -> Result<(
3458 NativeTensor,
3459 usize,
3460 usize,
3461 usize,
3462 Vec<DynIndex>,
3463 Vec<DynIndex>,
3464)> {
3465 let (matrix_inner, left_len, m, n, left_indices, right_indices) =
3466 unfold_split_inner(t, left_inds)?;
3467
3468 Ok((
3469 matrix_inner.data().clone(),
3470 left_len,
3471 m,
3472 n,
3473 left_indices,
3474 right_indices,
3475 ))
3476}
3477
3478pub(crate) fn unfold_split_inner(
3479 t: &TensorDynLen,
3480 left_inds: &[DynIndex],
3481) -> Result<UnfoldSplitInnerResult> {
3482 let rank = t.indices.len();
3483
3484 anyhow::ensure!(rank >= 2, "Tensor must have rank >= 2, got rank {}", rank);
3486
3487 let left_len = left_inds.len();
3488
3489 anyhow::ensure!(
3491 left_len > 0 && left_len < rank,
3492 "Left indices must be a non-empty proper subset of tensor indices (0 < left_len < rank), got left_len={}, rank={}",
3493 left_len,
3494 rank
3495 );
3496
3497 let tensor_set: HashSet<_> = t.indices.iter().collect();
3499 let mut left_set = HashSet::new();
3500
3501 for left_idx in left_inds {
3502 anyhow::ensure!(
3503 tensor_set.contains(left_idx),
3504 "Index in left_inds not found in tensor"
3505 );
3506 anyhow::ensure!(left_set.insert(left_idx), "Duplicate index in left_inds");
3507 }
3508
3509 let mut right_inds = Vec::new();
3511 for idx in &t.indices {
3512 if !left_set.contains(idx) {
3513 right_inds.push(idx.clone());
3514 }
3515 }
3516
3517 let mut new_indices = Vec::with_capacity(rank);
3519 new_indices.extend_from_slice(left_inds);
3520 new_indices.extend_from_slice(&right_inds);
3521
3522 let unfolded = t.permute_indices(&new_indices)?;
3524
3525 let unfolded_dims = unfolded.dims();
3527 let m: usize = unfolded_dims[..left_len].iter().product();
3528 let n: usize = unfolded_dims[left_len..].iter().product();
3529
3530 let matrix_tensor = unfolded.try_materialized_inner()?.reshape(&[m, n])?;
3531
3532 Ok((
3533 matrix_tensor,
3534 left_len,
3535 m,
3536 n,
3537 left_inds.to_vec(),
3538 right_inds,
3539 ))
3540}
3541
3542use crate::tensor_index::TensorIndex;
3547
3548impl TensorIndex for TensorDynLen {
3549 type Index = DynIndex;
3550
3551 fn external_indices(&self) -> Vec<DynIndex> {
3552 self.indices.clone()
3554 }
3555
3556 fn num_external_indices(&self) -> usize {
3557 self.indices.len()
3558 }
3559
3560 fn replaceind(&self, old_index: &DynIndex, new_index: &DynIndex) -> Result<Self> {
3561 TensorDynLen::replaceind(self, old_index, new_index)
3563 }
3564
3565 fn replaceinds(&self, old_indices: &[DynIndex], new_indices: &[DynIndex]) -> Result<Self> {
3566 TensorDynLen::replaceinds(self, old_indices, new_indices)
3568 }
3569}
3570
3571use crate::tensor_like::{
3576 FactorizeError, FactorizeOptions, FactorizeResult, TensorConstructionLike,
3577 TensorContractionLike, TensorFactorizationLike, TensorVectorSpace,
3578};
3579
3580impl TensorVectorSpace for TensorDynLen {
3581 fn norm_squared(&self) -> f64 {
3582 TensorDynLen::norm_squared(self)
3583 }
3584
3585 fn maxabs(&self) -> f64 {
3586 TensorDynLen::maxabs(self)
3587 }
3588
3589 fn axpby(&self, a: crate::AnyScalar, other: &Self, b: crate::AnyScalar) -> Result<Self> {
3590 TensorDynLen::axpby(self, a, other, b)
3591 }
3592
3593 fn scale(&self, scalar: crate::AnyScalar) -> Result<Self> {
3594 TensorDynLen::scale(self, scalar)
3595 }
3596
3597 fn inner_product(&self, other: &Self) -> Result<crate::AnyScalar> {
3598 TensorDynLen::inner_product(self, other)
3599 }
3600}
3601
3602impl TensorFactorizationLike for TensorDynLen {
3603 fn factorize(
3604 &self,
3605 left_inds: &[DynIndex],
3606 options: &FactorizeOptions,
3607 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
3608 crate::factorize::factorize(self, left_inds, options)
3609 }
3610
3611 fn factorize_full_rank(
3612 &self,
3613 left_inds: &[DynIndex],
3614 alg: crate::FactorizeAlg,
3615 canonical: crate::Canonical,
3616 ) -> std::result::Result<FactorizeResult<Self>, FactorizeError> {
3617 crate::factorize::factorize_full_rank(self, left_inds, alg, canonical)
3618 }
3619}
3620
3621impl TensorContractionLike for TensorDynLen {
3622 fn conj(&self) -> Self {
3623 TensorDynLen::conj(self)
3625 }
3626
3627 fn direct_sum(
3628 &self,
3629 other: &Self,
3630 pairs: &[(DynIndex, DynIndex)],
3631 ) -> Result<crate::tensor_like::DirectSumResult<Self>> {
3632 let (tensor, new_indices) = crate::direct_sum::direct_sum(self, other, pairs)?;
3633 Ok(crate::tensor_like::DirectSumResult {
3634 tensor,
3635 new_indices,
3636 })
3637 }
3638
3639 fn outer_product(&self, other: &Self) -> Result<Self> {
3640 super::contract::outer_product(self, other)
3641 }
3642
3643 fn permuteinds(&self, new_order: &[DynIndex]) -> Result<Self> {
3644 TensorDynLen::permute_indices(self, new_order)
3646 }
3647
3648 fn fuse_indices(
3649 &self,
3650 old_indices: &[DynIndex],
3651 new_index: DynIndex,
3652 order: LinearizationOrder,
3653 ) -> Result<Self> {
3654 TensorDynLen::fuse_indices(self, old_indices, new_index, order)
3655 }
3656
3657 fn contract(tensors: &[&Self]) -> Result<Self> {
3658 super::contract::contract(tensors)
3659 }
3660
3661 fn contract_pair(&self, other: &Self) -> Result<Self> {
3662 super::contract::contract_pair(self, other)
3663 }
3664}
3665
3666impl TensorConstructionLike for TensorDynLen {
3667 fn select_indices(&self, selected_indices: &[DynIndex], positions: &[usize]) -> Result<Self> {
3668 TensorDynLen::select_indices(self, selected_indices, positions)
3669 }
3670
3671 fn diagonal(input_index: &DynIndex, output_index: &DynIndex) -> Result<Self> {
3672 let dim = input_index.dim();
3673 if dim != output_index.dim() {
3674 return Err(anyhow::anyhow!(
3675 "Dimension mismatch: input index has dim {}, output has dim {}",
3676 dim,
3677 output_index.dim(),
3678 ));
3679 }
3680
3681 TensorDynLen::from_diag(
3682 vec![input_index.clone(), output_index.clone()],
3683 vec![1.0_f64; dim],
3684 )
3685 }
3686
3687 fn scalar_one() -> Result<Self> {
3688 TensorDynLen::from_dense(vec![], vec![1.0_f64])
3689 }
3690
3691 fn ones(indices: &[DynIndex]) -> Result<Self> {
3692 if indices.is_empty() {
3693 return Self::scalar_one();
3694 }
3695 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
3696 let total_size = checked_total_size(&dims)?;
3697 TensorDynLen::from_dense(indices.to_vec(), vec![1.0_f64; total_size])
3698 }
3699
3700 fn onehot(index_vals: &[(DynIndex, usize)]) -> Result<Self> {
3701 if index_vals.is_empty() {
3702 return Self::scalar_one();
3703 }
3704 let indices: Vec<DynIndex> = index_vals.iter().map(|(idx, _)| idx.clone()).collect();
3705 let vals: Vec<usize> = index_vals.iter().map(|(_, v)| *v).collect();
3706 let dims: Vec<usize> = indices.iter().map(|idx| idx.size()).collect();
3707
3708 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
3709 if v >= d {
3710 return Err(anyhow::anyhow!(
3711 "onehot: value {} at position {} is >= dimension {}",
3712 v,
3713 k,
3714 d
3715 ));
3716 }
3717 }
3718
3719 let total_size = checked_total_size(&dims)?;
3720 let mut data = vec![0.0_f64; total_size];
3721
3722 let offset = column_major_offset(&dims, &vals)?;
3723 data[offset] = 1.0;
3724
3725 Self::from_dense(indices, data)
3726 }
3727
3728 }
3730
3731fn checked_total_size(dims: &[usize]) -> Result<usize> {
3732 dims.iter().try_fold(1_usize, |acc, &d| {
3733 if d == 0 {
3734 return Err(anyhow::anyhow!("invalid dimension 0"));
3735 }
3736 acc.checked_mul(d)
3737 .ok_or_else(|| anyhow::anyhow!("tensor size overflow"))
3738 })
3739}
3740
3741fn column_major_offset(dims: &[usize], vals: &[usize]) -> Result<usize> {
3742 if dims.len() != vals.len() {
3743 return Err(anyhow::anyhow!(
3744 "column_major_offset: dims.len() != vals.len()"
3745 ));
3746 }
3747 checked_total_size(dims)?;
3748
3749 let mut offset = 0usize;
3750 let mut stride = 1usize;
3751 for (k, (&v, &d)) in vals.iter().zip(dims.iter()).enumerate() {
3752 if d == 0 {
3753 return Err(anyhow::anyhow!("invalid dimension 0 at position {}", k));
3754 }
3755 if v >= d {
3756 return Err(anyhow::anyhow!(
3757 "column_major_offset: value {} at position {} is >= dimension {}",
3758 v,
3759 k,
3760 d
3761 ));
3762 }
3763 let term = v
3764 .checked_mul(stride)
3765 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
3766 offset = offset
3767 .checked_add(term)
3768 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
3769 stride = stride
3770 .checked_mul(d)
3771 .ok_or_else(|| anyhow::anyhow!("column_major_offset: overflow"))?;
3772 }
3773 Ok(offset)
3774}
3775
3776impl TensorDynLen {
3781 fn any_scalar_payload_to_complex(data: Vec<AnyScalar>) -> Vec<Complex64> {
3782 data.into_iter()
3783 .map(|value| {
3784 value
3785 .as_c64()
3786 .unwrap_or_else(|| Complex64::new(value.real(), 0.0))
3787 })
3788 .collect()
3789 }
3790
3791 fn any_scalar_payload_to_real(data: Vec<AnyScalar>) -> Vec<f64> {
3792 data.into_iter().map(|value| value.real()).collect()
3793 }
3794
3795 fn validate_dense_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
3796 let expected_len = checked_total_size(dims)?;
3797 anyhow::ensure!(
3798 data_len == expected_len,
3799 "dense payload length {} does not match dims {:?} (expected {})",
3800 data_len,
3801 dims,
3802 expected_len
3803 );
3804 Ok(())
3805 }
3806
3807 fn validate_diag_payload_len(data_len: usize, dims: &[usize]) -> Result<()> {
3808 anyhow::ensure!(
3809 !dims.is_empty(),
3810 "diagonal tensor construction requires at least one index"
3811 );
3812 Self::validate_diag_dims(dims)?;
3813 anyhow::ensure!(
3814 data_len == dims[0],
3815 "diagonal payload length {} does not match diagonal dimension {}",
3816 data_len,
3817 dims[0]
3818 );
3819 Ok(())
3820 }
3821
3822 pub fn from_dense<T: TensorElement>(indices: Vec<DynIndex>, data: Vec<T>) -> Result<Self> {
3849 let dims = Self::expected_dims_from_indices(&indices);
3850 Self::validate_indices(&indices)?;
3851 Self::validate_dense_payload_len(data.len(), &dims)?;
3852 let native = dense_native_tensor_from_col_major(&data, &dims)?;
3853 Self::from_native(indices, native)
3854 }
3855
3856 pub fn from_dense_any(indices: Vec<DynIndex>, data: Vec<AnyScalar>) -> Result<Self> {
3882 if data.iter().any(AnyScalar::is_complex) {
3883 Self::from_dense(indices, Self::any_scalar_payload_to_complex(data))
3884 } else {
3885 Self::from_dense(indices, Self::any_scalar_payload_to_real(data))
3886 }
3887 }
3888
3889 pub fn from_diag<T: TensorElement>(indices: Vec<DynIndex>, data: Vec<T>) -> Result<Self> {
3917 let dims = Self::expected_dims_from_indices(&indices);
3918 Self::validate_indices(&indices)?;
3919 Self::validate_diag_payload_len(data.len(), &dims)?;
3920 let native = diag_native_tensor_from_col_major(&data, dims.len())?;
3921 Self::from_native_with_axis_classes(indices, native, Self::diag_axis_classes(dims.len()))
3922 }
3923
3924 pub fn from_diag_any(indices: Vec<DynIndex>, data: Vec<AnyScalar>) -> Result<Self> {
3946 if data.iter().any(AnyScalar::is_complex) {
3947 Self::from_diag(indices, Self::any_scalar_payload_to_complex(data))
3948 } else {
3949 Self::from_diag(indices, Self::any_scalar_payload_to_real(data))
3950 }
3951 }
3952
3953 pub fn copy_tensor(indices: Vec<DynIndex>, value: AnyScalar) -> Result<Self> {
3974 if indices.is_empty() {
3975 return Self::from_dense_any(vec![], vec![value]);
3976 }
3977 let dim = indices[0].dim();
3978 let data = vec![value; dim];
3979 Self::from_diag_any(indices, data)
3980 }
3981
3982 pub fn fuse_indices(
4036 &self,
4037 old_indices: &[DynIndex],
4038 new_index: DynIndex,
4039 order: LinearizationOrder,
4040 ) -> Result<Self> {
4041 anyhow::ensure!(
4042 !old_indices.is_empty(),
4043 "fuse_indices requires at least one index to fuse"
4044 );
4045
4046 let old_dims = self.dims();
4047 let mut seen_indices = HashSet::new();
4048 let mut old_axes = Vec::with_capacity(old_indices.len());
4049 for old_index in old_indices {
4050 anyhow::ensure!(
4051 seen_indices.insert(old_index),
4052 "duplicate index in old_indices"
4053 );
4054 let axis = self
4055 .indices
4056 .iter()
4057 .position(|idx| idx == old_index)
4058 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
4059 anyhow::ensure!(
4060 old_index.dim() == old_dims[axis],
4061 "old index dimension does not match tensor axis dimension"
4062 );
4063 old_axes.push(axis);
4064 }
4065
4066 let fused_dims: Vec<usize> = old_axes.iter().map(|&axis| old_dims[axis]).collect();
4067 let fused_product = checked_product(&fused_dims)?;
4068 anyhow::ensure!(
4069 fused_product == new_index.dim(),
4070 "product of old index dimensions must match the replacement index dimension"
4071 );
4072
4073 let insertion_axis =
4074 old_axes.iter().copied().min().ok_or_else(|| {
4075 anyhow::anyhow!("fuse_indices requires at least one index to fuse")
4076 })?;
4077 let old_axis_set: HashSet<usize> = old_axes.iter().copied().collect();
4078
4079 let mut result_indices =
4080 Vec::with_capacity(self.indices.len() - old_indices.len() + 1usize);
4081 for (axis, index) in self.indices.iter().enumerate() {
4082 if axis == insertion_axis {
4083 result_indices.push(new_index.clone());
4084 }
4085 if !old_axis_set.contains(&axis) {
4086 result_indices.push(index.clone());
4087 }
4088 }
4089 let mut result_seen = HashSet::new();
4090 for index in &result_indices {
4091 anyhow::ensure!(
4092 result_seen.insert(index),
4093 "fuse_indices result would contain duplicate index"
4094 );
4095 }
4096 Self::validate_indices(&result_indices)?;
4097
4098 let mut new_dims = Vec::with_capacity(old_dims.len() - old_indices.len() + 1usize);
4099 for (axis, dim) in old_dims.iter().copied().enumerate() {
4100 if axis == insertion_axis {
4101 new_dims.push(new_index.dim());
4102 }
4103 if !old_axis_set.contains(&axis) {
4104 new_dims.push(dim);
4105 }
4106 }
4107
4108 self.ensure_shape_packing_preserves_ad("fuse_indices")?;
4109
4110 let mut grouped_axes = old_axes.clone();
4111 if matches!(order, LinearizationOrder::RowMajor) {
4112 grouped_axes.reverse();
4113 }
4114 let mut perm = Vec::with_capacity(self.indices.len());
4115 perm.extend((0..insertion_axis).filter(|axis| !old_axis_set.contains(axis)));
4116 perm.extend(grouped_axes);
4117 perm.extend(
4118 ((insertion_axis + 1)..self.indices.len()).filter(|axis| !old_axis_set.contains(axis)),
4119 );
4120 debug_assert_eq!(perm.len(), self.indices.len());
4121
4122 let packed = self.permute(&perm)?;
4123 let reshaped = packed.try_materialized_inner()?.reshape(&new_dims)?;
4124 Self::from_inner(result_indices, reshaped)
4125 }
4126
4127 pub fn unfuse_index(
4149 &self,
4150 old_index: &DynIndex,
4151 new_indices: &[DynIndex],
4152 order: LinearizationOrder,
4153 ) -> Result<Self> {
4154 anyhow::ensure!(
4155 !new_indices.is_empty(),
4156 "unfuse_index requires at least one replacement index"
4157 );
4158
4159 let axis = self
4160 .indices
4161 .iter()
4162 .position(|idx| idx == old_index)
4163 .ok_or_else(|| anyhow::anyhow!("index {:?} not found in tensor", old_index))?;
4164
4165 let replacement_dims: Vec<usize> = new_indices.iter().map(DynIndex::dim).collect();
4166 let replacement_product = checked_product(&replacement_dims)?;
4167 anyhow::ensure!(
4168 replacement_product == old_index.dim(),
4169 "product of new index dimensions must match the replaced index dimension"
4170 );
4171
4172 let mut result_indices =
4173 Vec::with_capacity(self.indices.len() - 1usize + new_indices.len());
4174 result_indices.extend_from_slice(&self.indices[..axis]);
4175 result_indices.extend(new_indices.iter().cloned());
4176 result_indices.extend_from_slice(&self.indices[axis + 1..]);
4177 Self::validate_indices(&result_indices)?;
4178
4179 let old_dims = self.dims();
4180 let mut new_dims = Vec::with_capacity(old_dims.len() - 1usize + replacement_dims.len());
4181 new_dims.extend_from_slice(&old_dims[..axis]);
4182 new_dims.extend_from_slice(&replacement_dims);
4183 new_dims.extend_from_slice(&old_dims[axis + 1..]);
4184
4185 self.ensure_shape_packing_preserves_ad("unfuse_index")?;
4186
4187 let mut grouped_indices = new_indices.to_vec();
4188 let mut grouped_dims = replacement_dims.clone();
4189 if matches!(order, LinearizationOrder::RowMajor) {
4190 grouped_indices.reverse();
4191 grouped_dims.reverse();
4192 }
4193 let mut packed_indices =
4194 Vec::with_capacity(self.indices.len() - 1usize + grouped_indices.len());
4195 packed_indices.extend_from_slice(&self.indices[..axis]);
4196 packed_indices.extend(grouped_indices);
4197 packed_indices.extend_from_slice(&self.indices[axis + 1..]);
4198
4199 let mut packed_dims = Vec::with_capacity(old_dims.len() - 1usize + grouped_dims.len());
4200 packed_dims.extend_from_slice(&old_dims[..axis]);
4201 packed_dims.extend_from_slice(&grouped_dims);
4202 packed_dims.extend_from_slice(&old_dims[axis + 1..]);
4203
4204 let reshaped = self.try_materialized_inner()?.reshape(&packed_dims)?;
4205 let packed = Self::from_inner(packed_indices, reshaped)?;
4206 if matches!(order, LinearizationOrder::ColumnMajor) {
4207 Ok(packed)
4208 } else {
4209 packed.permute_indices(&result_indices)
4210 }
4211 }
4212
4213 pub fn scalar<T: TensorElement>(value: T) -> Result<Self> {
4224 Self::from_dense(vec![], vec![value])
4225 }
4226
4227 pub fn zeros<T: TensorElement + Zero + Clone>(indices: Vec<DynIndex>) -> Result<Self> {
4240 let dims: Vec<usize> = indices.iter().map(|idx| idx.dim()).collect();
4241 let size: usize = dims.iter().product();
4242 Self::from_dense(indices, vec![T::zero(); size])
4243 }
4244}
4245
4246impl TensorDynLen {
4251 pub fn to_vec<T: TensorElement>(&self) -> Result<Vec<T>> {
4273 native_tensor_primal_to_dense_col_major(self.as_native()?)
4274 }
4275
4276 pub fn into_dense_col_major_parts<T: TensorElement>(self) -> Result<(Vec<DynIndex>, Vec<T>)> {
4312 anyhow::ensure!(
4313 self.structured_ad.is_none() && !self.tracks_grad(),
4314 "TensorDynLen::into_dense_col_major_parts cannot consume tensors with tracked autodiff state"
4315 );
4316 let data = self.to_vec::<T>()?;
4317 Ok((self.indices, data))
4318 }
4319
4320 pub fn as_slice_f64(&self) -> Result<Vec<f64>> {
4325 self.to_vec::<f64>()
4326 }
4327
4328 pub fn as_slice_c64(&self) -> Result<Vec<Complex64>> {
4333 self.to_vec::<Complex64>()
4334 }
4335
4336 pub fn is_f64(&self) -> bool {
4349 self.storage.is_f64()
4350 }
4351
4352 pub fn is_diag(&self) -> bool {
4379 self.storage.is_diag()
4380 }
4381
4382 pub fn is_complex(&self) -> bool {
4401 self.storage.is_complex()
4402 }
4403}
4404
4405fn checked_product(dims: &[usize]) -> Result<usize> {
4406 dims.iter().try_fold(1usize, |acc, &dim| {
4407 acc.checked_mul(dim)
4408 .ok_or_else(|| anyhow::anyhow!("dimension product overflow"))
4409 })
4410}
4411
4412fn decode_col_major_linear(linear: usize, dims: &[usize]) -> Result<Vec<usize>> {
4413 let total = checked_product(dims)?;
4414 anyhow::ensure!(
4415 linear < total,
4416 "linear offset {} out of bounds for dims {:?}",
4417 linear,
4418 dims
4419 );
4420 let mut remaining = linear;
4421 let mut out = Vec::with_capacity(dims.len());
4422 for &dim in dims {
4423 out.push(remaining % dim);
4424 remaining /= dim;
4425 }
4426 Ok(out)
4427}
4428
4429fn encode_col_major_linear(indices: &[usize], dims: &[usize]) -> Result<usize> {
4430 anyhow::ensure!(
4431 indices.len() == dims.len(),
4432 "index rank {} does not match dims {:?}",
4433 indices.len(),
4434 dims
4435 );
4436 let mut linear = 0usize;
4437 let mut stride = 1usize;
4438 for (&index, &dim) in indices.iter().zip(dims.iter()) {
4439 anyhow::ensure!(
4440 index < dim,
4441 "index {} out of bounds for dimension {}",
4442 index,
4443 dim
4444 );
4445 linear += index * stride;
4446 stride = stride
4447 .checked_mul(dim)
4448 .ok_or_else(|| anyhow::anyhow!("stride overflow"))?;
4449 }
4450 Ok(linear)
4451}