Skip to main content

ft_core/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::fmt;
4use std::hash::{Hash, Hasher};
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7
8/// Half-precision float types re-exported from the `half` crate.
9/// Named `Float16`/`BFloat16` to avoid conflict with Rust 2024 primitive `f16`.
10pub type Float16 = half::f16;
11pub type BFloat16 = half::bf16;
12
13/// Complex number types re-exported from the `num_complex` crate.
14pub type Complex64 = num_complex::Complex<f32>;
15pub type Complex128 = num_complex::Complex<f64>;
16
17static NEXT_TENSOR_ID: AtomicU64 = AtomicU64::new(1);
18static NEXT_STORAGE_ID: AtomicU64 = AtomicU64::new(1);
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub enum DType {
22    F64,
23    F32,
24    F16,
25    BF16,
26    QInt8,
27    QUInt8,
28    I64,
29    I32,
30    Bool,
31    Complex64,
32    Complex128,
33}
34
35impl DType {
36    /// Size of one element in bytes.
37    #[must_use]
38    pub fn element_size(self) -> usize {
39        match self {
40            Self::Complex128 => 16,
41            Self::F64 | Self::I64 | Self::Complex64 => 8,
42            Self::F32 | Self::I32 => 4,
43            Self::F16 | Self::BF16 => 2,
44            Self::QInt8 | Self::QUInt8 | Self::Bool => 1,
45        }
46    }
47
48    /// Returns true for floating-point dtypes.
49    #[must_use]
50    pub fn is_floating_point(self) -> bool {
51        matches!(self, Self::F64 | Self::F32 | Self::F16 | Self::BF16)
52    }
53
54    /// Returns true for half-precision floating-point dtypes (F16 or BF16).
55    #[must_use]
56    pub fn is_half(self) -> bool {
57        matches!(self, Self::F16 | Self::BF16)
58    }
59
60    /// Returns true for integer dtypes (not bool).
61    #[must_use]
62    pub fn is_integer(self) -> bool {
63        matches!(self, Self::I32 | Self::I64)
64    }
65
66    /// Returns true for quantized storage dtypes.
67    #[must_use]
68    pub fn is_quantized(self) -> bool {
69        matches!(self, Self::QInt8 | Self::QUInt8)
70    }
71
72    /// Returns true for the boolean dtype.
73    #[must_use]
74    pub fn is_bool(self) -> bool {
75        matches!(self, Self::Bool)
76    }
77
78    /// Returns true for complex dtypes (Complex64 or Complex128).
79    #[must_use]
80    pub fn is_complex(self) -> bool {
81        matches!(self, Self::Complex64 | Self::Complex128)
82    }
83
84    /// Promote two floating-point or complex dtypes: F32+F64→F64, same→same.
85    /// Half-precision types promote to F32. F16+BF16→F32.
86    /// Complex types: Complex64+F32→Complex64, Complex64+F64→Complex128, etc.
87    /// Returns `None` for non-floating-point/non-complex dtypes.
88    #[must_use]
89    pub fn promote(self, other: Self) -> Option<Self> {
90        match (self, other) {
91            // Complex + Complex
92            (Self::Complex128, Self::Complex128) => Some(Self::Complex128),
93            (Self::Complex64, Self::Complex64) => Some(Self::Complex64),
94            (Self::Complex128, Self::Complex64) | (Self::Complex64, Self::Complex128) => {
95                Some(Self::Complex128)
96            }
97            // Complex + real float → complex (widen component if needed)
98            (Self::Complex128, Self::F64 | Self::F32 | Self::F16 | Self::BF16)
99            | (Self::F64 | Self::F32 | Self::F16 | Self::BF16, Self::Complex128) => {
100                Some(Self::Complex128)
101            }
102            (Self::Complex64, Self::F64) | (Self::F64, Self::Complex64) => Some(Self::Complex128),
103            (Self::Complex64, Self::F32 | Self::F16 | Self::BF16)
104            | (Self::F32 | Self::F16 | Self::BF16, Self::Complex64) => Some(Self::Complex64),
105            // Real floats
106            (Self::F64, Self::F64) => Some(Self::F64),
107            (Self::F32, Self::F32) => Some(Self::F32),
108            (Self::F64, Self::F32) | (Self::F32, Self::F64) => Some(Self::F64),
109            // Half-precision: same type stays, mixed half → F32
110            (Self::F16, Self::F16) => Some(Self::F16),
111            (Self::BF16, Self::BF16) => Some(Self::BF16),
112            (Self::F16, Self::BF16) | (Self::BF16, Self::F16) => Some(Self::F32),
113            // Half + wider float → wider float
114            (Self::F16 | Self::BF16, Self::F32) | (Self::F32, Self::F16 | Self::BF16) => {
115                Some(Self::F32)
116            }
117            (Self::F16 | Self::BF16, Self::F64) | (Self::F64, Self::F16 | Self::BF16) => {
118                Some(Self::F64)
119            }
120            _ => None,
121        }
122    }
123
124    /// Promote two dtypes following PyTorch's promotion hierarchy:
125    /// Bool → I32 → I64 → F16/BF16 → F32 → F64 → Complex64 → Complex128.
126    ///
127    /// Any pair of dtypes returns the wider type in this hierarchy.
128    /// Int + Float always promotes to the float type (or wider float).
129    /// F16 + BF16 promotes to F32 (matching PyTorch semantics).
130    /// Real + Complex promotes to Complex (widening component type if needed).
131    /// This matches PyTorch's `torch.promote_types()`.
132    #[must_use]
133    pub fn promote_types(self, other: Self) -> Self {
134        if self == other {
135            return self;
136        }
137        // Special case: F16 + BF16 → F32 (PyTorch semantics)
138        if matches!(
139            (self, other),
140            (Self::F16, Self::BF16) | (Self::BF16, Self::F16)
141        ) {
142            return Self::F32;
143        }
144        // Special case: Complex64 + F64 → Complex128 (widen component)
145        if matches!(
146            (self, other),
147            (Self::Complex64, Self::F64) | (Self::F64, Self::Complex64)
148        ) {
149            return Self::Complex128;
150        }
151        // Assign a rank following PyTorch's promotion hierarchy.
152        let rank = |d: Self| -> u8 {
153            match d {
154                Self::Bool => 0,
155                Self::QInt8 | Self::QUInt8 => 1,
156                Self::I32 => 2,
157                Self::I64 => 3,
158                Self::F16 | Self::BF16 => 4,
159                Self::F32 => 5,
160                Self::F64 => 6,
161                Self::Complex64 => 7,
162                Self::Complex128 => 8,
163            }
164        };
165        if rank(self) >= rank(other) {
166            self
167        } else {
168            other
169        }
170    }
171}
172
173#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
174pub enum Device {
175    Cpu,
176    Cuda,
177}
178
179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
180pub enum ExecutionMode {
181    Strict,
182    Hardened,
183}
184
185#[derive(Debug, Clone, PartialEq, Eq, Hash)]
186pub struct QuantizationParams {
187    scale_bits: Vec<u64>,
188    zero_points: Vec<i64>,
189    axis: Option<usize>,
190}
191
192impl QuantizationParams {
193    /// Create affine quantization parameters.
194    ///
195    /// `scale` must be finite and strictly positive. It is stored by bit pattern
196    /// so tensor metadata remains equality/hash stable.
197    pub fn new(scale: f64, zero_point: i64) -> Result<Self, TensorMetaError> {
198        let scale_bits = Self::validate_scale(scale)?;
199        Ok(Self {
200            scale_bits: vec![scale_bits],
201            zero_points: vec![zero_point],
202            axis: None,
203        })
204    }
205
206    /// Create per-channel affine quantization parameters.
207    ///
208    /// `axis` identifies the quantized channel dimension in the tensor shape.
209    /// `scales` and `zero_points` must have the same non-zero length.
210    pub fn per_channel(
211        scales: Vec<f64>,
212        zero_points: Vec<i64>,
213        axis: usize,
214    ) -> Result<Self, TensorMetaError> {
215        if scales.is_empty() {
216            return Err(TensorMetaError::EmptyQuantizationChannels);
217        }
218        if scales.len() != zero_points.len() {
219            return Err(TensorMetaError::QuantizationVectorLengthMismatch {
220                scales: scales.len(),
221                zero_points: zero_points.len(),
222            });
223        }
224
225        let mut scale_bits = Vec::with_capacity(scales.len());
226        for scale in scales {
227            scale_bits.push(Self::validate_scale(scale)?);
228        }
229
230        Ok(Self {
231            scale_bits,
232            zero_points,
233            axis: Some(axis),
234        })
235    }
236
237    fn validate_scale(scale: f64) -> Result<u64, TensorMetaError> {
238        if !scale.is_finite() || scale <= 0.0 {
239            return Err(TensorMetaError::InvalidQuantizationScale {
240                scale_bits: scale.to_bits(),
241            });
242        }
243        Ok(scale.to_bits())
244    }
245
246    #[must_use]
247    pub fn scale(&self) -> f64 {
248        self.scale_at(0).unwrap_or(f64::NAN)
249    }
250
251    #[must_use]
252    pub fn zero_point(&self) -> i64 {
253        self.zero_point_at(0).unwrap_or_default()
254    }
255
256    fn scale_at(&self, channel: usize) -> Option<f64> {
257        self.scale_bits
258            .get(channel)
259            .map(|bits| f64::from_bits(*bits))
260    }
261
262    fn zero_point_at(&self, channel: usize) -> Option<i64> {
263        self.zero_points.get(channel).copied()
264    }
265
266    #[must_use]
267    pub fn scales(&self) -> Vec<f64> {
268        self.scale_bits
269            .iter()
270            .map(|&bits| f64::from_bits(bits))
271            .collect()
272    }
273
274    #[must_use]
275    pub fn zero_points(&self) -> &[i64] {
276        &self.zero_points
277    }
278
279    #[must_use]
280    pub fn axis(&self) -> Option<usize> {
281        self.axis
282    }
283
284    #[must_use]
285    pub fn len(&self) -> usize {
286        self.scale_bits.len()
287    }
288
289    #[must_use]
290    pub fn is_empty(&self) -> bool {
291        self.scale_bits.is_empty()
292    }
293
294    fn validate_for_shape(&self, shape: &[usize]) -> Result<(), TensorMetaError> {
295        if self.scale_bits.is_empty() {
296            return Err(TensorMetaError::EmptyQuantizationChannels);
297        }
298        if self.scale_bits.len() != self.zero_points.len() {
299            return Err(TensorMetaError::QuantizationVectorLengthMismatch {
300                scales: self.scale_bits.len(),
301                zero_points: self.zero_points.len(),
302            });
303        }
304
305        if let Some(axis) = self.axis {
306            if axis >= shape.len() {
307                return Err(TensorMetaError::InvalidQuantizationAxis {
308                    axis,
309                    rank: shape.len(),
310                });
311            }
312            let Some(&expected) = shape.get(axis) else {
313                return Err(TensorMetaError::InvalidQuantizationAxis {
314                    axis,
315                    rank: shape.len(),
316                });
317            };
318            let actual = self.scale_bits.len();
319            if expected != actual {
320                return Err(TensorMetaError::QuantizationChannelCountMismatch {
321                    axis,
322                    expected,
323                    actual,
324                });
325            }
326        } else if self.scale_bits.len() != 1 {
327            return Err(TensorMetaError::QuantizationVectorLengthMismatch {
328                scales: self.scale_bits.len(),
329                zero_points: self.zero_points.len(),
330            });
331        }
332
333        Ok(())
334    }
335}
336
337#[derive(Debug, Clone, PartialEq, Eq)]
338pub struct TensorMeta {
339    shape: Vec<usize>,
340    strides: Vec<usize>,
341    numel: usize,
342    storage_offset: usize,
343    dtype: DType,
344    device: Device,
345    quantization: Option<QuantizationParams>,
346}
347
348impl TensorMeta {
349    #[must_use]
350    pub fn scalar(dtype: DType, device: Device) -> Self {
351        Self {
352            shape: Vec::new(),
353            strides: Vec::new(),
354            numel: 1,
355            storage_offset: 0,
356            dtype,
357            device,
358            quantization: None,
359        }
360    }
361
362    #[must_use]
363    pub fn from_shape(shape: Vec<usize>, dtype: DType, device: Device) -> Self {
364        let strides = contiguous_strides(&shape);
365        let numel = saturated_numel(&shape);
366        Self {
367            shape,
368            strides,
369            numel,
370            storage_offset: 0,
371            dtype,
372            device,
373            quantization: None,
374        }
375    }
376
377    pub fn quantized_from_shape(
378        shape: Vec<usize>,
379        dtype: DType,
380        device: Device,
381        scale: f64,
382        zero_point: i64,
383    ) -> Result<Self, TensorMetaError> {
384        let quantization = QuantizationParams::new(scale, zero_point)?;
385        let meta = Self::from_shape(shape, dtype, device).with_quantization(quantization);
386        meta.validate()?;
387        Ok(meta)
388    }
389
390    pub fn quantized_per_channel_from_shape(
391        shape: Vec<usize>,
392        dtype: DType,
393        device: Device,
394        scales: Vec<f64>,
395        zero_points: Vec<i64>,
396        axis: usize,
397    ) -> Result<Self, TensorMetaError> {
398        let quantization = QuantizationParams::per_channel(scales, zero_points, axis)?;
399        let meta = Self::from_shape(shape, dtype, device).with_quantization(quantization);
400        meta.validate()?;
401        Ok(meta)
402    }
403
404    pub fn from_shape_and_strides(
405        shape: Vec<usize>,
406        strides: Vec<usize>,
407        storage_offset: usize,
408        dtype: DType,
409        device: Device,
410    ) -> Result<Self, TensorMetaError> {
411        let numel = saturated_numel(&shape);
412        let meta = Self {
413            shape,
414            strides,
415            numel,
416            storage_offset,
417            dtype,
418            device,
419            quantization: None,
420        };
421        meta.validate()?;
422        Ok(meta)
423    }
424
425    pub fn quantized_from_shape_and_strides(
426        shape: Vec<usize>,
427        strides: Vec<usize>,
428        storage_offset: usize,
429        dtype: DType,
430        device: Device,
431        scale: f64,
432        zero_point: i64,
433    ) -> Result<Self, TensorMetaError> {
434        let quantization = QuantizationParams::new(scale, zero_point)?;
435        let numel = saturated_numel(&shape);
436        let meta = Self {
437            shape,
438            strides,
439            numel,
440            storage_offset,
441            dtype,
442            device,
443            quantization: Some(quantization),
444        };
445        meta.validate()?;
446        Ok(meta)
447    }
448
449    #[allow(clippy::too_many_arguments)]
450    pub fn quantized_per_channel_from_shape_and_strides(
451        shape: Vec<usize>,
452        strides: Vec<usize>,
453        storage_offset: usize,
454        dtype: DType,
455        device: Device,
456        scales: Vec<f64>,
457        zero_points: Vec<i64>,
458        axis: usize,
459    ) -> Result<Self, TensorMetaError> {
460        let quantization = QuantizationParams::per_channel(scales, zero_points, axis)?;
461        let numel = saturated_numel(&shape);
462        let meta = Self {
463            shape,
464            strides,
465            numel,
466            storage_offset,
467            dtype,
468            device,
469            quantization: Some(quantization),
470        };
471        meta.validate()?;
472        Ok(meta)
473    }
474
475    #[must_use]
476    pub fn with_storage_offset(mut self, storage_offset: usize) -> Self {
477        self.storage_offset = storage_offset;
478        self
479    }
480
481    #[must_use]
482    pub fn with_dtype(mut self, dtype: DType) -> Self {
483        self.dtype = dtype;
484        if !dtype.is_quantized() {
485            self.quantization = None;
486        }
487        self
488    }
489
490    #[must_use]
491    pub fn with_quantization(mut self, quantization: QuantizationParams) -> Self {
492        self.quantization = Some(quantization);
493        self
494    }
495
496    pub fn validate(&self) -> Result<(), TensorMetaError> {
497        match (self.dtype.is_quantized(), self.quantization.as_ref()) {
498            (true, Some(quantization)) => quantization.validate_for_shape(&self.shape)?,
499            (true, None) => {
500                return Err(TensorMetaError::MissingQuantizationParams { dtype: self.dtype });
501            }
502            (false, Some(_)) => {
503                return Err(TensorMetaError::UnexpectedQuantizationParams { dtype: self.dtype });
504            }
505            (false, None) => {}
506        }
507
508        if self.shape.len() != self.strides.len() {
509            return Err(TensorMetaError::RankStrideMismatch {
510                rank: self.shape.len(),
511                strides: self.strides.len(),
512            });
513        }
514
515        let mut max_linear_offset = 0usize;
516        for (size, stride) in self.shape.iter().copied().zip(self.strides.iter().copied()) {
517            if size == 0 {
518                continue;
519            }
520
521            let span = stride
522                .checked_mul(size.saturating_sub(1))
523                .ok_or(TensorMetaError::StrideOverflow { size, stride })?;
524            max_linear_offset = max_linear_offset.checked_add(span).ok_or(
525                TensorMetaError::StorageOffsetOverflow {
526                    storage_offset: self.storage_offset,
527                    max_linear_offset,
528                },
529            )?;
530        }
531
532        let _ = self.storage_offset.checked_add(max_linear_offset).ok_or(
533            TensorMetaError::StorageOffsetOverflow {
534                storage_offset: self.storage_offset,
535                max_linear_offset,
536            },
537        )?;
538
539        Ok(())
540    }
541
542    #[must_use]
543    pub fn shape(&self) -> &[usize] {
544        &self.shape
545    }
546
547    #[must_use]
548    pub fn strides(&self) -> &[usize] {
549        &self.strides
550    }
551
552    #[must_use]
553    pub fn storage_offset(&self) -> usize {
554        self.storage_offset
555    }
556
557    #[must_use]
558    pub fn dtype(&self) -> DType {
559        self.dtype
560    }
561
562    #[must_use]
563    pub fn device(&self) -> Device {
564        self.device
565    }
566
567    #[must_use]
568    pub fn quantization(&self) -> Option<&QuantizationParams> {
569        self.quantization.as_ref()
570    }
571
572    #[must_use]
573    pub fn numel(&self) -> usize {
574        self.numel
575    }
576
577    #[must_use]
578    pub fn is_contiguous(&self) -> bool {
579        if self.shape.len() != self.strides.len() {
580            return false;
581        }
582
583        let mut expected_stride = 1usize;
584        for (size, stride) in self
585            .shape
586            .iter()
587            .copied()
588            .zip(self.strides.iter().copied())
589            .rev()
590        {
591            // Match PyTorch semantics: singleton dimensions are contiguous
592            // regardless of stride.
593            if size == 1 {
594                continue;
595            }
596            if stride != expected_stride {
597                return false;
598            }
599            let Some(next_expected) = expected_stride.checked_mul(size) else {
600                return false;
601            };
602            expected_stride = next_expected;
603        }
604        true
605    }
606
607    pub fn storage_index_for(&self, index: &[usize]) -> Result<usize, TensorMetaError> {
608        if index.len() != self.shape.len() {
609            return Err(TensorMetaError::IndexRankMismatch {
610                expected: self.shape.len(),
611                actual: index.len(),
612            });
613        }
614
615        let mut linear = self.storage_offset;
616        for (dim, ((idx, dim_size), stride)) in index
617            .iter()
618            .copied()
619            .zip(self.shape.iter().copied())
620            .zip(self.strides.iter().copied())
621            .enumerate()
622        {
623            if idx >= dim_size {
624                return Err(TensorMetaError::IndexOutOfBounds {
625                    dim,
626                    index: idx,
627                    size: dim_size,
628                });
629            }
630
631            let step = idx
632                .checked_mul(stride)
633                .ok_or(TensorMetaError::StrideOverflow { size: idx, stride })?;
634            linear = linear
635                .checked_add(step)
636                .ok_or(TensorMetaError::StorageOffsetOverflow {
637                    storage_offset: self.storage_offset,
638                    max_linear_offset: step,
639                })?;
640        }
641
642        Ok(linear)
643    }
644
645    #[must_use]
646    pub fn fingerprint64(&self) -> u64 {
647        let mut hasher = DetHasher::new();
648        self.shape.hash(&mut hasher);
649        self.strides.hash(&mut hasher);
650        self.storage_offset.hash(&mut hasher);
651        self.dtype.hash(&mut hasher);
652        self.device.hash(&mut hasher);
653        self.quantization.hash(&mut hasher);
654        hasher.finish()
655    }
656}
657
658struct DetHasher(u64);
659
660impl DetHasher {
661    fn new() -> Self {
662        Self(0xcbf2_9ce4_8422_2325)
663    }
664}
665
666impl Hasher for DetHasher {
667    fn finish(&self) -> u64 {
668        self.0
669    }
670
671    fn write(&mut self, bytes: &[u8]) {
672        for &byte in bytes {
673            self.0 ^= u64::from(byte);
674            self.0 = self.0.wrapping_mul(0x0000_0100_0000_01b3);
675        }
676    }
677}
678
679#[derive(Debug, Clone, Copy, PartialEq, Eq)]
680pub enum TensorMetaError {
681    RankStrideMismatch {
682        rank: usize,
683        strides: usize,
684    },
685    StrideOverflow {
686        size: usize,
687        stride: usize,
688    },
689    StorageOffsetOverflow {
690        storage_offset: usize,
691        max_linear_offset: usize,
692    },
693    IndexRankMismatch {
694        expected: usize,
695        actual: usize,
696    },
697    IndexOutOfBounds {
698        dim: usize,
699        index: usize,
700        size: usize,
701    },
702    MissingQuantizationParams {
703        dtype: DType,
704    },
705    UnexpectedQuantizationParams {
706        dtype: DType,
707    },
708    InvalidQuantizationScale {
709        scale_bits: u64,
710    },
711    EmptyQuantizationChannels,
712    QuantizationVectorLengthMismatch {
713        scales: usize,
714        zero_points: usize,
715    },
716    InvalidQuantizationAxis {
717        axis: usize,
718        rank: usize,
719    },
720    QuantizationChannelCountMismatch {
721        axis: usize,
722        expected: usize,
723        actual: usize,
724    },
725}
726
727impl fmt::Display for TensorMetaError {
728    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
729        match self {
730            Self::RankStrideMismatch { rank, strides } => {
731                write!(f, "shape rank {rank} does not match strides rank {strides}")
732            }
733            Self::StrideOverflow { size, stride } => {
734                write!(f, "stride overflow for size={size}, stride={stride}")
735            }
736            Self::StorageOffsetOverflow {
737                storage_offset,
738                max_linear_offset,
739            } => write!(
740                f,
741                "storage offset overflow for storage_offset={storage_offset}, max_linear_offset={max_linear_offset}"
742            ),
743            Self::IndexRankMismatch { expected, actual } => {
744                write!(
745                    f,
746                    "index rank mismatch expected={expected}, actual={actual}"
747                )
748            }
749            Self::IndexOutOfBounds { dim, index, size } => {
750                write!(
751                    f,
752                    "index out of bounds at dim={dim}: index={index}, size={size}"
753                )
754            }
755            Self::MissingQuantizationParams { dtype } => {
756                write!(
757                    f,
758                    "quantized dtype {dtype:?} requires quantization metadata"
759                )
760            }
761            Self::UnexpectedQuantizationParams { dtype } => {
762                write!(
763                    f,
764                    "non-quantized dtype {dtype:?} cannot carry quantization metadata"
765                )
766            }
767            Self::InvalidQuantizationScale { scale_bits } => {
768                write!(
769                    f,
770                    "quantization scale must be finite and > 0: bits={scale_bits:#x}"
771                )
772            }
773            Self::EmptyQuantizationChannels => {
774                write!(f, "per-channel quantization requires at least one channel")
775            }
776            Self::QuantizationVectorLengthMismatch {
777                scales,
778                zero_points,
779            } => write!(
780                f,
781                "quantization scales/zero_points length mismatch: scales={scales}, zero_points={zero_points}"
782            ),
783            Self::InvalidQuantizationAxis { axis, rank } => {
784                write!(
785                    f,
786                    "quantization axis out of range: axis={axis}, rank={rank}"
787                )
788            }
789            Self::QuantizationChannelCountMismatch {
790                axis,
791                expected,
792                actual,
793            } => write!(
794                f,
795                "quantization channel count mismatch at axis={axis}: expected={expected}, actual={actual}"
796            ),
797        }
798    }
799}
800
801impl std::error::Error for TensorMetaError {}
802
803#[derive(Debug, Clone, Copy, PartialEq, Eq)]
804pub enum TensorCompatError {
805    DTypeMismatch { lhs: DType, rhs: DType },
806    DeviceMismatch { lhs: Device, rhs: Device },
807}
808
809impl fmt::Display for TensorCompatError {
810    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
811        match self {
812            Self::DTypeMismatch { lhs, rhs } => {
813                write!(f, "dtype mismatch: lhs={lhs:?}, rhs={rhs:?}")
814            }
815            Self::DeviceMismatch { lhs, rhs } => {
816                write!(f, "device mismatch: lhs={lhs:?}, rhs={rhs:?}")
817            }
818        }
819    }
820}
821
822impl std::error::Error for TensorCompatError {}
823
824#[derive(Debug, Clone, PartialEq)]
825pub struct ScalarTensor {
826    id: u64,
827    storage_id: u64,
828    meta: TensorMeta,
829    value: f64,
830    version: u64,
831}
832
833impl ScalarTensor {
834    #[must_use]
835    pub fn new(value: f64, dtype: DType, device: Device) -> Self {
836        Self {
837            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
838            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
839            meta: TensorMeta::scalar(dtype, device),
840            value,
841            version: 0,
842        }
843    }
844
845    #[must_use]
846    pub fn with_value(&self, value: f64) -> Self {
847        Self {
848            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
849            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
850            meta: self.meta.clone(),
851            value,
852            version: self.version.saturating_add(1),
853        }
854    }
855
856    pub fn alias_view(&self, storage_offset: usize) -> Result<Self, TensorMetaError> {
857        let meta = self.meta.clone().with_storage_offset(storage_offset);
858        meta.validate()?;
859        Ok(Self {
860            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
861            storage_id: self.storage_id,
862            meta,
863            value: self.value,
864            version: self.version,
865        })
866    }
867
868    pub fn set_in_place(&mut self, value: f64) {
869        self.value = value;
870        self.version = self.version.saturating_add(1);
871    }
872
873    #[must_use]
874    pub fn id(&self) -> u64 {
875        self.id
876    }
877
878    #[must_use]
879    pub fn storage_id(&self) -> u64 {
880        self.storage_id
881    }
882
883    #[must_use]
884    pub fn value(&self) -> f64 {
885        self.value
886    }
887
888    #[must_use]
889    pub fn meta(&self) -> &TensorMeta {
890        &self.meta
891    }
892
893    #[must_use]
894    pub fn version(&self) -> u64 {
895        self.version
896    }
897
898    #[must_use]
899    pub fn evidence_fingerprint64(&self) -> u64 {
900        let mut hasher = DetHasher::new();
901        self.id.hash(&mut hasher);
902        self.storage_id.hash(&mut hasher);
903        self.version.hash(&mut hasher);
904        self.meta.fingerprint64().hash(&mut hasher);
905        self.value.to_bits().hash(&mut hasher);
906        hasher.finish()
907    }
908}
909
910// ── Typed Tensor Storage ────────────────────────────────────────────────
911
912#[derive(Debug, Clone, PartialEq)]
913pub enum TensorStorage {
914    F32(Arc<Vec<f32>>),
915    F64(Arc<Vec<f64>>),
916    F64Inline4([f64; 4]),
917    F16(Arc<Vec<Float16>>),
918    BF16(Arc<Vec<BFloat16>>),
919    QInt8(Arc<Vec<i8>>),
920    QUInt8(Arc<Vec<u8>>),
921    Complex64(Arc<Vec<Complex64>>),
922    Complex128(Arc<Vec<Complex128>>),
923}
924
925impl TensorStorage {
926    #[must_use]
927    pub fn len(&self) -> usize {
928        match self {
929            Self::F32(v) => v.len(),
930            Self::F64(v) => v.len(),
931            Self::F64Inline4(v) => v.len(),
932            Self::F16(v) => v.len(),
933            Self::BF16(v) => v.len(),
934            Self::QInt8(v) => v.len(),
935            Self::QUInt8(v) => v.len(),
936            Self::Complex64(v) => v.len(),
937            Self::Complex128(v) => v.len(),
938        }
939    }
940
941    #[must_use]
942    pub fn is_empty(&self) -> bool {
943        self.len() == 0
944    }
945
946    #[must_use]
947    pub fn dtype(&self) -> DType {
948        match self {
949            Self::F32(_) => DType::F32,
950            Self::F64(_) | Self::F64Inline4(_) => DType::F64,
951            Self::F16(_) => DType::F16,
952            Self::BF16(_) => DType::BF16,
953            Self::QInt8(_) => DType::QInt8,
954            Self::QUInt8(_) => DType::QUInt8,
955            Self::Complex64(_) => DType::Complex64,
956            Self::Complex128(_) => DType::Complex128,
957        }
958    }
959
960    #[must_use]
961    pub fn as_f64(&self) -> Option<&[f64]> {
962        match self {
963            Self::F64(v) => Some(v.as_slice()),
964            Self::F64Inline4(v) => Some(v.as_slice()),
965            _ => None,
966        }
967    }
968
969    #[must_use]
970    pub fn as_f32(&self) -> Option<&[f32]> {
971        match self {
972            Self::F32(v) => Some(v.as_slice()),
973            _ => None,
974        }
975    }
976
977    #[must_use]
978    pub fn as_f16(&self) -> Option<&[Float16]> {
979        match self {
980            Self::F16(v) => Some(v.as_slice()),
981            _ => None,
982        }
983    }
984
985    #[must_use]
986    pub fn as_bf16(&self) -> Option<&[BFloat16]> {
987        match self {
988            Self::BF16(v) => Some(v.as_slice()),
989            _ => None,
990        }
991    }
992
993    #[must_use]
994    pub fn as_qint8(&self) -> Option<&[i8]> {
995        match self {
996            Self::QInt8(v) => Some(v.as_slice()),
997            _ => None,
998        }
999    }
1000
1001    #[must_use]
1002    pub fn as_quint8(&self) -> Option<&[u8]> {
1003        match self {
1004            Self::QUInt8(v) => Some(v.as_slice()),
1005            _ => None,
1006        }
1007    }
1008
1009    #[must_use]
1010    pub fn as_complex64(&self) -> Option<&[Complex64]> {
1011        match self {
1012            Self::Complex64(v) => Some(v.as_slice()),
1013            _ => None,
1014        }
1015    }
1016
1017    #[must_use]
1018    pub fn as_complex128(&self) -> Option<&[Complex128]> {
1019        match self {
1020            Self::Complex128(v) => Some(v.as_slice()),
1021            _ => None,
1022        }
1023    }
1024
1025    /// Convert storage to f64 values, promoting from any float type.
1026    /// Complex types extract the real part.
1027    #[must_use]
1028    pub fn to_f64_vec(&self) -> Vec<f64> {
1029        match self {
1030            Self::F64(v) => v.as_ref().clone(),
1031            Self::F64Inline4(v) => v.to_vec(),
1032            Self::F32(v) => v.iter().map(|&x| f64::from(x)).collect(),
1033            Self::F16(v) => v.iter().map(|&x| f64::from(x.to_f32())).collect(),
1034            Self::BF16(v) => v.iter().map(|&x| f64::from(x.to_f32())).collect(),
1035            Self::QInt8(v) => v.iter().map(|&x| f64::from(x)).collect(),
1036            Self::QUInt8(v) => v.iter().map(|&x| f64::from(x)).collect(),
1037            Self::Complex64(v) => v.iter().map(|z| f64::from(z.re)).collect(),
1038            Self::Complex128(v) => v.iter().map(|z| z.re).collect(),
1039        }
1040    }
1041
1042    /// Convert storage to f32 values, promoting from half or demoting from f64.
1043    /// Complex types extract the real part.
1044    #[must_use]
1045    pub fn to_f32_vec(&self) -> Vec<f32> {
1046        match self {
1047            Self::F32(v) => v.as_ref().clone(),
1048            Self::F64(v) => v.iter().map(|&x| x as f32).collect(),
1049            Self::F64Inline4(v) => v.iter().map(|&x| x as f32).collect(),
1050            Self::F16(v) => v.iter().map(|&x| x.to_f32()).collect(),
1051            Self::BF16(v) => v.iter().map(|&x| x.to_f32()).collect(),
1052            Self::QInt8(v) => v.iter().map(|&x| f32::from(x)).collect(),
1053            Self::QUInt8(v) => v.iter().map(|&x| f32::from(x)).collect(),
1054            Self::Complex64(v) => v.iter().map(|z| z.re).collect(),
1055            Self::Complex128(v) => v.iter().map(|z| z.re as f32).collect(),
1056        }
1057    }
1058}
1059
1060#[derive(Debug, Clone, PartialEq)]
1061pub struct DenseTensor {
1062    id: u64,
1063    storage_id: u64,
1064    meta: TensorMeta,
1065    storage: TensorStorage,
1066    version: u64,
1067}
1068
1069#[derive(Debug, Clone, PartialEq, Eq)]
1070pub enum DenseTensorError {
1071    Meta(TensorMetaError),
1072    UnsupportedDType(DType),
1073    UnsupportedLayout,
1074    UnsupportedStorageAccess { dtype: DType },
1075    StorageSpanOverflow { storage_offset: usize, numel: usize },
1076    InsufficientStorage { needed: usize, actual: usize },
1077    ShapeOverflow { shape: Vec<usize> },
1078}
1079
1080impl fmt::Display for DenseTensorError {
1081    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1082        match self {
1083            Self::Meta(error) => write!(f, "invalid tensor metadata: {error}"),
1084            Self::UnsupportedDType(dtype) => write!(
1085                f,
1086                "unsupported tensor dtype for this storage type: {dtype:?}"
1087            ),
1088            Self::UnsupportedLayout => {
1089                write!(f, "dense tensor requires contiguous layout")
1090            }
1091            Self::UnsupportedStorageAccess { dtype } => write!(
1092                f,
1093                "raw f64 storage access is unsupported for tensor dtype {dtype:?}; use typed_storage() or contiguous_values_as_f64()"
1094            ),
1095            Self::StorageSpanOverflow {
1096                storage_offset,
1097                numel,
1098            } => write!(
1099                f,
1100                "dense tensor storage span overflow for storage_offset={storage_offset}, numel={numel}"
1101            ),
1102            Self::InsufficientStorage { needed, actual } => write!(
1103                f,
1104                "dense tensor storage length is insufficient: needed={needed}, actual={actual}"
1105            ),
1106            Self::ShapeOverflow { shape } => {
1107                write!(f, "dense tensor shape volume overflow for shape={shape:?}")
1108            }
1109        }
1110    }
1111}
1112
1113impl std::error::Error for DenseTensorError {}
1114
1115impl From<TensorMetaError> for DenseTensorError {
1116    fn from(value: TensorMetaError) -> Self {
1117        Self::Meta(value)
1118    }
1119}
1120
1121fn contiguous_required_len(meta: &TensorMeta) -> Result<usize, DenseTensorError> {
1122    meta.storage_offset()
1123        .checked_add(meta.numel())
1124        .ok_or(DenseTensorError::StorageSpanOverflow {
1125            storage_offset: meta.storage_offset(),
1126            numel: meta.numel(),
1127        })
1128}
1129
1130fn checked_shape_numel(shape: &[usize]) -> Result<usize, DenseTensorError> {
1131    if shape.is_empty() {
1132        return Ok(1);
1133    }
1134    let mut product = 1usize;
1135    for &dim in shape {
1136        if dim == 0 {
1137            return Ok(0);
1138        }
1139        product = product
1140            .checked_mul(dim)
1141            .ok_or_else(|| DenseTensorError::ShapeOverflow {
1142                shape: shape.to_vec(),
1143            })?;
1144    }
1145    Ok(product)
1146}
1147
1148impl DenseTensor {
1149    pub fn from_typed_storage(
1150        meta: TensorMeta,
1151        storage: TensorStorage,
1152    ) -> Result<Self, DenseTensorError> {
1153        meta.validate()?;
1154        if meta.dtype() != storage.dtype() {
1155            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1156        }
1157        if !meta.dtype().is_floating_point()
1158            && !meta.dtype().is_complex()
1159            && !meta.dtype().is_quantized()
1160        {
1161            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1162        }
1163
1164        let needed = Self::storage_span_required_len(&meta)?;
1165        if storage.len() < needed {
1166            return Err(DenseTensorError::InsufficientStorage {
1167                needed,
1168                actual: storage.len(),
1169            });
1170        }
1171
1172        Ok(Self {
1173            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
1174            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
1175            meta,
1176            storage,
1177            version: 0,
1178        })
1179    }
1180
1181    pub fn from_storage(meta: TensorMeta, storage: Vec<f64>) -> Result<Self, DenseTensorError> {
1182        if meta.dtype() != DType::F64 {
1183            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1184        }
1185        Self::from_typed_storage(meta, TensorStorage::F64(Arc::new(storage)))
1186    }
1187
1188    pub fn from_storage_f64_inline4(
1189        meta: TensorMeta,
1190        storage: [f64; 4],
1191    ) -> Result<Self, DenseTensorError> {
1192        if meta.dtype() != DType::F64 {
1193            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1194        }
1195        Self::from_typed_storage(meta, TensorStorage::F64Inline4(storage))
1196    }
1197
1198    pub fn from_storage_f32(meta: TensorMeta, storage: Vec<f32>) -> Result<Self, DenseTensorError> {
1199        if meta.dtype() != DType::F32 {
1200            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1201        }
1202        Self::from_typed_storage(meta, TensorStorage::F32(Arc::new(storage)))
1203    }
1204
1205    pub fn from_storage_f16(
1206        meta: TensorMeta,
1207        storage: Vec<Float16>,
1208    ) -> Result<Self, DenseTensorError> {
1209        if meta.dtype() != DType::F16 {
1210            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1211        }
1212        Self::from_typed_storage(meta, TensorStorage::F16(Arc::new(storage)))
1213    }
1214
1215    pub fn from_storage_bf16(
1216        meta: TensorMeta,
1217        storage: Vec<BFloat16>,
1218    ) -> Result<Self, DenseTensorError> {
1219        if meta.dtype() != DType::BF16 {
1220            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1221        }
1222        Self::from_typed_storage(meta, TensorStorage::BF16(Arc::new(storage)))
1223    }
1224
1225    pub fn from_storage_qint8(
1226        meta: TensorMeta,
1227        storage: Vec<i8>,
1228    ) -> Result<Self, DenseTensorError> {
1229        if meta.dtype() != DType::QInt8 {
1230            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1231        }
1232        Self::from_typed_storage(meta, TensorStorage::QInt8(Arc::new(storage)))
1233    }
1234
1235    pub fn from_storage_quint8(
1236        meta: TensorMeta,
1237        storage: Vec<u8>,
1238    ) -> Result<Self, DenseTensorError> {
1239        if meta.dtype() != DType::QUInt8 {
1240            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1241        }
1242        Self::from_typed_storage(meta, TensorStorage::QUInt8(Arc::new(storage)))
1243    }
1244
1245    pub fn from_contiguous_values(
1246        values: Vec<f64>,
1247        shape: Vec<usize>,
1248        device: Device,
1249    ) -> Result<Self, DenseTensorError> {
1250        let meta = TensorMeta::from_shape(shape, DType::F64, device);
1251        Self::from_storage(meta, values)
1252    }
1253
1254    pub fn from_contiguous_values_f32(
1255        values: Vec<f32>,
1256        shape: Vec<usize>,
1257        device: Device,
1258    ) -> Result<Self, DenseTensorError> {
1259        let meta = TensorMeta::from_shape(shape, DType::F32, device);
1260        Self::from_storage_f32(meta, values)
1261    }
1262
1263    pub fn from_contiguous_values_f16(
1264        values: Vec<Float16>,
1265        shape: Vec<usize>,
1266        device: Device,
1267    ) -> Result<Self, DenseTensorError> {
1268        let meta = TensorMeta::from_shape(shape, DType::F16, device);
1269        Self::from_storage_f16(meta, values)
1270    }
1271
1272    pub fn from_contiguous_values_bf16(
1273        values: Vec<BFloat16>,
1274        shape: Vec<usize>,
1275        device: Device,
1276    ) -> Result<Self, DenseTensorError> {
1277        let meta = TensorMeta::from_shape(shape, DType::BF16, device);
1278        Self::from_storage_bf16(meta, values)
1279    }
1280
1281    pub fn from_contiguous_values_qint8(
1282        values: Vec<i8>,
1283        shape: Vec<usize>,
1284        device: Device,
1285        scale: f64,
1286        zero_point: i64,
1287    ) -> Result<Self, DenseTensorError> {
1288        let meta =
1289            TensorMeta::quantized_from_shape(shape, DType::QInt8, device, scale, zero_point)?;
1290        Self::from_storage_qint8(meta, values)
1291    }
1292
1293    pub fn from_contiguous_values_qint8_per_channel(
1294        values: Vec<i8>,
1295        shape: Vec<usize>,
1296        device: Device,
1297        scales: Vec<f64>,
1298        zero_points: Vec<i64>,
1299        axis: usize,
1300    ) -> Result<Self, DenseTensorError> {
1301        let meta = TensorMeta::quantized_per_channel_from_shape(
1302            shape,
1303            DType::QInt8,
1304            device,
1305            scales,
1306            zero_points,
1307            axis,
1308        )?;
1309        Self::from_storage_qint8(meta, values)
1310    }
1311
1312    pub fn from_contiguous_values_quint8(
1313        values: Vec<u8>,
1314        shape: Vec<usize>,
1315        device: Device,
1316        scale: f64,
1317        zero_point: i64,
1318    ) -> Result<Self, DenseTensorError> {
1319        let meta =
1320            TensorMeta::quantized_from_shape(shape, DType::QUInt8, device, scale, zero_point)?;
1321        Self::from_storage_quint8(meta, values)
1322    }
1323
1324    pub fn from_contiguous_values_quint8_per_channel(
1325        values: Vec<u8>,
1326        shape: Vec<usize>,
1327        device: Device,
1328        scales: Vec<f64>,
1329        zero_points: Vec<i64>,
1330        axis: usize,
1331    ) -> Result<Self, DenseTensorError> {
1332        let meta = TensorMeta::quantized_per_channel_from_shape(
1333            shape,
1334            DType::QUInt8,
1335            device,
1336            scales,
1337            zero_points,
1338            axis,
1339        )?;
1340        Self::from_storage_quint8(meta, values)
1341    }
1342
1343    fn contiguous_required_len(meta: &TensorMeta) -> Result<usize, DenseTensorError> {
1344        contiguous_required_len(meta)
1345    }
1346
1347    fn storage_span_required_len(meta: &TensorMeta) -> Result<usize, DenseTensorError> {
1348        let mut max_linear_offset = 0usize;
1349        for (size, stride) in meta
1350            .shape()
1351            .iter()
1352            .copied()
1353            .zip(meta.strides().iter().copied())
1354        {
1355            if size == 0 {
1356                continue;
1357            }
1358            let span = stride.checked_mul(size.saturating_sub(1)).ok_or(
1359                DenseTensorError::StorageSpanOverflow {
1360                    storage_offset: meta.storage_offset(),
1361                    numel: meta.numel(),
1362                },
1363            )?;
1364            max_linear_offset = max_linear_offset.checked_add(span).ok_or(
1365                DenseTensorError::StorageSpanOverflow {
1366                    storage_offset: meta.storage_offset(),
1367                    numel: meta.numel(),
1368                },
1369            )?;
1370        }
1371
1372        if meta.numel() == 0 {
1373            return Ok(meta.storage_offset());
1374        }
1375
1376        let max_index = meta.storage_offset().checked_add(max_linear_offset).ok_or(
1377            DenseTensorError::StorageSpanOverflow {
1378                storage_offset: meta.storage_offset(),
1379                numel: meta.numel(),
1380            },
1381        )?;
1382        max_index
1383            .checked_add(1)
1384            .ok_or(DenseTensorError::StorageSpanOverflow {
1385                storage_offset: meta.storage_offset(),
1386                numel: meta.numel(),
1387            })
1388    }
1389
1390    pub fn dispatch_values(&self) -> Result<&[f64], DenseTensorError> {
1391        let start = self.meta.storage_offset();
1392        let end = Self::storage_span_required_len(&self.meta)?;
1393        match &self.storage {
1394            TensorStorage::F64(v) => Ok(&v[start..end]),
1395            TensorStorage::F64Inline4(v) => Ok(&v[start..end]),
1396            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1397        }
1398    }
1399
1400    pub fn contiguous_values(&self) -> Result<&[f64], DenseTensorError> {
1401        if !self.meta.is_contiguous() {
1402            return Err(DenseTensorError::UnsupportedLayout);
1403        }
1404        self.dispatch_values()
1405    }
1406
1407    pub fn contiguous_values_f32(&self) -> Result<&[f32], DenseTensorError> {
1408        if !self.meta.is_contiguous() {
1409            return Err(DenseTensorError::UnsupportedLayout);
1410        }
1411        let start = self.meta.storage_offset();
1412        let end = Self::storage_span_required_len(&self.meta)?;
1413        match &self.storage {
1414            TensorStorage::F32(v) => Ok(&v[start..end]),
1415            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1416        }
1417    }
1418
1419    /// Mutable view of a contiguous f32 tensor's values — the in-place counterpart
1420    /// of [`Self::contiguous_values_f32`]. Lets an inference-only consumer rewrite a
1421    /// dead intermediate's storage (e.g. a fused softmax / GELU) without
1422    /// materializing a fresh tensor. The caller is responsible for the value being
1423    /// used as a single-owner intermediate (no aliasing tape node depends on the
1424    /// pre-mutation contents).
1425    pub fn contiguous_values_f32_mut(&mut self) -> Result<&mut [f32], DenseTensorError> {
1426        if !self.meta.is_contiguous() {
1427            return Err(DenseTensorError::UnsupportedLayout);
1428        }
1429        let start = self.meta.storage_offset();
1430        let end = Self::storage_span_required_len(&self.meta)?;
1431        let dtype = self.meta.dtype();
1432        match &mut self.storage {
1433            // `make_mut` mutates in place when this is the sole owner (the common
1434            // case for a fresh single-use intermediate) and performs a safe
1435            // copy-on-write clone if the storage is shared, so no other tensor's
1436            // values are ever clobbered.
1437            TensorStorage::F32(v) => Ok(&mut Arc::make_mut(v)[start..end]),
1438            _ => Err(DenseTensorError::UnsupportedDType(dtype)),
1439        }
1440    }
1441
1442    pub fn contiguous_values_qint8(&self) -> Result<&[i8], DenseTensorError> {
1443        if !self.meta.is_contiguous() {
1444            return Err(DenseTensorError::UnsupportedLayout);
1445        }
1446        let start = self.meta.storage_offset();
1447        let end = Self::storage_span_required_len(&self.meta)?;
1448        match &self.storage {
1449            TensorStorage::QInt8(v) => Ok(&v[start..end]),
1450            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1451        }
1452    }
1453
1454    pub fn contiguous_values_quint8(&self) -> Result<&[u8], DenseTensorError> {
1455        if !self.meta.is_contiguous() {
1456            return Err(DenseTensorError::UnsupportedLayout);
1457        }
1458        let start = self.meta.storage_offset();
1459        let end = Self::storage_span_required_len(&self.meta)?;
1460        match &self.storage {
1461            TensorStorage::QUInt8(v) => Ok(&v[start..end]),
1462            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1463        }
1464    }
1465
1466    /// Returns contiguous values as f64, converting from any float type.
1467    /// Used by backward pass to keep gradient computation in f64.
1468    pub fn contiguous_values_as_f64(&self) -> Result<Vec<f64>, DenseTensorError> {
1469        if !self.meta.is_contiguous() {
1470            return Err(DenseTensorError::UnsupportedLayout);
1471        }
1472        let start = self.meta.storage_offset();
1473        let end = Self::storage_span_required_len(&self.meta)?;
1474        match &self.storage {
1475            TensorStorage::F64(v) => Ok(v[start..end].to_vec()),
1476            TensorStorage::F64Inline4(v) => Ok(v[start..end].to_vec()),
1477            TensorStorage::F32(v) => Ok(v[start..end].iter().map(|&x| f64::from(x)).collect()),
1478            TensorStorage::F16(v) => Ok(v[start..end]
1479                .iter()
1480                .map(|&x| f64::from(x.to_f32()))
1481                .collect()),
1482            TensorStorage::BF16(v) => Ok(v[start..end]
1483                .iter()
1484                .map(|&x| f64::from(x.to_f32()))
1485                .collect()),
1486            TensorStorage::Complex64(v) => {
1487                Ok(v[start..end].iter().map(|z| f64::from(z.re)).collect())
1488            }
1489            TensorStorage::Complex128(v) => Ok(v[start..end].iter().map(|z| z.re).collect()),
1490            TensorStorage::QInt8(v) => Ok(v[start..end].iter().map(|&x| f64::from(x)).collect()),
1491            TensorStorage::QUInt8(v) => Ok(v[start..end].iter().map(|&x| f64::from(x)).collect()),
1492        }
1493    }
1494
1495    fn contiguous_complex_values_as_complex128(&self) -> Result<Vec<Complex128>, DenseTensorError> {
1496        if !self.meta.is_contiguous() {
1497            return Err(DenseTensorError::UnsupportedLayout);
1498        }
1499        let start = self.meta.storage_offset();
1500        let end = Self::storage_span_required_len(&self.meta)?;
1501        match &self.storage {
1502            TensorStorage::Complex64(v) => Ok(v[start..end]
1503                .iter()
1504                .map(|z| Complex128::new(f64::from(z.re), f64::from(z.im)))
1505                .collect()),
1506            TensorStorage::Complex128(v) => Ok(v[start..end].to_vec()),
1507            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1508        }
1509    }
1510
1511    pub fn dequantized_values_as_f64(&self) -> Result<Vec<f64>, DenseTensorError> {
1512        if !self.meta.is_contiguous() {
1513            return Err(DenseTensorError::UnsupportedLayout);
1514        }
1515        let Some(qparams) = self.meta.quantization() else {
1516            return Err(DenseTensorError::UnsupportedDType(self.meta.dtype()));
1517        };
1518        let start = self.meta.storage_offset();
1519        let end = Self::storage_span_required_len(&self.meta)?;
1520        // Per-channel layout (`channel_count` = shape[axis], `inner` =
1521        // shape[axis+1..].product()) is shape/axis-derived, NOT element-dependent.
1522        // Hoist it out of the per-element closure so the inner-dims product loop
1523        // runs ONCE instead of for every element (was O(numel * rank); now O(numel)
1524        // + O(rank) once). Bit-for-bit identical to the prior `channel_index_for_flat`
1525        // path: the same Option checks and the same per-element index formula
1526        // `(flat_idx / inner) % channel_count` are preserved, only the product
1527        // recomputation moves out.
1528        let rank = self.meta.shape().len();
1529        let channel_layout = qparams.axis().map(|axis| {
1530            let shape = self.meta.shape();
1531            let channel_count = shape.get(axis).copied();
1532            let inner = shape
1533                .get(axis.saturating_add(1)..)
1534                .map(|rest| rest.iter().copied().product::<usize>());
1535            (axis, channel_count, inner)
1536        });
1537        let dequantize = |flat_idx: usize, qvalue: f64| -> Result<f64, DenseTensorError> {
1538            let channel = match channel_layout {
1539                Some((axis, channel_count, inner)) => {
1540                    let invalid_axis = || {
1541                        DenseTensorError::Meta(TensorMetaError::InvalidQuantizationAxis {
1542                            axis,
1543                            rank,
1544                        })
1545                    };
1546                    let channel_count = channel_count.ok_or_else(invalid_axis)?;
1547                    let inner = inner.ok_or_else(invalid_axis)?;
1548                    if channel_count == 0 || inner == 0 {
1549                        return Err(invalid_axis());
1550                    }
1551                    (flat_idx / inner) % channel_count
1552                }
1553                None => 0,
1554            };
1555            let scale = qparams.scale_at(channel).ok_or(DenseTensorError::Meta(
1556                TensorMetaError::QuantizationChannelCountMismatch {
1557                    axis: qparams.axis().unwrap_or(0),
1558                    expected: qparams.len(),
1559                    actual: channel.saturating_add(1),
1560                },
1561            ))?;
1562            let zero_point = qparams
1563                .zero_point_at(channel)
1564                .ok_or(DenseTensorError::Meta(
1565                    TensorMetaError::QuantizationVectorLengthMismatch {
1566                        scales: qparams.scale_bits.len(),
1567                        zero_points: qparams.zero_points.len(),
1568                    },
1569                ))?;
1570            // frankentorch-9rvxq: compute in f32 and widen ONCE, matching
1571            // `torch.Tensor.dequantize()`, which produces float32 — callers wanting f64 widen
1572            // afterwards, so torch's f32 rounding is baked into the value. Doing this in f64 is
1573            // strictly MORE accurate and therefore diverges: a quantized linear consuming
1574            // f64-dequantized weights landed 1.16e-6 off torch on an f64 dot product, because
1575            // every weight carried ~16 significant digits here against torch's ~7. Verified
1576            // bit-exact (full-tensor equality, not a tolerance) against
1577            // `torch.quantize_per_tensor(...).dequantize()`. Being more precise than the
1578            // reference is still a parity bug when the reference's rounding is observable.
1579            #[allow(clippy::cast_possible_truncation)]
1580            let scale = scale as f32;
1581            #[allow(clippy::cast_precision_loss)]
1582            let zero_point = zero_point as f32;
1583            #[allow(clippy::cast_possible_truncation)]
1584            let qvalue = qvalue as f32;
1585            Ok(f64::from((qvalue - zero_point) * scale))
1586        };
1587        match &self.storage {
1588            TensorStorage::QInt8(v) => Ok(v[start..end]
1589                .iter()
1590                .enumerate()
1591                .map(|(flat_idx, &q)| dequantize(flat_idx, f64::from(q)))
1592                .collect::<Result<Vec<_>, _>>()?),
1593            TensorStorage::QUInt8(v) => Ok(v[start..end]
1594                .iter()
1595                .enumerate()
1596                .map(|(flat_idx, &q)| dequantize(flat_idx, f64::from(q)))
1597                .collect::<Result<Vec<_>, _>>()?),
1598            _ => Err(DenseTensorError::UnsupportedDType(self.meta.dtype())),
1599        }
1600    }
1601
1602    #[must_use]
1603    pub fn typed_storage(&self) -> &TensorStorage {
1604        &self.storage
1605    }
1606
1607    /// Returns the raw f64 storage slice.
1608    ///
1609    /// This is only available for `F64` tensors. Callers that need a dtype-agnostic
1610    /// path should use `typed_storage()` or `contiguous_values_as_f64()`.
1611    pub fn storage(&self) -> Result<&[f64], DenseTensorError> {
1612        match &self.storage {
1613            TensorStorage::F64(v) => Ok(v.as_slice()),
1614            TensorStorage::F64Inline4(v) => Ok(v.as_slice()),
1615            other => Err(DenseTensorError::UnsupportedStorageAccess {
1616                dtype: other.dtype(),
1617            }),
1618        }
1619    }
1620
1621    #[must_use]
1622    pub fn meta(&self) -> &TensorMeta {
1623        &self.meta
1624    }
1625
1626    #[must_use]
1627    pub fn id(&self) -> u64 {
1628        self.id
1629    }
1630
1631    #[must_use]
1632    pub fn storage_id(&self) -> u64 {
1633        self.storage_id
1634    }
1635
1636    #[must_use]
1637    pub fn version(&self) -> u64 {
1638        self.version
1639    }
1640
1641    /// Cast this tensor to a different floating-point or complex dtype.
1642    pub fn to_dtype(&self, dtype: DType) -> Result<Self, DenseTensorError> {
1643        if !dtype.is_floating_point() && !dtype.is_complex() {
1644            return Err(DenseTensorError::UnsupportedDType(dtype));
1645        }
1646        if self.meta.dtype() == dtype {
1647            return Ok(self.clone());
1648        }
1649        let new_meta =
1650            TensorMeta::from_shape(self.meta.shape().to_vec(), dtype, self.meta.device());
1651        let logical_f64 = if self.meta.dtype().is_quantized() {
1652            self.dequantized_values_as_f64()?
1653        } else {
1654            self.contiguous_values_as_f64()?
1655        };
1656        let as_f64 = || -> Vec<f64> { logical_f64.clone() };
1657        let as_f32 = || -> Vec<f32> {
1658            #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
1659            {
1660                as_f64().into_iter().map(|value| value as f32).collect()
1661            }
1662        };
1663        let as_complex128 = || -> Result<Vec<Complex128>, DenseTensorError> {
1664            if self.meta.dtype().is_complex() {
1665                self.contiguous_complex_values_as_complex128()
1666            } else {
1667                Ok(as_f64()
1668                    .into_iter()
1669                    .map(|r| Complex128::new(r, 0.0))
1670                    .collect())
1671            }
1672        };
1673        let new_storage = match dtype {
1674            DType::F64 => TensorStorage::F64(Arc::new(as_f64())),
1675            DType::F32 => TensorStorage::F32(Arc::new(as_f32())),
1676            DType::F16 => {
1677                let vals: Vec<Float16> = as_f32().into_iter().map(Float16::from_f32).collect();
1678                TensorStorage::F16(Arc::new(vals))
1679            }
1680            DType::BF16 => {
1681                let vals: Vec<BFloat16> = as_f32().into_iter().map(BFloat16::from_f32).collect();
1682                TensorStorage::BF16(Arc::new(vals))
1683            }
1684            DType::Complex64 => {
1685                let vals: Vec<Complex64> = if self.meta.dtype().is_complex() {
1686                    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
1687                    {
1688                        as_complex128()?
1689                            .into_iter()
1690                            .map(|z| Complex64::new(z.re as f32, z.im as f32))
1691                            .collect()
1692                    }
1693                } else {
1694                    as_f32()
1695                        .into_iter()
1696                        .map(|r| Complex64::new(r, 0.0))
1697                        .collect()
1698                };
1699                TensorStorage::Complex64(Arc::new(vals))
1700            }
1701            DType::Complex128 => {
1702                let vals: Vec<Complex128> = as_complex128()?;
1703                TensorStorage::Complex128(Arc::new(vals))
1704            }
1705            _ => return Err(DenseTensorError::UnsupportedDType(dtype)),
1706        };
1707        Self::from_typed_storage(new_meta, new_storage)
1708    }
1709
1710    /// Update the contiguous values in-place and bump the version counter.
1711    ///
1712    /// The new values must exactly match the length of the contiguous slice.
1713    pub fn update_contiguous_values(&mut self, new_values: &[f64]) -> Result<(), DenseTensorError> {
1714        if !self.meta.is_contiguous() {
1715            return Err(DenseTensorError::UnsupportedLayout);
1716        }
1717        let start = self.meta.storage_offset();
1718        let end = Self::contiguous_required_len(&self.meta)?;
1719        match &mut self.storage {
1720            TensorStorage::F64(v) => {
1721                let buf = Arc::make_mut(v);
1722                let slice = &mut buf[start..end];
1723                if new_values.len() != slice.len() {
1724                    return Err(DenseTensorError::InsufficientStorage {
1725                        needed: slice.len(),
1726                        actual: new_values.len(),
1727                    });
1728                }
1729                slice.copy_from_slice(new_values);
1730            }
1731            TensorStorage::F64Inline4(v) => {
1732                let slice = &mut v[start..end];
1733                if new_values.len() != slice.len() {
1734                    return Err(DenseTensorError::InsufficientStorage {
1735                        needed: slice.len(),
1736                        actual: new_values.len(),
1737                    });
1738                }
1739                slice.copy_from_slice(new_values);
1740            }
1741            _ => {
1742                return Err(DenseTensorError::UnsupportedDType(self.meta.dtype()));
1743            }
1744        }
1745        self.version += 1;
1746        Ok(())
1747    }
1748
1749    /// Mutate the contiguous f64 values in-place and bump the version counter.
1750    pub fn update_contiguous_values_with<F>(&mut self, update: F) -> Result<(), DenseTensorError>
1751    where
1752        F: FnOnce(&mut [f64]),
1753    {
1754        if !self.meta.is_contiguous() {
1755            return Err(DenseTensorError::UnsupportedLayout);
1756        }
1757        let start = self.meta.storage_offset();
1758        let end = Self::contiguous_required_len(&self.meta)?;
1759        match &mut self.storage {
1760            TensorStorage::F64(v) => {
1761                let buf = Arc::make_mut(v);
1762                update(&mut buf[start..end]);
1763            }
1764            TensorStorage::F64Inline4(v) => {
1765                update(&mut v[start..end]);
1766            }
1767            _ => {
1768                return Err(DenseTensorError::UnsupportedDType(self.meta.dtype()));
1769            }
1770        }
1771        self.version += 1;
1772        Ok(())
1773    }
1774
1775    /// Update the contiguous f32 values in-place and bump the version counter.
1776    pub fn update_contiguous_values_f32(
1777        &mut self,
1778        new_values: &[f32],
1779    ) -> Result<(), DenseTensorError> {
1780        if !self.meta.is_contiguous() {
1781            return Err(DenseTensorError::UnsupportedLayout);
1782        }
1783        let start = self.meta.storage_offset();
1784        let end = Self::contiguous_required_len(&self.meta)?;
1785        match &mut self.storage {
1786            TensorStorage::F32(v) => {
1787                let buf = Arc::make_mut(v);
1788                let slice = &mut buf[start..end];
1789                if new_values.len() != slice.len() {
1790                    return Err(DenseTensorError::InsufficientStorage {
1791                        needed: slice.len(),
1792                        actual: new_values.len(),
1793                    });
1794                }
1795                slice.copy_from_slice(new_values);
1796            }
1797            _ => {
1798                return Err(DenseTensorError::UnsupportedDType(self.meta.dtype()));
1799            }
1800        }
1801        self.version += 1;
1802        Ok(())
1803    }
1804
1805    /// Create a view of this tensor with a new shape.
1806    /// The view shares the same underlying storage (zero-copy).
1807    /// Only works for contiguous tensors where the new shape has the same numel.
1808    pub fn view(&self, new_shape: Vec<usize>) -> Result<Self, DenseTensorError> {
1809        if !self.meta.is_contiguous() {
1810            return Err(DenseTensorError::UnsupportedLayout);
1811        }
1812        let new_numel = checked_shape_numel(&new_shape)?;
1813        if new_numel != self.meta.numel() {
1814            return Err(DenseTensorError::InsufficientStorage {
1815                needed: new_numel,
1816                actual: self.meta.numel(),
1817            });
1818        }
1819        let new_meta = TensorMeta {
1820            strides: contiguous_strides(&new_shape),
1821            shape: new_shape,
1822            numel: new_numel,
1823            storage_offset: self.meta.storage_offset(),
1824            dtype: self.meta.dtype(),
1825            device: self.meta.device(),
1826            quantization: self.meta.quantization.clone(),
1827        };
1828        new_meta.validate()?;
1829        Ok(Self {
1830            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
1831            storage_id: self.storage_id, // same storage
1832            meta: new_meta,
1833            storage: self.storage.clone(), // Arc clone = cheap refcount bump
1834            version: self.version,
1835        })
1836    }
1837
1838    /// Returns true if this tensor shares storage with another.
1839    #[must_use]
1840    pub fn shares_storage_with(&self, other: &Self) -> bool {
1841        self.storage_id == other.storage_id
1842    }
1843}
1844
1845// ── Integer Tensor Types ───────────────────────────────────────────────
1846
1847/// Dense tensor backed by `Vec<i64>` storage.
1848///
1849/// Integer tensors do NOT participate in autograd (`requires_grad` is always false).
1850/// They are used for indexing, class labels, and shape operations.
1851#[derive(Debug, Clone, PartialEq, Eq)]
1852pub struct DenseI64Tensor {
1853    id: u64,
1854    storage_id: u64,
1855    meta: TensorMeta,
1856    storage: Vec<i64>,
1857    version: u64,
1858}
1859
1860impl DenseI64Tensor {
1861    pub fn from_storage(meta: TensorMeta, storage: Vec<i64>) -> Result<Self, DenseTensorError> {
1862        meta.validate()?;
1863        if meta.dtype() != DType::I64 {
1864            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1865        }
1866        if !meta.is_contiguous() {
1867            return Err(DenseTensorError::UnsupportedLayout);
1868        }
1869        let needed = contiguous_required_len(&meta)?;
1870        if storage.len() < needed {
1871            return Err(DenseTensorError::InsufficientStorage {
1872                needed,
1873                actual: storage.len(),
1874            });
1875        }
1876        Ok(Self {
1877            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
1878            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
1879            meta,
1880            storage,
1881            version: 0,
1882        })
1883    }
1884
1885    pub fn from_contiguous_values(
1886        values: Vec<i64>,
1887        shape: Vec<usize>,
1888        device: Device,
1889    ) -> Result<Self, DenseTensorError> {
1890        let meta = TensorMeta::from_shape(shape, DType::I64, device);
1891        Self::from_storage(meta, values)
1892    }
1893
1894    #[must_use]
1895    pub fn id(&self) -> u64 {
1896        self.id
1897    }
1898
1899    #[must_use]
1900    pub fn storage_id(&self) -> u64 {
1901        self.storage_id
1902    }
1903
1904    #[must_use]
1905    pub fn meta(&self) -> &TensorMeta {
1906        &self.meta
1907    }
1908
1909    #[must_use]
1910    pub fn storage(&self) -> &[i64] {
1911        &self.storage
1912    }
1913
1914    /// Return the contiguous values slice.
1915    pub fn contiguous_values(&self) -> Result<&[i64], DenseTensorError> {
1916        if !self.meta.is_contiguous() {
1917            return Err(DenseTensorError::UnsupportedLayout);
1918        }
1919        let start = self.meta.storage_offset();
1920        let end = contiguous_required_len(&self.meta)?;
1921        Ok(&self.storage[start..end])
1922    }
1923
1924    #[must_use]
1925    pub fn version(&self) -> u64 {
1926        self.version
1927    }
1928}
1929
1930/// Dense tensor backed by `Vec<i32>` storage.
1931///
1932/// Integer tensors do NOT participate in autograd (`requires_grad` is always false).
1933#[derive(Debug, Clone, PartialEq, Eq)]
1934pub struct DenseI32Tensor {
1935    id: u64,
1936    storage_id: u64,
1937    meta: TensorMeta,
1938    storage: Vec<i32>,
1939    version: u64,
1940}
1941
1942impl DenseI32Tensor {
1943    pub fn from_storage(meta: TensorMeta, storage: Vec<i32>) -> Result<Self, DenseTensorError> {
1944        meta.validate()?;
1945        if meta.dtype() != DType::I32 {
1946            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
1947        }
1948        if !meta.is_contiguous() {
1949            return Err(DenseTensorError::UnsupportedLayout);
1950        }
1951        let needed = contiguous_required_len(&meta)?;
1952        if storage.len() < needed {
1953            return Err(DenseTensorError::InsufficientStorage {
1954                needed,
1955                actual: storage.len(),
1956            });
1957        }
1958        Ok(Self {
1959            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
1960            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
1961            meta,
1962            storage,
1963            version: 0,
1964        })
1965    }
1966
1967    pub fn from_contiguous_values(
1968        values: Vec<i32>,
1969        shape: Vec<usize>,
1970        device: Device,
1971    ) -> Result<Self, DenseTensorError> {
1972        let meta = TensorMeta::from_shape(shape, DType::I32, device);
1973        Self::from_storage(meta, values)
1974    }
1975
1976    #[must_use]
1977    pub fn id(&self) -> u64 {
1978        self.id
1979    }
1980
1981    #[must_use]
1982    pub fn storage_id(&self) -> u64 {
1983        self.storage_id
1984    }
1985
1986    #[must_use]
1987    pub fn meta(&self) -> &TensorMeta {
1988        &self.meta
1989    }
1990
1991    #[must_use]
1992    pub fn storage(&self) -> &[i32] {
1993        &self.storage
1994    }
1995
1996    /// Return the contiguous values slice.
1997    pub fn contiguous_values(&self) -> Result<&[i32], DenseTensorError> {
1998        if !self.meta.is_contiguous() {
1999            return Err(DenseTensorError::UnsupportedLayout);
2000        }
2001        let start = self.meta.storage_offset();
2002        let end = contiguous_required_len(&self.meta)?;
2003        Ok(&self.storage[start..end])
2004    }
2005
2006    #[must_use]
2007    pub fn version(&self) -> u64 {
2008        self.version
2009    }
2010}
2011
2012/// Dense tensor backed by `Vec<u8>` storage (0=false, 1=true).
2013///
2014/// Bool tensors do NOT participate in autograd (`requires_grad` is always false).
2015#[derive(Debug, Clone, PartialEq, Eq)]
2016pub struct DenseBoolTensor {
2017    id: u64,
2018    storage_id: u64,
2019    meta: TensorMeta,
2020    storage: Vec<u8>,
2021    version: u64,
2022}
2023
2024impl DenseBoolTensor {
2025    pub fn from_storage(meta: TensorMeta, storage: Vec<u8>) -> Result<Self, DenseTensorError> {
2026        meta.validate()?;
2027        if meta.dtype() != DType::Bool {
2028            return Err(DenseTensorError::UnsupportedDType(meta.dtype()));
2029        }
2030        if !meta.is_contiguous() {
2031            return Err(DenseTensorError::UnsupportedLayout);
2032        }
2033        let needed = contiguous_required_len(&meta)?;
2034        if storage.len() < needed {
2035            return Err(DenseTensorError::InsufficientStorage {
2036                needed,
2037                actual: storage.len(),
2038            });
2039        }
2040        Ok(Self {
2041            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
2042            storage_id: NEXT_STORAGE_ID.fetch_add(1, Ordering::Relaxed),
2043            meta,
2044            storage,
2045            version: 0,
2046        })
2047    }
2048
2049    pub fn from_bools(
2050        values: &[bool],
2051        shape: Vec<usize>,
2052        device: Device,
2053    ) -> Result<Self, DenseTensorError> {
2054        let storage: Vec<u8> = values.iter().map(|&b| u8::from(b)).collect();
2055        let meta = TensorMeta::from_shape(shape, DType::Bool, device);
2056        Self::from_storage(meta, storage)
2057    }
2058
2059    #[must_use]
2060    pub fn id(&self) -> u64 {
2061        self.id
2062    }
2063
2064    #[must_use]
2065    pub fn storage_id(&self) -> u64 {
2066        self.storage_id
2067    }
2068
2069    #[must_use]
2070    pub fn meta(&self) -> &TensorMeta {
2071        &self.meta
2072    }
2073
2074    #[must_use]
2075    pub fn storage(&self) -> &[u8] {
2076        &self.storage
2077    }
2078
2079    /// Return the contiguous values slice as u8 (0=false, 1=true).
2080    pub fn contiguous_values(&self) -> Result<&[u8], DenseTensorError> {
2081        if !self.meta.is_contiguous() {
2082            return Err(DenseTensorError::UnsupportedLayout);
2083        }
2084        let start = self.meta.storage_offset();
2085        let end = contiguous_required_len(&self.meta)?;
2086        Ok(&self.storage[start..end])
2087    }
2088
2089    /// Return the contiguous values as a Vec<bool>.
2090    pub fn contiguous_bools(&self) -> Result<Vec<bool>, DenseTensorError> {
2091        let values = self.contiguous_values()?;
2092        Ok(values.iter().map(|&v| v != 0).collect())
2093    }
2094
2095    #[must_use]
2096    pub fn version(&self) -> u64 {
2097        self.version
2098    }
2099}
2100
2101pub fn ensure_compatible(lhs: &ScalarTensor, rhs: &ScalarTensor) -> Result<(), TensorCompatError> {
2102    if lhs.meta().dtype() != rhs.meta().dtype() {
2103        return Err(TensorCompatError::DTypeMismatch {
2104            lhs: lhs.meta().dtype(),
2105            rhs: rhs.meta().dtype(),
2106        });
2107    }
2108
2109    if lhs.meta().device() != rhs.meta().device() {
2110        return Err(TensorCompatError::DeviceMismatch {
2111            lhs: lhs.meta().device(),
2112            rhs: rhs.meta().device(),
2113        });
2114    }
2115
2116    Ok(())
2117}
2118
2119#[must_use]
2120pub fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
2121    if shape.is_empty() {
2122        return Vec::new();
2123    }
2124
2125    let mut strides = vec![1; shape.len()];
2126    let mut running = 1usize;
2127    for idx in (0..shape.len()).rev() {
2128        strides[idx] = running;
2129        running = running.saturating_mul(shape[idx]);
2130    }
2131    strides
2132}
2133
2134fn saturated_numel(shape: &[usize]) -> usize {
2135    if shape.is_empty() {
2136        return 1;
2137    }
2138    if shape.contains(&0) {
2139        return 0;
2140    }
2141    let mut product = 1usize;
2142    for dim in shape.iter().copied() {
2143        let Some(next) = product.checked_mul(dim) else {
2144            return usize::MAX;
2145        };
2146        product = next;
2147    }
2148    product
2149}
2150
2151// ── Sparse Tensor Types ────────────────────────────────────────────────
2152
2153/// Error type for sparse tensor operations.
2154#[derive(Debug, Clone, PartialEq, Eq)]
2155pub enum SparseTensorError {
2156    /// Indices tensor has wrong rank (expected 2 for COO).
2157    InvalidIndicesRank { expected: usize, actual: usize },
2158    /// Indices tensor has wrong dtype (must be I64).
2159    InvalidIndicesDType { actual: DType },
2160    /// Indices tensor is on a different device than values.
2161    DeviceMismatch { expected: Device, actual: Device },
2162    /// Indices sparse_dim doesn't match the dense shape.
2163    SparseDimMismatch {
2164        indices_sparse_dim: usize,
2165        expected: usize,
2166    },
2167    /// Number of non-zero entries doesn't match between indices and values.
2168    NnzMismatch {
2169        indices_nnz: usize,
2170        values_nnz: usize,
2171    },
2172    /// Values tensor shape doesn't match expected [nnz, *dense_dims].
2173    InvalidValuesShape {
2174        expected: Vec<usize>,
2175        actual: Vec<usize>,
2176    },
2177    /// Index out of bounds for the dense shape.
2178    IndexOutOfBounds { dim: usize, index: i64, size: usize },
2179    /// Negative index in indices tensor.
2180    NegativeIndex { dim: usize, index: i64 },
2181    /// CSR crow_indices has wrong length (expected nrows + 1).
2182    InvalidCrowIndicesLen { expected: usize, actual: usize },
2183    /// CSR col_indices has wrong length (expected nnz).
2184    InvalidColIndicesLen { expected: usize, actual: usize },
2185    /// CSR crow_indices values are not monotonically increasing.
2186    NonMonotonicCrowIndices { row: usize, prev: i64, curr: i64 },
2187    /// CSR crow_indices contains an invalid value.
2188    InvalidCrowIndexValue { index: usize, value: i64 },
2189    /// CSR column index out of bounds.
2190    ColIndexOutOfBounds { index: i64, ncols: usize },
2191    /// CSR row contains the same column index more than once.
2192    DuplicateCsrColumn { row: usize, col: i64 },
2193    /// A coalesced COO tensor contains duplicate coordinates.
2194    DuplicateCooIndex { position: usize, coord: Vec<i64> },
2195    /// A coalesced COO tensor contains coordinates out of lexicographic order.
2196    UnsortedCooIndex {
2197        position: usize,
2198        previous: Vec<i64>,
2199        current: Vec<i64>,
2200    },
2201    /// Dense tensor error during conversion.
2202    DenseTensor(DenseTensorError),
2203    /// Only 2D sparse CSR tensors are supported.
2204    UnsupportedRank { rank: usize },
2205}
2206
2207impl fmt::Display for SparseTensorError {
2208    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2209        match self {
2210            Self::InvalidIndicesRank { expected, actual } => {
2211                write!(f, "indices tensor has rank {actual}, expected {expected}")
2212            }
2213            Self::InvalidIndicesDType { actual } => {
2214                write!(f, "indices tensor has dtype {actual:?}, expected I64")
2215            }
2216            Self::DeviceMismatch { expected, actual } => {
2217                write!(
2218                    f,
2219                    "sparse tensor device mismatch: expected {expected:?}, got {actual:?}"
2220                )
2221            }
2222            Self::SparseDimMismatch {
2223                indices_sparse_dim,
2224                expected,
2225            } => {
2226                write!(
2227                    f,
2228                    "indices sparse_dim is {indices_sparse_dim}, expected {expected}"
2229                )
2230            }
2231            Self::NnzMismatch {
2232                indices_nnz,
2233                values_nnz,
2234            } => {
2235                write!(
2236                    f,
2237                    "indices nnz ({indices_nnz}) != values nnz ({values_nnz})"
2238                )
2239            }
2240            Self::InvalidValuesShape { expected, actual } => {
2241                write!(f, "values shape is {actual:?}, expected {expected:?}")
2242            }
2243            Self::IndexOutOfBounds { dim, index, size } => {
2244                write!(
2245                    f,
2246                    "index {index} at dim {dim} out of bounds for size {size}"
2247                )
2248            }
2249            Self::NegativeIndex { dim, index } => {
2250                write!(f, "negative index {index} at dim {dim}")
2251            }
2252            Self::InvalidCrowIndicesLen { expected, actual } => {
2253                write!(f, "crow_indices length is {actual}, expected {expected}")
2254            }
2255            Self::InvalidColIndicesLen { expected, actual } => {
2256                write!(f, "col_indices length is {actual}, expected {expected}")
2257            }
2258            Self::NonMonotonicCrowIndices { row, prev, curr } => {
2259                write!(
2260                    f,
2261                    "crow_indices not monotonic at row {row}: {prev} > {curr}"
2262                )
2263            }
2264            Self::InvalidCrowIndexValue { index, value } => {
2265                write!(f, "crow_indices[{index}] has invalid value {value}")
2266            }
2267            Self::ColIndexOutOfBounds { index, ncols } => {
2268                write!(f, "column index {index} out of bounds for {ncols} columns")
2269            }
2270            Self::DuplicateCsrColumn { row, col } => {
2271                write!(f, "CSR row {row} contains duplicate column index {col}")
2272            }
2273            Self::DuplicateCooIndex { position, coord } => {
2274                write!(
2275                    f,
2276                    "coalesced COO index at position {position} duplicates coordinate {coord:?}"
2277                )
2278            }
2279            Self::UnsortedCooIndex {
2280                position,
2281                previous,
2282                current,
2283            } => {
2284                write!(
2285                    f,
2286                    "coalesced COO index at position {position} is out of order: {previous:?} > {current:?}"
2287                )
2288            }
2289            Self::DenseTensor(err) => write!(f, "dense tensor error: {err}"),
2290            Self::UnsupportedRank { rank } => {
2291                write!(f, "sparse CSR only supports 2D tensors, got rank {rank}")
2292            }
2293        }
2294    }
2295}
2296
2297impl std::error::Error for SparseTensorError {}
2298
2299impl From<DenseTensorError> for SparseTensorError {
2300    fn from(value: DenseTensorError) -> Self {
2301        Self::DenseTensor(value)
2302    }
2303}
2304
2305fn coo_coordinate(
2306    indices_values: &[i64],
2307    sparse_dim: usize,
2308    nnz: usize,
2309    position: usize,
2310) -> Vec<i64> {
2311    (0..sparse_dim)
2312        .map(|dim| indices_values[dim * nnz + position])
2313        .collect()
2314}
2315
2316fn dense_tensor_from_complex128_values(
2317    values: Vec<Complex128>,
2318    shape: Vec<usize>,
2319    dtype: DType,
2320    device: Device,
2321) -> Result<DenseTensor, DenseTensorError> {
2322    let meta = TensorMeta::from_shape(shape, dtype, device);
2323    let storage = match dtype {
2324        DType::Complex64 => {
2325            #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
2326            {
2327                TensorStorage::Complex64(Arc::new(
2328                    values
2329                        .into_iter()
2330                        .map(|value| Complex64::new(value.re as f32, value.im as f32))
2331                        .collect(),
2332                ))
2333            }
2334        }
2335        DType::Complex128 => TensorStorage::Complex128(Arc::new(values)),
2336        _ => return Err(DenseTensorError::UnsupportedDType(dtype)),
2337    };
2338    DenseTensor::from_typed_storage(meta, storage)
2339}
2340
2341/// Sparse tensor in COO (Coordinate) format.
2342///
2343/// COO format stores sparse tensors as a pair of tensors:
2344/// - `indices`: shape [sparse_dim, nnz], dtype I64 — the coordinates of non-zero elements
2345/// - `values`: shape [nnz, *dense_dims] — the values at those coordinates
2346///
2347/// For a sparse tensor with shape [3, 4, 5] and sparse_dim=2:
2348/// - indices has shape [2, nnz] (row and column indices)
2349/// - values has shape [nnz, 5] (the remaining dense dimension)
2350///
2351/// The `coalesced` flag indicates whether duplicate indices have been merged.
2352/// A coalesced tensor has unique, sorted indices.
2353#[derive(Debug, Clone, PartialEq)]
2354pub struct SparseCOOTensor {
2355    id: u64,
2356    /// The indices of non-zero elements, shape [sparse_dim, nnz], dtype I64.
2357    indices: DenseI64Tensor,
2358    /// The values at the indexed positions, shape [nnz, *dense_dims].
2359    values: DenseTensor,
2360    /// The full dense shape this sparse tensor represents.
2361    dense_shape: Vec<usize>,
2362    /// Number of sparse dimensions (indices.shape[0]).
2363    sparse_dim: usize,
2364    /// Whether the tensor is coalesced (no duplicate indices, sorted).
2365    coalesced: bool,
2366    device: Device,
2367    version: u64,
2368}
2369
2370impl SparseCOOTensor {
2371    /// Create a new sparse COO tensor.
2372    ///
2373    /// # Arguments
2374    /// * `indices` - shape [sparse_dim, nnz], dtype I64
2375    /// * `values` - shape [nnz, *dense_dims], any floating-point dtype
2376    /// * `dense_shape` - the full shape if this were a dense tensor
2377    /// * `coalesced` - whether indices are unique and sorted
2378    ///
2379    /// # Errors
2380    /// Returns error if indices/values shapes don't match or indices are out of bounds.
2381    pub fn new(
2382        indices: DenseI64Tensor,
2383        values: DenseTensor,
2384        dense_shape: Vec<usize>,
2385        coalesced: bool,
2386    ) -> Result<Self, SparseTensorError> {
2387        // Validate indices shape: must be [sparse_dim, nnz]
2388        let indices_shape = indices.meta().shape();
2389        if indices_shape.len() != 2 {
2390            return Err(SparseTensorError::InvalidIndicesRank {
2391                expected: 2,
2392                actual: indices_shape.len(),
2393            });
2394        }
2395
2396        let sparse_dim = indices_shape[0];
2397        let nnz = indices_shape[1];
2398
2399        // sparse_dim must not exceed dense_shape rank
2400        if sparse_dim > dense_shape.len() {
2401            return Err(SparseTensorError::SparseDimMismatch {
2402                indices_sparse_dim: sparse_dim,
2403                expected: dense_shape.len(),
2404            });
2405        }
2406
2407        // Validate values shape: must be [nnz, *dense_dims]
2408        let values_shape = values.meta().shape();
2409        if values_shape.is_empty() {
2410            return Err(SparseTensorError::InvalidValuesShape {
2411                expected: vec![nnz],
2412                actual: values_shape.to_vec(),
2413            });
2414        }
2415
2416        if values_shape[0] != nnz {
2417            return Err(SparseTensorError::NnzMismatch {
2418                indices_nnz: nnz,
2419                values_nnz: values_shape[0],
2420            });
2421        }
2422
2423        // dense_dims are the remaining dimensions after sparse_dim
2424        let dense_dims = &dense_shape[sparse_dim..];
2425        let expected_values_shape: Vec<usize> = std::iter::once(nnz)
2426            .chain(dense_dims.iter().copied())
2427            .collect();
2428
2429        if values_shape != expected_values_shape.as_slice() {
2430            return Err(SparseTensorError::InvalidValuesShape {
2431                expected: expected_values_shape,
2432                actual: values_shape.to_vec(),
2433            });
2434        }
2435
2436        let device = values.meta().device();
2437        if indices.meta().device() != device {
2438            return Err(SparseTensorError::DeviceMismatch {
2439                expected: device,
2440                actual: indices.meta().device(),
2441            });
2442        }
2443
2444        // Validate indices are within bounds even when not coalesced.
2445        let indices_values = indices.contiguous_values()?;
2446        for d in 0..sparse_dim {
2447            let dim_size = dense_shape[d];
2448            for i in 0..nnz {
2449                let idx = indices_values[d * nnz + i];
2450                if idx < 0 {
2451                    return Err(SparseTensorError::NegativeIndex { dim: d, index: idx });
2452                }
2453                if (idx as usize) >= dim_size {
2454                    return Err(SparseTensorError::IndexOutOfBounds {
2455                        dim: d,
2456                        index: idx,
2457                        size: dim_size,
2458                    });
2459                }
2460            }
2461        }
2462
2463        if coalesced && nnz > 1 {
2464            for position in 1..nnz {
2465                let mut ordering = std::cmp::Ordering::Equal;
2466                for dim in 0..sparse_dim {
2467                    let previous = indices_values[dim * nnz + position - 1];
2468                    let current = indices_values[dim * nnz + position];
2469                    ordering = previous.cmp(&current);
2470                    if !ordering.is_eq() {
2471                        break;
2472                    }
2473                }
2474                match ordering {
2475                    std::cmp::Ordering::Less => {}
2476                    std::cmp::Ordering::Equal => {
2477                        return Err(SparseTensorError::DuplicateCooIndex {
2478                            position,
2479                            coord: coo_coordinate(indices_values, sparse_dim, nnz, position),
2480                        });
2481                    }
2482                    std::cmp::Ordering::Greater => {
2483                        return Err(SparseTensorError::UnsortedCooIndex {
2484                            position,
2485                            previous: coo_coordinate(indices_values, sparse_dim, nnz, position - 1),
2486                            current: coo_coordinate(indices_values, sparse_dim, nnz, position),
2487                        });
2488                    }
2489                }
2490            }
2491        }
2492
2493        Ok(Self {
2494            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
2495            indices,
2496            values,
2497            dense_shape,
2498            sparse_dim,
2499            coalesced,
2500            device,
2501            version: 0,
2502        })
2503    }
2504
2505    /// Create a sparse COO tensor from coordinate lists.
2506    ///
2507    /// # Arguments
2508    /// * `coords` - list of coordinate tuples, each of length sparse_dim
2509    /// * `values` - flat values for each coordinate
2510    /// * `dense_shape` - the full dense shape
2511    /// * `dtype` - dtype for the values
2512    /// * `device` - device for the tensor
2513    pub fn from_coords(
2514        coords: &[Vec<i64>],
2515        values: Vec<f64>,
2516        dense_shape: Vec<usize>,
2517        dtype: DType,
2518        device: Device,
2519    ) -> Result<Self, SparseTensorError> {
2520        let nnz = coords.len();
2521        if nnz == 0 {
2522            // Empty sparse tensor
2523            let sparse_dim = dense_shape.len();
2524            let indices =
2525                DenseI64Tensor::from_contiguous_values(vec![], vec![sparse_dim, 0], device)?;
2526            let values_tensor =
2527                DenseTensor::from_contiguous_values(vec![], vec![0], device)?.to_dtype(dtype)?;
2528            return Self::new(indices, values_tensor, dense_shape, true);
2529        }
2530
2531        let sparse_dim = coords[0].len();
2532        if sparse_dim > dense_shape.len() {
2533            return Err(SparseTensorError::SparseDimMismatch {
2534                indices_sparse_dim: sparse_dim,
2535                expected: dense_shape.len(),
2536            });
2537        }
2538
2539        // Build indices tensor [sparse_dim, nnz]
2540        let mut indices_data = vec![0i64; sparse_dim * nnz];
2541        for (i, coord) in coords.iter().enumerate() {
2542            if coord.len() != sparse_dim {
2543                return Err(SparseTensorError::SparseDimMismatch {
2544                    indices_sparse_dim: coord.len(),
2545                    expected: sparse_dim,
2546                });
2547            }
2548            for (d, &idx) in coord.iter().enumerate() {
2549                indices_data[d * nnz + i] = idx;
2550            }
2551        }
2552
2553        let indices =
2554            DenseI64Tensor::from_contiguous_values(indices_data, vec![sparse_dim, nnz], device)?;
2555
2556        let values_shape: Vec<usize> = std::iter::once(nnz)
2557            .chain(dense_shape[sparse_dim..].iter().copied())
2558            .collect();
2559        let values_tensor = DenseTensor::from_contiguous_values(values, values_shape, device)?;
2560        let values_tensor = values_tensor.to_dtype(dtype)?;
2561
2562        Self::new(indices, values_tensor, dense_shape, false)
2563    }
2564
2565    #[must_use]
2566    pub fn id(&self) -> u64 {
2567        self.id
2568    }
2569
2570    #[must_use]
2571    pub fn indices(&self) -> &DenseI64Tensor {
2572        &self.indices
2573    }
2574
2575    #[must_use]
2576    pub fn values(&self) -> &DenseTensor {
2577        &self.values
2578    }
2579
2580    #[must_use]
2581    pub fn dense_shape(&self) -> &[usize] {
2582        &self.dense_shape
2583    }
2584
2585    #[must_use]
2586    pub fn sparse_dim(&self) -> usize {
2587        self.sparse_dim
2588    }
2589
2590    #[must_use]
2591    pub fn nnz(&self) -> usize {
2592        self.indices.meta().shape().get(1).copied().unwrap_or(0)
2593    }
2594
2595    #[must_use]
2596    pub fn is_coalesced(&self) -> bool {
2597        self.coalesced
2598    }
2599
2600    #[must_use]
2601    pub fn dtype(&self) -> DType {
2602        self.values.meta().dtype()
2603    }
2604
2605    #[must_use]
2606    pub fn device(&self) -> Device {
2607        self.device
2608    }
2609
2610    #[must_use]
2611    pub fn version(&self) -> u64 {
2612        self.version
2613    }
2614
2615    /// Convert this sparse tensor to a dense tensor.
2616    ///
2617    /// This allocates a new dense tensor with zeros and fills in the non-zero values.
2618    pub fn to_dense(&self) -> Result<DenseTensor, SparseTensorError> {
2619        let numel = checked_shape_numel(&self.dense_shape)?;
2620        if self.dtype().is_complex() {
2621            return self.to_dense_complex(numel);
2622        }
2623
2624        let mut dense_data = vec![0.0f64; numel];
2625
2626        let nnz = self.nnz();
2627        if nnz == 0 {
2628            let result = DenseTensor::from_contiguous_values(
2629                dense_data,
2630                self.dense_shape.clone(),
2631                self.device,
2632            )?;
2633            return Ok(result.to_dtype(self.dtype())?);
2634        }
2635
2636        let indices_data = self.indices.contiguous_values()?;
2637        let values_data = self.values.contiguous_values_as_f64()?;
2638        let strides = contiguous_strides(&self.dense_shape);
2639
2640        // For now, only support fully sparse tensors (sparse_dim == rank)
2641        // where values are scalars
2642        if self.sparse_dim == self.dense_shape.len() {
2643            for i in 0..nnz {
2644                let mut linear_idx = 0usize;
2645                for d in 0..self.sparse_dim {
2646                    let idx = indices_data[d * nnz + i];
2647                    if idx < 0 || (idx as usize) >= self.dense_shape[d] {
2648                        return Err(SparseTensorError::IndexOutOfBounds {
2649                            dim: d,
2650                            index: idx,
2651                            size: self.dense_shape[d],
2652                        });
2653                    }
2654                    linear_idx += (idx as usize) * strides[d];
2655                }
2656                dense_data[linear_idx] += values_data[i];
2657            }
2658        } else {
2659            // Hybrid sparse-dense: values have shape [nnz, *dense_dims]
2660            let dense_dims = &self.dense_shape[self.sparse_dim..];
2661            let dense_numel = checked_shape_numel(dense_dims)?;
2662
2663            for i in 0..nnz {
2664                let mut sparse_linear_idx = 0usize;
2665                for d in 0..self.sparse_dim {
2666                    let idx = indices_data[d * nnz + i];
2667                    if idx < 0 || (idx as usize) >= self.dense_shape[d] {
2668                        return Err(SparseTensorError::IndexOutOfBounds {
2669                            dim: d,
2670                            index: idx,
2671                            size: self.dense_shape[d],
2672                        });
2673                    }
2674                    sparse_linear_idx += (idx as usize) * strides[d];
2675                }
2676                // Copy the dense slice
2677                for j in 0..dense_numel {
2678                    dense_data[sparse_linear_idx + j] += values_data[i * dense_numel + j];
2679                }
2680            }
2681        }
2682
2683        let result =
2684            DenseTensor::from_contiguous_values(dense_data, self.dense_shape.clone(), self.device)?;
2685        Ok(result.to_dtype(self.dtype())?)
2686    }
2687
2688    fn to_dense_complex(&self, numel: usize) -> Result<DenseTensor, SparseTensorError> {
2689        let mut dense_data = vec![Complex128::new(0.0, 0.0); numel];
2690
2691        let nnz = self.nnz();
2692        if nnz == 0 {
2693            return Ok(dense_tensor_from_complex128_values(
2694                dense_data,
2695                self.dense_shape.clone(),
2696                self.dtype(),
2697                self.device,
2698            )?);
2699        }
2700
2701        let indices_data = self.indices.contiguous_values()?;
2702        let values_data = self.values.contiguous_complex_values_as_complex128()?;
2703        let strides = contiguous_strides(&self.dense_shape);
2704
2705        if self.sparse_dim == self.dense_shape.len() {
2706            for i in 0..nnz {
2707                let mut linear_idx = 0usize;
2708                for d in 0..self.sparse_dim {
2709                    let idx = indices_data[d * nnz + i];
2710                    if idx < 0 || (idx as usize) >= self.dense_shape[d] {
2711                        return Err(SparseTensorError::IndexOutOfBounds {
2712                            dim: d,
2713                            index: idx,
2714                            size: self.dense_shape[d],
2715                        });
2716                    }
2717                    linear_idx += (idx as usize) * strides[d];
2718                }
2719                dense_data[linear_idx] += values_data[i];
2720            }
2721        } else {
2722            let dense_dims = &self.dense_shape[self.sparse_dim..];
2723            let dense_numel = checked_shape_numel(dense_dims)?;
2724
2725            for i in 0..nnz {
2726                let mut sparse_linear_idx = 0usize;
2727                for d in 0..self.sparse_dim {
2728                    let idx = indices_data[d * nnz + i];
2729                    if idx < 0 || (idx as usize) >= self.dense_shape[d] {
2730                        return Err(SparseTensorError::IndexOutOfBounds {
2731                            dim: d,
2732                            index: idx,
2733                            size: self.dense_shape[d],
2734                        });
2735                    }
2736                    sparse_linear_idx += (idx as usize) * strides[d];
2737                }
2738                for j in 0..dense_numel {
2739                    dense_data[sparse_linear_idx + j] += values_data[i * dense_numel + j];
2740                }
2741            }
2742        }
2743
2744        Ok(dense_tensor_from_complex128_values(
2745            dense_data,
2746            self.dense_shape.clone(),
2747            self.dtype(),
2748            self.device,
2749        )?)
2750    }
2751}
2752
2753/// Sparse tensor in CSR (Compressed Sparse Row) format.
2754///
2755/// CSR format is efficient for row-wise operations on 2D sparse matrices:
2756/// - `crow_indices`: shape [nrows + 1], dtype I64 — row pointers
2757/// - `col_indices`: shape [nnz], dtype I64 — column indices for each non-zero
2758/// - `values`: shape [nnz] — the non-zero values
2759///
2760/// For row `i`, the non-zeros are at positions crow_indices[i]..crow_indices[i+1]
2761/// in col_indices and values.
2762#[derive(Debug, Clone, PartialEq)]
2763pub struct SparseCSRTensor {
2764    id: u64,
2765    /// Row pointers, shape [nrows + 1], dtype I64.
2766    crow_indices: DenseI64Tensor,
2767    /// Column indices for each non-zero, shape [nnz], dtype I64.
2768    col_indices: DenseI64Tensor,
2769    /// Values at each non-zero position, shape [nnz].
2770    values: DenseTensor,
2771    /// The 2D shape [nrows, ncols].
2772    shape: [usize; 2],
2773    device: Device,
2774    version: u64,
2775}
2776
2777impl SparseCSRTensor {
2778    /// Create a new sparse CSR tensor.
2779    ///
2780    /// # Arguments
2781    /// * `crow_indices` - shape [nrows + 1], dtype I64
2782    /// * `col_indices` - shape [nnz], dtype I64
2783    /// * `values` - shape [nnz], any floating-point dtype
2784    /// * `shape` - [nrows, ncols]
2785    ///
2786    /// # Errors
2787    /// Returns error if shapes don't match or indices are invalid.
2788    pub fn new(
2789        crow_indices: DenseI64Tensor,
2790        col_indices: DenseI64Tensor,
2791        values: DenseTensor,
2792        shape: [usize; 2],
2793    ) -> Result<Self, SparseTensorError> {
2794        let [nrows, ncols] = shape;
2795
2796        // Validate crow_indices: shape [nrows + 1]
2797        let crow_shape = crow_indices.meta().shape();
2798        if crow_shape.len() != 1 {
2799            return Err(SparseTensorError::UnsupportedRank {
2800                rank: crow_shape.len(),
2801            });
2802        }
2803        let expected_crow_len = nrows.checked_add(1).ok_or_else(|| {
2804            SparseTensorError::DenseTensor(DenseTensorError::ShapeOverflow {
2805                shape: vec![nrows, ncols],
2806            })
2807        })?;
2808        if crow_shape[0] != expected_crow_len {
2809            return Err(SparseTensorError::InvalidCrowIndicesLen {
2810                expected: expected_crow_len,
2811                actual: crow_shape[0],
2812            });
2813        }
2814
2815        let crow_data = crow_indices.contiguous_values()?;
2816
2817        if crow_data[0] != 0 {
2818            return Err(SparseTensorError::InvalidCrowIndexValue {
2819                index: 0,
2820                value: crow_data[0],
2821            });
2822        }
2823
2824        let nnz = usize::try_from(crow_data[nrows]).map_err(|_| {
2825            SparseTensorError::InvalidCrowIndexValue {
2826                index: nrows,
2827                value: crow_data[nrows],
2828            }
2829        })?;
2830
2831        // Validate col_indices: shape [nnz]
2832        let col_shape = col_indices.meta().shape();
2833        if col_shape.len() != 1 {
2834            return Err(SparseTensorError::UnsupportedRank {
2835                rank: col_shape.len(),
2836            });
2837        }
2838        if col_shape[0] != nnz {
2839            return Err(SparseTensorError::InvalidColIndicesLen {
2840                expected: nnz,
2841                actual: col_shape[0],
2842            });
2843        }
2844
2845        // Validate values: shape [nnz]
2846        let values_shape = values.meta().shape();
2847        if values_shape.len() != 1 {
2848            return Err(SparseTensorError::UnsupportedRank {
2849                rank: values_shape.len(),
2850            });
2851        }
2852        if values_shape[0] != nnz {
2853            return Err(SparseTensorError::NnzMismatch {
2854                indices_nnz: nnz,
2855                values_nnz: values_shape[0],
2856            });
2857        }
2858
2859        let device = values.meta().device();
2860        if crow_indices.meta().device() != device {
2861            return Err(SparseTensorError::DeviceMismatch {
2862                expected: device,
2863                actual: crow_indices.meta().device(),
2864            });
2865        }
2866        if col_indices.meta().device() != device {
2867            return Err(SparseTensorError::DeviceMismatch {
2868                expected: device,
2869                actual: col_indices.meta().device(),
2870            });
2871        }
2872
2873        // Validate crow_indices is monotonically increasing and within bounds.
2874        for i in 0..=nrows {
2875            let value = crow_data[i];
2876            if value < 0 {
2877                return Err(SparseTensorError::InvalidCrowIndexValue { index: i, value });
2878            }
2879            if i > 0 && crow_data[i - 1] > value {
2880                return Err(SparseTensorError::NonMonotonicCrowIndices {
2881                    row: i - 1,
2882                    prev: crow_data[i - 1],
2883                    curr: value,
2884                });
2885            }
2886        }
2887
2888        // Validate col_indices are in bounds
2889        let col_data = col_indices.contiguous_values()?;
2890        for &col in col_data {
2891            if col < 0 || (col as usize) >= ncols {
2892                return Err(SparseTensorError::ColIndexOutOfBounds { index: col, ncols });
2893            }
2894        }
2895        for row in 0..nrows {
2896            let start = usize::try_from(crow_data[row]).map_err(|_| {
2897                SparseTensorError::InvalidCrowIndexValue {
2898                    index: row,
2899                    value: crow_data[row],
2900                }
2901            })?;
2902            let end = usize::try_from(crow_data[row + 1]).map_err(|_| {
2903                SparseTensorError::InvalidCrowIndexValue {
2904                    index: row + 1,
2905                    value: crow_data[row + 1],
2906                }
2907            })?;
2908            for idx in start..end {
2909                let col = col_data[idx];
2910                if col_data[start..idx].contains(&col) {
2911                    return Err(SparseTensorError::DuplicateCsrColumn { row, col });
2912                }
2913            }
2914        }
2915
2916        Ok(Self {
2917            id: NEXT_TENSOR_ID.fetch_add(1, Ordering::Relaxed),
2918            crow_indices,
2919            col_indices,
2920            values,
2921            shape,
2922            device,
2923            version: 0,
2924        })
2925    }
2926
2927    #[must_use]
2928    pub fn id(&self) -> u64 {
2929        self.id
2930    }
2931
2932    #[must_use]
2933    pub fn crow_indices(&self) -> &DenseI64Tensor {
2934        &self.crow_indices
2935    }
2936
2937    #[must_use]
2938    pub fn col_indices(&self) -> &DenseI64Tensor {
2939        &self.col_indices
2940    }
2941
2942    #[must_use]
2943    pub fn values(&self) -> &DenseTensor {
2944        &self.values
2945    }
2946
2947    #[must_use]
2948    pub fn shape(&self) -> [usize; 2] {
2949        self.shape
2950    }
2951
2952    #[must_use]
2953    pub fn nrows(&self) -> usize {
2954        self.shape[0]
2955    }
2956
2957    #[must_use]
2958    pub fn ncols(&self) -> usize {
2959        self.shape[1]
2960    }
2961
2962    #[must_use]
2963    pub fn nnz(&self) -> usize {
2964        self.col_indices
2965            .meta()
2966            .shape()
2967            .first()
2968            .copied()
2969            .unwrap_or(0)
2970    }
2971
2972    #[must_use]
2973    pub fn dtype(&self) -> DType {
2974        self.values.meta().dtype()
2975    }
2976
2977    #[must_use]
2978    pub fn device(&self) -> Device {
2979        self.device
2980    }
2981
2982    #[must_use]
2983    pub fn version(&self) -> u64 {
2984        self.version
2985    }
2986
2987    /// Convert this sparse CSR tensor to a dense tensor.
2988    pub fn to_dense(&self) -> Result<DenseTensor, SparseTensorError> {
2989        let [nrows, ncols] = self.shape;
2990        let numel = nrows
2991            .checked_mul(ncols)
2992            .ok_or_else(|| DenseTensorError::ShapeOverflow {
2993                shape: vec![nrows, ncols],
2994            })?;
2995        if self.dtype().is_complex() {
2996            return self.to_dense_complex(numel);
2997        }
2998
2999        let mut dense_data = vec![0.0f64; numel];
3000
3001        let crow_data = self.crow_indices.contiguous_values()?;
3002        let col_data = self.col_indices.contiguous_values()?;
3003        let values_data = self.values.contiguous_values_as_f64()?;
3004
3005        for row in 0..nrows {
3006            let start = crow_data[row] as usize;
3007            let end = crow_data[row + 1] as usize;
3008            for idx in start..end {
3009                let col = col_data[idx] as usize;
3010                dense_data[row * ncols + col] = values_data[idx];
3011            }
3012        }
3013
3014        let result =
3015            DenseTensor::from_contiguous_values(dense_data, vec![nrows, ncols], self.device)?;
3016        Ok(result.to_dtype(self.dtype())?)
3017    }
3018
3019    fn to_dense_complex(&self, numel: usize) -> Result<DenseTensor, SparseTensorError> {
3020        let [nrows, ncols] = self.shape;
3021        let mut dense_data = vec![Complex128::new(0.0, 0.0); numel];
3022
3023        let crow_data = self.crow_indices.contiguous_values()?;
3024        let col_data = self.col_indices.contiguous_values()?;
3025        let values_data = self.values.contiguous_complex_values_as_complex128()?;
3026
3027        for row in 0..nrows {
3028            let start = crow_data[row] as usize;
3029            let end = crow_data[row + 1] as usize;
3030            for idx in start..end {
3031                let col = col_data[idx] as usize;
3032                dense_data[row * ncols + col] = values_data[idx];
3033            }
3034        }
3035
3036        Ok(dense_tensor_from_complex128_values(
3037            dense_data,
3038            vec![nrows, ncols],
3039            self.dtype(),
3040            self.device,
3041        )?)
3042    }
3043}
3044
3045#[cfg(test)]
3046mod tests {
3047    use std::collections::BTreeMap;
3048
3049    use proptest::prelude::*;
3050
3051    use std::sync::Arc;
3052
3053    use super::{
3054        BFloat16, Complex64, Complex128, DType, DenseBoolTensor, DenseI32Tensor, DenseI64Tensor,
3055        DenseTensor, DenseTensorError, Device, Float16, QuantizationParams, ScalarTensor,
3056        SparseCOOTensor, SparseCSRTensor, SparseTensorError, TensorMeta, TensorMetaError,
3057        TensorStorage, contiguous_strides, ensure_compatible,
3058    };
3059
3060    fn det_seed(parts: &[usize]) -> u64 {
3061        let mut hash = 0xcbf2_9ce4_8422_2325u64;
3062        for value in parts {
3063            for byte in value.to_le_bytes() {
3064                hash ^= u64::from(byte);
3065                hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
3066            }
3067        }
3068        hash
3069    }
3070
3071    fn build_property_log(
3072        test_id: &str,
3073        mode: &str,
3074        seed: u64,
3075        input_digest: u64,
3076        output_digest: u64,
3077        reason_code: &str,
3078    ) -> BTreeMap<String, String> {
3079        let mut log = BTreeMap::new();
3080        let scenario_id = format!("ft_core_property/{mode}:{test_id}");
3081        log.insert("ts_utc".to_string(), "1970-01-01T00:00:00Z".to_string());
3082        log.insert("suite_id".to_string(), "ft_core_property".to_string());
3083        log.insert("test_id".to_string(), test_id.to_string());
3084        log.insert("packet_id".to_string(), "FT-P2C-001".to_string());
3085        log.insert(
3086            "fixture_id".to_string(),
3087            "ft_core_property_generated".to_string(),
3088        );
3089        log.insert("scenario_id".to_string(), scenario_id);
3090        log.insert("mode".to_string(), mode.to_string());
3091        log.insert("seed".to_string(), seed.to_string());
3092        log.insert(
3093            "input_digest".to_string(),
3094            format!("det64:{input_digest:016x}"),
3095        );
3096        log.insert(
3097            "output_digest".to_string(),
3098            format!("det64:{output_digest:016x}"),
3099        );
3100        log.insert(
3101            "env_fingerprint".to_string(),
3102            "det64:ft-core-test".to_string(),
3103        );
3104        log.insert(
3105            "artifact_refs".to_string(),
3106            "artifacts/phase2c/FT-P2C-001/fixture_manifest.json".to_string(),
3107        );
3108        log.insert(
3109            "replay_command".to_string(),
3110            format!("cargo test -p ft-core {test_id} -- --nocapture"),
3111        );
3112        log.insert("duration_ms".to_string(), "0".to_string());
3113        log.insert("outcome".to_string(), "pass".to_string());
3114        log.insert("contract_id".to_string(), reason_code.to_string());
3115        log.insert("shrink_trace".to_string(), "none".to_string());
3116        log.insert("reason_code".to_string(), reason_code.to_string());
3117        log
3118    }
3119
3120    fn assert_log_contract(log: &BTreeMap<String, String>) {
3121        for key in [
3122            "ts_utc",
3123            "suite_id",
3124            "test_id",
3125            "packet_id",
3126            "fixture_id",
3127            "scenario_id",
3128            "mode",
3129            "seed",
3130            "input_digest",
3131            "output_digest",
3132            "env_fingerprint",
3133            "artifact_refs",
3134            "replay_command",
3135            "duration_ms",
3136            "outcome",
3137            "contract_id",
3138            "shrink_trace",
3139            "reason_code",
3140        ] {
3141            assert!(
3142                log.contains_key(key),
3143                "property log missing required key '{key}'"
3144            );
3145        }
3146    }
3147
3148    #[test]
3149    fn scalar_meta_is_valid() {
3150        let meta = TensorMeta::scalar(DType::F64, Device::Cpu);
3151        assert!(meta.validate().is_ok());
3152        assert!(meta.shape().is_empty());
3153        assert!(meta.strides().is_empty());
3154        assert_eq!(meta.numel(), 1);
3155        assert!(meta.is_contiguous());
3156    }
3157
3158    #[test]
3159    fn dense_tensor_from_contiguous_values_accepts_matching_storage() {
3160        let tensor = DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0], vec![3], Device::Cpu)
3161            .expect("contiguous dense tensor should build");
3162        assert_eq!(
3163            tensor.contiguous_values().expect("slice should resolve"),
3164            &[1.0, 2.0, 3.0]
3165        );
3166        assert_eq!(tensor.meta().shape(), &[3]);
3167        assert_eq!(tensor.meta().dtype(), DType::F64);
3168    }
3169
3170    #[test]
3171    fn dense_tensor_contiguous_view_rejects_non_contiguous_layout() {
3172        let meta =
3173            TensorMeta::from_shape_and_strides(vec![2, 2], vec![4, 1], 0, DType::F64, Device::Cpu)
3174                .expect("meta should validate");
3175        let tensor = DenseTensor::from_storage(meta, vec![1.0; 6])
3176            .expect("non-contiguous metadata should still construct tensor storage");
3177        let err = tensor
3178            .contiguous_values()
3179            .expect_err("non-contiguous layout must fail contiguous view");
3180        assert!(matches!(err, DenseTensorError::UnsupportedLayout));
3181    }
3182
3183    #[test]
3184    fn dense_tensor_rejects_insufficient_storage_for_offset() {
3185        let meta = TensorMeta::from_shape(vec![3], DType::F64, Device::Cpu).with_storage_offset(2);
3186        let err = DenseTensor::from_storage(meta, vec![1.0, 2.0, 3.0])
3187            .expect_err("offset span must require enough storage");
3188        assert!(matches!(
3189            err,
3190            DenseTensorError::InsufficientStorage {
3191                needed: 5,
3192                actual: 3
3193            }
3194        ));
3195    }
3196
3197    #[test]
3198    fn dense_tensor_rejects_insufficient_storage_for_strided_layout() {
3199        let meta =
3200            TensorMeta::from_shape_and_strides(vec![2, 2], vec![4, 1], 0, DType::F64, Device::Cpu)
3201                .expect("strided meta should validate");
3202        let err = DenseTensor::from_storage(meta, vec![1.0; 5])
3203            .expect_err("strided span must require enough backing storage");
3204        assert!(matches!(
3205            err,
3206            DenseTensorError::InsufficientStorage {
3207                needed: 6,
3208                actual: 5
3209            }
3210        ));
3211    }
3212
3213    #[test]
3214    fn shape_builds_contiguous_strides() {
3215        let meta = TensorMeta::from_shape(vec![2, 3, 4], DType::F64, Device::Cpu);
3216        assert_eq!(meta.strides(), &[12, 4, 1]);
3217        assert_eq!(meta.numel(), 24);
3218        assert!(meta.is_contiguous());
3219    }
3220
3221    #[test]
3222    fn singleton_dim_stride_variation_is_still_contiguous() {
3223        let broadcast_meta =
3224            TensorMeta::from_shape_and_strides(vec![1, 3], vec![0, 1], 0, DType::F64, Device::Cpu)
3225                .expect("broadcast shape should validate");
3226        assert!(
3227            broadcast_meta.is_contiguous(),
3228            "singleton stride should not break contiguous semantics"
3229        );
3230
3231        let interior_singleton_meta = TensorMeta::from_shape_and_strides(
3232            vec![2, 1, 4],
3233            vec![4, 99, 1],
3234            0,
3235            DType::F64,
3236            Device::Cpu,
3237        )
3238        .expect("interior singleton stride should validate");
3239        assert!(
3240            interior_singleton_meta.is_contiguous(),
3241            "interior singleton stride should be ignored for contiguity"
3242        );
3243    }
3244
3245    #[test]
3246    fn non_singleton_stride_mismatch_is_not_contiguous() {
3247        let meta =
3248            TensorMeta::from_shape_and_strides(vec![2, 3], vec![0, 1], 0, DType::F64, Device::Cpu)
3249                .expect("meta should validate");
3250        assert!(
3251            !meta.is_contiguous(),
3252            "non-singleton zero stride should fail contiguous check"
3253        );
3254    }
3255
3256    #[test]
3257    fn custom_strides_validate_and_index_into_storage() {
3258        let meta =
3259            TensorMeta::from_shape_and_strides(vec![2, 2], vec![4, 1], 3, DType::F64, Device::Cpu)
3260                .expect("meta should validate");
3261
3262        assert_eq!(meta.storage_index_for(&[0, 0]).expect("index 0,0"), 3);
3263        assert_eq!(meta.storage_index_for(&[1, 1]).expect("index 1,1"), 8);
3264    }
3265
3266    #[test]
3267    fn index_rank_and_bounds_are_guarded() {
3268        let meta = TensorMeta::from_shape(vec![2, 3], DType::F64, Device::Cpu);
3269
3270        let rank_err = meta
3271            .storage_index_for(&[1])
3272            .expect_err("rank mismatch should fail");
3273        assert!(matches!(
3274            rank_err,
3275            TensorMetaError::IndexRankMismatch {
3276                expected: 2,
3277                actual: 1
3278            }
3279        ));
3280
3281        let oob_err = meta
3282            .storage_index_for(&[2, 0])
3283            .expect_err("out-of-bounds index should fail");
3284        assert!(matches!(
3285            oob_err,
3286            TensorMetaError::IndexOutOfBounds {
3287                dim: 0,
3288                index: 2,
3289                size: 2
3290            }
3291        ));
3292    }
3293
3294    #[test]
3295    fn validate_rejects_stride_overflow() {
3296        let err = TensorMeta::from_shape_and_strides(
3297            vec![3],
3298            vec![usize::MAX],
3299            0,
3300            DType::F64,
3301            Device::Cpu,
3302        )
3303        .expect_err("overflowing stride span must fail validation");
3304
3305        assert!(matches!(
3306            err,
3307            TensorMetaError::StrideOverflow { size: 3, stride } if stride == usize::MAX
3308        ));
3309    }
3310
3311    #[test]
3312    fn storage_index_for_rejects_storage_offset_overflow() {
3313        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu)
3314            .with_storage_offset(usize::MAX);
3315        let err = meta
3316            .storage_index_for(&[1])
3317            .expect_err("overflowing storage offset accumulation must fail");
3318
3319        assert!(matches!(
3320            err,
3321            TensorMetaError::StorageOffsetOverflow {
3322                storage_offset,
3323                max_linear_offset: 1
3324            } if storage_offset == usize::MAX
3325        ));
3326    }
3327
3328    #[test]
3329    fn compatibility_checks_dtype_and_device() {
3330        let lhs = ScalarTensor::new(1.0, DType::F64, Device::Cpu);
3331        let rhs = ScalarTensor::new(2.0, DType::F64, Device::Cpu);
3332        assert!(ensure_compatible(&lhs, &rhs).is_ok());
3333    }
3334
3335    #[test]
3336    fn compatibility_checks_reject_dtype_mismatch() {
3337        let lhs = ScalarTensor::new(1.0, DType::F64, Device::Cpu);
3338        let rhs = ScalarTensor::new(2.0, DType::F32, Device::Cpu);
3339        let err = ensure_compatible(&lhs, &rhs).expect_err("dtype mismatch must fail");
3340        assert!(matches!(
3341            err,
3342            super::TensorCompatError::DTypeMismatch {
3343                lhs: DType::F64,
3344                rhs: DType::F32
3345            }
3346        ));
3347    }
3348
3349    #[test]
3350    fn compatibility_checks_reject_device_mismatch() {
3351        let lhs = ScalarTensor::new(1.0, DType::F64, Device::Cpu);
3352        let rhs = ScalarTensor::new(2.0, DType::F64, Device::Cuda);
3353        let err = ensure_compatible(&lhs, &rhs).expect_err("device mismatch must fail");
3354        assert!(matches!(
3355            err,
3356            super::TensorCompatError::DeviceMismatch {
3357                lhs: Device::Cpu,
3358                rhs: Device::Cuda
3359            }
3360        ));
3361    }
3362
3363    #[test]
3364    fn contiguous_stride_helper_handles_scalar() {
3365        assert_eq!(contiguous_strides(&[]), Vec::<usize>::new());
3366    }
3367
3368    #[test]
3369    fn numel_saturates_on_overflow() {
3370        let meta = TensorMeta::from_shape(vec![usize::MAX, 2], DType::F64, Device::Cpu);
3371        assert_eq!(meta.numel(), usize::MAX);
3372    }
3373
3374    #[test]
3375    fn numel_zero_dimension_short_circuits_before_overflow() {
3376        let meta = TensorMeta::from_shape(vec![usize::MAX, 2, 0], DType::F64, Device::Cpu);
3377        assert_eq!(meta.numel(), 0);
3378    }
3379
3380    #[test]
3381    fn tensor_meta_numel_golden_summary_matches_fixture() {
3382        use std::fmt::Write as _;
3383
3384        let scalar = TensorMeta::scalar(DType::F64, Device::Cpu);
3385        let normal_shape = vec![2, 3, 5, 7, 11, 13, 17, 19];
3386        let normal = TensorMeta::from_shape(normal_shape.clone(), DType::F64, Device::Cpu);
3387        let zero_shape = vec![8, 0, usize::MAX];
3388        let zero = TensorMeta::from_shape(zero_shape.clone(), DType::F64, Device::Cpu);
3389        let overflow_shape = vec![usize::MAX, 2];
3390        let overflow = TensorMeta::from_shape(overflow_shape.clone(), DType::F64, Device::Cpu);
3391
3392        let mut summary = String::new();
3393        summary.push_str("ft_core_numel_pass22\n");
3394        let _ = writeln!(&mut summary, "scalar={}", scalar.numel());
3395        let _ = writeln!(&mut summary, "normal_shape={normal_shape:?}");
3396        let _ = writeln!(&mut summary, "normal_numel={}", normal.numel());
3397        let _ = writeln!(&mut summary, "normal_strides={:?}", normal.strides());
3398        let _ = writeln!(
3399            &mut summary,
3400            "normal_fingerprint={:016x}",
3401            normal.fingerprint64()
3402        );
3403        let _ = writeln!(&mut summary, "zero_shape={zero_shape:?}");
3404        let _ = writeln!(&mut summary, "zero_numel={}", zero.numel());
3405        let _ = writeln!(&mut summary, "overflow_shape={overflow_shape:?}");
3406        let _ = writeln!(&mut summary, "overflow_numel={}", overflow.numel());
3407
3408        assert_eq!(
3409            summary,
3410            include_str!("../../../artifacts/optimization/golden_outputs/ft_core_numel_pass22.txt")
3411        );
3412    }
3413
3414    #[test]
3415    fn dense_view_rejects_shape_overflow() {
3416        let tensor = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
3417        let err = tensor
3418            .view(vec![usize::MAX, 2])
3419            .expect_err("overflowing view shape must fail");
3420        assert!(matches!(err, DenseTensorError::ShapeOverflow { .. }));
3421    }
3422
3423    #[test]
3424    fn dense_view_preserves_storage_offset() {
3425        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu).with_storage_offset(1);
3426        let tensor = DenseTensor::from_storage(meta, vec![99.0, 1.0, 2.0]).unwrap();
3427
3428        let view = tensor.view(vec![1, 2]).unwrap();
3429
3430        assert_eq!(view.storage_id(), tensor.storage_id());
3431        assert_eq!(view.meta().storage_offset(), 1);
3432        assert_eq!(view.meta().shape(), &[1, 2]);
3433        assert_eq!(view.contiguous_values().unwrap(), &[1.0, 2.0]);
3434    }
3435
3436    #[test]
3437    fn dense_matrix_view_preserves_storage_offset() {
3438        let meta =
3439            TensorMeta::from_shape(vec![2, 2], DType::F64, Device::Cpu).with_storage_offset(2);
3440        let tensor = DenseTensor::from_storage(meta, vec![99.0, 98.0, 1.0, 2.0, 3.0, 4.0]).unwrap();
3441
3442        let view = tensor.view(vec![4, 1]).unwrap();
3443
3444        assert_eq!(view.storage_id(), tensor.storage_id());
3445        assert_eq!(view.meta().storage_offset(), 2);
3446        assert_eq!(view.meta().shape(), &[4, 1]);
3447        assert_eq!(view.contiguous_values().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
3448    }
3449
3450    #[test]
3451    fn dense_view_preserves_version_counter() {
3452        let mut tensor =
3453            DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0, 4.0], vec![4], Device::Cpu)
3454                .unwrap();
3455        tensor
3456            .update_contiguous_values(&[5.0, 6.0, 7.0, 8.0])
3457            .unwrap();
3458
3459        let view = tensor.view(vec![2, 2]).unwrap();
3460
3461        assert_ne!(view.id(), tensor.id());
3462        assert_eq!(view.storage_id(), tensor.storage_id());
3463        assert_eq!(view.version(), tensor.version());
3464        assert_eq!(view.version(), 1);
3465        assert_eq!(view.contiguous_values().unwrap(), &[5.0, 6.0, 7.0, 8.0]);
3466    }
3467
3468    #[test]
3469    fn dense_scalar_view_preserves_storage_offset() {
3470        let meta = TensorMeta::from_shape(vec![1], DType::F64, Device::Cpu).with_storage_offset(1);
3471        let tensor = DenseTensor::from_storage(meta, vec![99.0, 7.0]).unwrap();
3472
3473        let view = tensor.view(vec![]).unwrap();
3474
3475        assert_eq!(view.storage_id(), tensor.storage_id());
3476        assert_eq!(view.meta().storage_offset(), 1);
3477        assert_eq!(view.meta().shape(), &[] as &[usize]);
3478        assert_eq!(view.contiguous_values().unwrap(), &[7.0]);
3479    }
3480
3481    #[test]
3482    fn dense_view_preserves_quantization_metadata() {
3483        let tensor = DenseTensor::from_contiguous_values_qint8(
3484            vec![0, 2, 4, 6],
3485            vec![4],
3486            Device::Cpu,
3487            0.5,
3488            -2,
3489        )
3490        .unwrap();
3491
3492        let view = tensor.view(vec![2, 2]).unwrap();
3493
3494        assert_eq!(view.storage_id(), tensor.storage_id());
3495        assert_eq!(view.meta().shape(), &[2, 2]);
3496        assert!(view.meta().quantization().is_some());
3497        assert_eq!(view.contiguous_values_qint8().unwrap(), &[0, 2, 4, 6]);
3498        assert_eq!(
3499            view.dequantized_values_as_f64().unwrap(),
3500            &[1.0, 2.0, 3.0, 4.0]
3501        );
3502    }
3503
3504    #[test]
3505    fn dense_view_preserves_quint8_quantization_metadata() {
3506        let tensor = DenseTensor::from_contiguous_values_quint8(
3507            vec![10, 14, 18, 22],
3508            vec![4],
3509            Device::Cpu,
3510            0.25,
3511            10,
3512        )
3513        .unwrap();
3514
3515        let view = tensor.view(vec![2, 2]).unwrap();
3516
3517        assert_eq!(view.storage_id(), tensor.storage_id());
3518        assert_eq!(view.meta().shape(), &[2, 2]);
3519        assert!(view.meta().quantization().is_some());
3520        assert_eq!(view.contiguous_values_quint8().unwrap(), &[10, 14, 18, 22]);
3521        assert_eq!(
3522            view.dequantized_values_as_f64().unwrap(),
3523            &[0.0, 1.0, 2.0, 3.0]
3524        );
3525    }
3526
3527    #[test]
3528    fn dense_quantized_offset_view_preserves_metadata() {
3529        let meta = TensorMeta::quantized_from_shape_and_strides(
3530            vec![2],
3531            vec![1],
3532            1,
3533            DType::QInt8,
3534            Device::Cpu,
3535            0.5,
3536            -2,
3537        )
3538        .unwrap();
3539        let tensor = DenseTensor::from_storage_qint8(meta, vec![99, 0, 2]).unwrap();
3540
3541        let view = tensor.view(vec![1, 2]).unwrap();
3542
3543        assert_eq!(view.storage_id(), tensor.storage_id());
3544        assert_eq!(view.meta().storage_offset(), 1);
3545        assert!(view.meta().quantization().is_some());
3546        assert_eq!(view.contiguous_values_qint8().unwrap(), &[0, 2]);
3547        assert_eq!(view.dequantized_values_as_f64().unwrap(), &[1.0, 2.0]);
3548    }
3549
3550    #[test]
3551    fn sparse_coo_to_dense_rejects_shape_overflow() {
3552        let coords = vec![vec![0, 0]];
3553        let sparse = SparseCOOTensor::from_coords(
3554            &coords,
3555            vec![1.0],
3556            vec![usize::MAX, 2],
3557            DType::F64,
3558            Device::Cpu,
3559        )
3560        .expect("sparse COO should build with oversized shape");
3561        let err = sparse
3562            .to_dense()
3563            .expect_err("overflowing dense shape must fail");
3564        assert!(matches!(
3565            err,
3566            SparseTensorError::DenseTensor(DenseTensorError::ShapeOverflow { .. })
3567        ));
3568    }
3569
3570    #[test]
3571    fn out_of_place_result_gets_new_storage_and_version_bump() {
3572        let source = ScalarTensor::new(2.0, DType::F64, Device::Cpu);
3573        let derived = source.with_value(5.0);
3574
3575        assert_ne!(source.id(), derived.id());
3576        assert_ne!(source.storage_id(), derived.storage_id());
3577        assert_eq!(derived.version(), source.version() + 1);
3578    }
3579
3580    #[test]
3581    fn alias_view_shares_storage_identity() {
3582        let source = ScalarTensor::new(2.0, DType::F64, Device::Cpu);
3583        let alias = source.alias_view(0).expect("alias with zero offset");
3584
3585        assert_ne!(source.id(), alias.id());
3586        assert_eq!(source.storage_id(), alias.storage_id());
3587        assert_eq!(source.version(), alias.version());
3588        assert_eq!(source.value(), alias.value());
3589    }
3590
3591    #[test]
3592    fn in_place_updates_bump_version_and_fingerprint() {
3593        let mut tensor = ScalarTensor::new(2.0, DType::F64, Device::Cpu);
3594        let before = tensor.evidence_fingerprint64();
3595        tensor.set_in_place(7.0);
3596        let after = tensor.evidence_fingerprint64();
3597
3598        assert_eq!(tensor.value(), 7.0);
3599        assert_eq!(tensor.version(), 1);
3600        assert_ne!(before, after);
3601    }
3602
3603    #[test]
3604    fn meta_fingerprint_changes_when_offset_changes() {
3605        let a = TensorMeta::from_shape(vec![2, 2], DType::F64, Device::Cpu);
3606        let b = a.clone().with_storage_offset(1);
3607        assert_ne!(a.fingerprint64(), b.fingerprint64());
3608    }
3609
3610    proptest! {
3611        #[test]
3612        fn prop_contiguous_stride_contract(shape in prop::collection::vec(1usize..=4, 1..=4)) {
3613            let strides = contiguous_strides(shape.as_slice());
3614            prop_assert_eq!(strides.len(), shape.len());
3615            prop_assert_eq!(strides.last().copied(), Some(1));
3616
3617            let seed = det_seed(shape.as_slice());
3618            let log = build_property_log(
3619                "prop_contiguous_stride_contract",
3620                "strict",
3621                seed,
3622                seed,
3623                det_seed(strides.as_slice()),
3624                "contiguous_stride_contract_ok",
3625            );
3626            assert_log_contract(&log);
3627        }
3628
3629        #[test]
3630        fn prop_numel_matches_shape_product(shape in prop::collection::vec(1usize..=6, 1..=4)) {
3631            let meta = TensorMeta::from_shape(shape.clone(), DType::F64, Device::Cpu);
3632            let expected: usize = shape.iter().copied().product();
3633            prop_assert_eq!(meta.numel(), expected);
3634
3635            let seed = det_seed(shape.as_slice());
3636            let log = build_property_log(
3637                "prop_numel_matches_shape_product",
3638                "strict",
3639                seed,
3640                seed,
3641                expected as u64,
3642                "numel_product_contract_ok",
3643            );
3644            assert_log_contract(&log);
3645        }
3646
3647        #[test]
3648        fn prop_contiguous_index_bounds(shape in prop::collection::vec(1usize..=4, 1..=4)) {
3649            let meta = TensorMeta::from_shape(shape.clone(), DType::F64, Device::Cpu);
3650            let zero_index = vec![0; shape.len()];
3651            let max_index = shape.iter().map(|dim| dim - 1).collect::<Vec<_>>();
3652
3653            let zero_linear = meta.storage_index_for(zero_index.as_slice()).expect("zero index must be valid");
3654            let max_linear = meta.storage_index_for(max_index.as_slice()).expect("max index must be valid");
3655
3656            prop_assert_eq!(zero_linear, 0);
3657            prop_assert!(max_linear < meta.numel());
3658
3659            let seed = det_seed(shape.as_slice());
3660            let log = build_property_log(
3661                "prop_contiguous_index_bounds",
3662                "strict",
3663                seed,
3664                seed,
3665                max_linear as u64,
3666                "index_bounds_contract_ok",
3667            );
3668            assert_log_contract(&log);
3669        }
3670
3671        #[test]
3672        fn prop_with_value_bumps_version_and_storage(
3673            source_value in -1_000.0f64..1_000.0f64,
3674            derived_value in -1_000.0f64..1_000.0f64,
3675        ) {
3676            let source = ScalarTensor::new(source_value, DType::F64, Device::Cpu);
3677            let derived = source.with_value(derived_value);
3678
3679            prop_assert_eq!(derived.version(), source.version() + 1);
3680            prop_assert_ne!(derived.storage_id(), source.storage_id());
3681
3682            let source_seed = source_value.to_bits() as usize;
3683            let derived_seed = derived_value.to_bits() as usize;
3684            let seed = det_seed([source_seed, derived_seed].as_slice());
3685            let log = build_property_log(
3686                "prop_with_value_bumps_version_and_storage",
3687                "strict",
3688                seed,
3689                source.evidence_fingerprint64(),
3690                derived.evidence_fingerprint64(),
3691                "version_storage_contract_ok",
3692            );
3693            assert_log_contract(&log);
3694        }
3695
3696        #[test]
3697        fn prop_rank_stride_mismatch_fail_closed(
3698            shape in prop::collection::vec(1usize..=4, 1..=4),
3699            extra in 1usize..=3,
3700        ) {
3701            let strides = vec![1usize; shape.len() + extra];
3702            let err = TensorMeta::from_shape_and_strides(
3703                shape.clone(),
3704                strides,
3705                0,
3706                DType::F64,
3707                Device::Cpu,
3708            )
3709            .expect_err("rank/stride mismatch must fail");
3710
3711            match err {
3712                TensorMetaError::RankStrideMismatch { .. } => {}
3713                other => prop_assert!(false, "expected RankStrideMismatch, got {other:?}"),
3714            }
3715
3716            let seed = det_seed(shape.as_slice());
3717            let log = build_property_log(
3718                "prop_rank_stride_mismatch_fail_closed",
3719                "strict",
3720                seed,
3721                seed,
3722                0,
3723                "rank_stride_mismatch_fail_closed",
3724            );
3725            assert_log_contract(&log);
3726        }
3727    }
3728
3729    // ── bd-2wwr: DenseTensor accessors and methods ──
3730
3731    #[test]
3732    fn dense_tensor_id_storage_id_version_accessors() {
3733        let dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0], vec![3], Device::Cpu)
3734            .expect("create dense tensor");
3735        // id and storage_id should be nonzero (unique)
3736        assert!(dt.id() > 0);
3737        assert!(dt.storage_id() > 0);
3738        assert_eq!(dt.version(), 0);
3739    }
3740
3741    #[test]
3742    fn dense_tensor_storage_accessor() {
3743        let vals = vec![10.0, 20.0, 30.0, 40.0];
3744        let dt = DenseTensor::from_contiguous_values(vals.clone(), vec![4], Device::Cpu)
3745            .expect("create dense tensor");
3746        assert_eq!(
3747            dt.storage().expect("f64 storage should be accessible"),
3748            &[10.0, 20.0, 30.0, 40.0]
3749        );
3750    }
3751
3752    #[test]
3753    fn dense_tensor_storage_accessor_rejects_non_f64_dtype() {
3754        let dt = DenseTensor::from_storage_f32(
3755            TensorMeta::from_shape(vec![2], DType::F32, Device::Cpu),
3756            vec![1.0f32, 2.0],
3757        )
3758        .expect("create f32 dense tensor");
3759        let err = dt
3760            .storage()
3761            .expect_err("raw f64 storage access should reject f32 tensors");
3762        assert!(matches!(
3763            err,
3764            DenseTensorError::UnsupportedStorageAccess { dtype: DType::F32 }
3765        ));
3766    }
3767
3768    #[test]
3769    fn dense_tensor_replace_storage_success() {
3770        let mut dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0], vec![3], Device::Cpu)
3771            .expect("create dense tensor");
3772        assert_eq!(dt.version(), 0);
3773
3774        dt.update_contiguous_values(&[4.0, 5.0, 6.0])
3775            .expect("replace with same-length storage should succeed");
3776        assert_eq!(dt.version(), 1);
3777        assert_eq!(
3778            dt.storage().expect("f64 storage should be accessible"),
3779            &[4.0, 5.0, 6.0]
3780        );
3781    }
3782
3783    #[test]
3784    fn dense_tensor_update_contiguous_values_with_mutates_in_place() {
3785        let mut dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0], vec![3], Device::Cpu)
3786            .expect("create dense tensor");
3787        assert_eq!(dt.version(), 0);
3788
3789        dt.update_contiguous_values_with(|values| {
3790            for value in values {
3791                *value *= 2.0;
3792            }
3793        })
3794        .expect("closure update should succeed");
3795
3796        assert_eq!(dt.version(), 1);
3797        assert_eq!(
3798            dt.storage().expect("f64 storage should be accessible"),
3799            &[2.0, 4.0, 6.0]
3800        );
3801    }
3802
3803    #[test]
3804    fn dense_tensor_replace_storage_wrong_length() {
3805        let mut dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0], vec![3], Device::Cpu)
3806            .expect("create dense tensor");
3807        let err = dt
3808            .update_contiguous_values(&[1.0, 2.0])
3809            .expect_err("replace with different-length storage should fail");
3810        assert!(
3811            matches!(
3812                err,
3813                DenseTensorError::InsufficientStorage {
3814                    needed: 3,
3815                    actual: 2
3816                }
3817            ),
3818            "expected InsufficientStorage, got {err:?}"
3819        );
3820        // version should NOT have been bumped
3821        assert_eq!(dt.version(), 0);
3822    }
3823
3824    #[test]
3825    fn dense_tensor_dispatch_values_returns_storage_slice() {
3826        let dt = DenseTensor::from_contiguous_values(vec![5.0, 6.0, 7.0], vec![3], Device::Cpu)
3827            .expect("create dense tensor");
3828        let vals = dt
3829            .dispatch_values()
3830            .expect("dispatch_values should succeed");
3831        assert_eq!(vals, &[5.0, 6.0, 7.0]);
3832    }
3833
3834    #[test]
3835    fn dense_tensor_from_storage_rejects_dtype_mismatch() {
3836        // from_storage takes Vec<f64>, so passing F32 meta is a mismatch
3837        let meta = TensorMeta::from_shape(vec![2], DType::F32, Device::Cpu);
3838        let err = DenseTensor::from_storage(meta, vec![1.0, 2.0])
3839            .expect_err("F32 meta with f64 storage should be rejected");
3840        assert!(
3841            matches!(err, DenseTensorError::UnsupportedDType(DType::F32)),
3842            "expected UnsupportedDType(F32), got {err:?}"
3843        );
3844
3845        // from_storage_f32 with F64 meta is also a mismatch
3846        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu);
3847        let err = DenseTensor::from_storage_f32(meta, vec![1.0f32, 2.0])
3848            .expect_err("F64 meta with f32 storage should be rejected");
3849        assert!(
3850            matches!(err, DenseTensorError::UnsupportedDType(DType::F64)),
3851            "expected UnsupportedDType(F64), got {err:?}"
3852        );
3853    }
3854
3855    #[test]
3856    fn dense_tensor_from_storage_rejects_insufficient_storage() {
3857        let meta = TensorMeta::from_shape(vec![5], DType::F64, Device::Cpu);
3858        let err = DenseTensor::from_storage(meta, vec![1.0, 2.0])
3859            .expect_err("too-short storage should be rejected");
3860        assert!(
3861            matches!(
3862                err,
3863                DenseTensorError::InsufficientStorage {
3864                    needed: 5,
3865                    actual: 2
3866                }
3867            ),
3868            "expected InsufficientStorage, got {err:?}"
3869        );
3870    }
3871
3872    #[test]
3873    fn dense_tensor_replace_storage_bumps_version_multiple() {
3874        let mut dt = DenseTensor::from_contiguous_values(vec![0.0, 0.0], vec![2], Device::Cpu)
3875            .expect("create dense tensor");
3876        for i in 1..=5 {
3877            dt.update_contiguous_values(&[i as f64, i as f64 * 2.0])
3878                .expect("replace should succeed");
3879        }
3880        assert_eq!(dt.version(), 5);
3881    }
3882
3883    // ── Integer DType tests ───────────────────────────────────────────
3884
3885    #[test]
3886    fn dtype_element_sizes() {
3887        assert_eq!(DType::F64.element_size(), 8);
3888        assert_eq!(DType::F32.element_size(), 4);
3889        assert_eq!(DType::QInt8.element_size(), 1);
3890        assert_eq!(DType::QUInt8.element_size(), 1);
3891        assert_eq!(DType::I64.element_size(), 8);
3892        assert_eq!(DType::I32.element_size(), 4);
3893    }
3894
3895    #[test]
3896    fn dtype_is_floating_point() {
3897        assert!(DType::F64.is_floating_point());
3898        assert!(DType::F32.is_floating_point());
3899        assert!(!DType::QInt8.is_floating_point());
3900        assert!(!DType::QUInt8.is_floating_point());
3901        assert!(!DType::I64.is_floating_point());
3902        assert!(!DType::I32.is_floating_point());
3903    }
3904
3905    #[test]
3906    fn dtype_is_integer() {
3907        assert!(!DType::F64.is_integer());
3908        assert!(!DType::F32.is_integer());
3909        assert!(!DType::QInt8.is_integer());
3910        assert!(!DType::QUInt8.is_integer());
3911        assert!(DType::I64.is_integer());
3912        assert!(DType::I32.is_integer());
3913    }
3914
3915    #[test]
3916    fn dtype_is_quantized() {
3917        assert!(DType::QInt8.is_quantized());
3918        assert!(DType::QUInt8.is_quantized());
3919        assert!(!DType::F32.is_quantized());
3920        assert!(!DType::I64.is_quantized());
3921    }
3922
3923    #[test]
3924    fn tensor_meta_requires_quantization_params_for_quantized_dtype() {
3925        let missing = TensorMeta::from_shape(vec![2], DType::QInt8, Device::Cpu)
3926            .validate()
3927            .expect_err("qint8 metadata requires qparams");
3928        assert!(matches!(
3929            missing,
3930            TensorMetaError::MissingQuantizationParams {
3931                dtype: DType::QInt8
3932            }
3933        ));
3934
3935        let unexpected = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu)
3936            .with_quantization(QuantizationParams::new(0.25, 3).expect("qparams"))
3937            .validate()
3938            .expect_err("f64 metadata cannot carry qparams");
3939        assert!(matches!(
3940            unexpected,
3941            TensorMetaError::UnexpectedQuantizationParams { dtype: DType::F64 }
3942        ));
3943
3944        let invalid_scale =
3945            QuantizationParams::new(0.0, 0).expect_err("zero quantization scale must be rejected");
3946        assert!(matches!(
3947            invalid_scale,
3948            TensorMetaError::InvalidQuantizationScale { .. }
3949        ));
3950    }
3951
3952    #[test]
3953    fn dense_tensor_accepts_qint8_storage_with_quantized_meta() {
3954        let tensor = DenseTensor::from_contiguous_values_qint8(
3955            vec![-10, 0, 10],
3956            vec![3],
3957            Device::Cpu,
3958            0.5,
3959            0,
3960        )
3961        .expect("qint8 dense tensor");
3962
3963        assert_eq!(tensor.meta().dtype(), DType::QInt8);
3964        let qparams = tensor
3965            .meta()
3966            .quantization()
3967            .expect("quantization params should be present");
3968        assert_eq!(qparams.scale(), 0.5);
3969        assert_eq!(qparams.zero_point(), 0);
3970        assert_eq!(
3971            tensor.contiguous_values_qint8().expect("raw qint8"),
3972            &[-10, 0, 10]
3973        );
3974        assert_eq!(
3975            tensor
3976                .dequantized_values_as_f64()
3977                .expect("dequantized qint8"),
3978            vec![-5.0, 0.0, 5.0]
3979        );
3980    }
3981
3982    #[test]
3983    fn dense_tensor_dequantizes_qint8_per_channel_by_axis() {
3984        let tensor = DenseTensor::from_contiguous_values_qint8_per_channel(
3985            vec![1, 2, 3, 4, 5, 6],
3986            vec![2, 3],
3987            Device::Cpu,
3988            vec![0.5, 0.25],
3989            vec![1, 2],
3990            0,
3991        )
3992        .expect("per-channel qint8 dense tensor");
3993
3994        let qparams = tensor
3995            .meta()
3996            .quantization()
3997            .expect("per-channel quantization params");
3998        assert_eq!(qparams.axis(), Some(0));
3999        assert_eq!(qparams.scales(), vec![0.5, 0.25]);
4000        assert_eq!(qparams.zero_points(), &[1, 2]);
4001        assert_eq!(
4002            tensor
4003                .dequantized_values_as_f64()
4004                .expect("dequantized per-channel qint8"),
4005            vec![0.0, 0.5, 1.0, 0.5, 0.75, 1.0]
4006        );
4007    }
4008
4009    #[test]
4010    fn dense_tensor_dequantizes_per_channel_mid_axis_rank3() {
4011        // axis=1 on a rank-3 [2,2,2] tensor: `inner` = product(shape[2..]) = 2 spans
4012        // the trailing dim, exercising the hoisted multi-dim inner-product path.
4013        // channel index = (flat_idx / 2) % 2.
4014        let tensor = DenseTensor::from_contiguous_values_quint8_per_channel(
4015            vec![2, 4, 6, 8, 10, 12, 14, 16],
4016            vec![2, 2, 2],
4017            Device::Cpu,
4018            vec![0.5, 0.25],
4019            vec![0, 0],
4020            1,
4021        )
4022        .expect("per-channel quint8 rank-3 dense tensor");
4023
4024        assert_eq!(
4025            tensor
4026                .dequantized_values_as_f64()
4027                .expect("dequantized per-channel quint8 rank-3"),
4028            vec![1.0, 2.0, 1.5, 2.0, 5.0, 6.0, 3.5, 4.0]
4029        );
4030    }
4031
4032    #[test]
4033    fn tensor_meta_rejects_per_channel_qparams_with_wrong_axis_size() {
4034        let err = TensorMeta::quantized_per_channel_from_shape(
4035            vec![2, 3],
4036            DType::QInt8,
4037            Device::Cpu,
4038            vec![0.5, 0.25],
4039            vec![1, 2],
4040            1,
4041        )
4042        .expect_err("axis=1 requires three qparam channels");
4043
4044        assert!(matches!(
4045            err,
4046            TensorMetaError::QuantizationChannelCountMismatch {
4047                axis: 1,
4048                expected: 3,
4049                actual: 2
4050            }
4051        ));
4052    }
4053
4054    #[test]
4055    fn dense_tensor_accepts_quint8_storage_with_quantized_meta() {
4056        let tensor = DenseTensor::from_contiguous_values_quint8(
4057            vec![0, 128, 255],
4058            vec![3],
4059            Device::Cpu,
4060            0.25,
4061            128,
4062        )
4063        .expect("quint8 dense tensor");
4064
4065        assert_eq!(tensor.meta().dtype(), DType::QUInt8);
4066        assert_eq!(
4067            tensor.contiguous_values_quint8().expect("raw quint8"),
4068            &[0, 128, 255]
4069        );
4070        assert_eq!(
4071            tensor
4072                .dequantized_values_as_f64()
4073                .expect("dequantized quint8"),
4074            vec![-32.0, 0.0, 31.75]
4075        );
4076    }
4077
4078    proptest! {
4079        // Metamorphic/property coverage for the freshly-churned sparse COO path
4080        // (frankentorch-b2yob). (1) to_dense scatter-adds and SUMS duplicate
4081        // coordinates bit-exactly (torch non-coalesced parity); (2) coalesced
4082        // construction accepts exactly the strictly lex-sorted-unique inputs.
4083        #[test]
4084        fn prop_coo_to_dense_scatter_adds_and_sums_duplicates(
4085            rows in 1usize..=4,
4086            cols in 1usize..=4,
4087            raw in prop::collection::vec((0usize..16, 0usize..16, -100.0f64..100.0f64), 1..=12),
4088        ) {
4089            let nnz = raw.len();
4090            let mut row_idx: Vec<i64> = Vec::with_capacity(nnz);
4091            let mut col_idx: Vec<i64> = Vec::with_capacity(nnz);
4092            let mut vals: Vec<f64> = Vec::with_capacity(nnz);
4093            let mut expected = vec![0.0f64; rows * cols];
4094            for (r, c, v) in &raw {
4095                let rr = r % rows;
4096                let cc = c % cols;
4097                row_idx.push(rr as i64);
4098                col_idx.push(cc as i64);
4099                vals.push(*v);
4100                // Independent scatter-add in the SAME nnz order to_dense uses, so
4101                // the float summation is bit-for-bit comparable.
4102                expected[rr * cols + cc] += *v;
4103            }
4104            let mut flat = row_idx;
4105            flat.extend(col_idx.iter().copied());
4106            let indices = DenseI64Tensor::from_contiguous_values(flat, vec![2, nnz], Device::Cpu)
4107                .expect("indices");
4108            let values = DenseTensor::from_contiguous_values(vals, vec![nnz], Device::Cpu)
4109                .expect("values");
4110            // coalesced=false: duplicate coordinates are allowed and must be summed.
4111            let coo = SparseCOOTensor::new(indices, values, vec![rows, cols], false)
4112                .expect("coo construct");
4113            let dense = coo.to_dense().expect("to_dense");
4114            let got = dense.contiguous_values_as_f64().expect("dense values");
4115            for k in 0..rows * cols {
4116                prop_assert_eq!(got[k].to_bits(), expected[k].to_bits());
4117            }
4118            let seed = det_seed(&[rows, cols, nnz]);
4119            let log = build_property_log(
4120                "prop_coo_to_dense_scatter_adds_and_sums_duplicates",
4121                "strict",
4122                seed,
4123                seed,
4124                det_seed(&[got.len()]),
4125                "coo_to_dense_scatter_sum_ok",
4126            );
4127            assert_log_contract(&log);
4128        }
4129
4130        #[test]
4131        fn prop_coalesced_coo_accepts_iff_strictly_sorted_unique(
4132            rows in 1usize..=4,
4133            cols in 1usize..=4,
4134            raw in prop::collection::vec((0usize..8, 0usize..8), 1..=8),
4135        ) {
4136            let nnz = raw.len();
4137            let coords: Vec<(i64, i64)> = raw
4138                .iter()
4139                .map(|(r, c)| ((r % rows) as i64, (c % cols) as i64))
4140                .collect();
4141            // Tuple Ord is lexicographic on (row, col) — exactly the dim-0-then-dim-1
4142            // order the coalesced validator enforces.
4143            let strictly_sorted_unique = (1..nnz).all(|i| coords[i - 1] < coords[i]);
4144            let mut flat: Vec<i64> = coords.iter().map(|(r, _)| *r).collect();
4145            flat.extend(coords.iter().map(|(_, c)| *c));
4146            let indices = DenseI64Tensor::from_contiguous_values(flat, vec![2, nnz], Device::Cpu)
4147                .expect("indices");
4148            let values =
4149                DenseTensor::from_contiguous_values(vec![1.0f64; nnz], vec![nnz], Device::Cpu)
4150                    .expect("values");
4151            let result = SparseCOOTensor::new(indices, values, vec![rows, cols], true);
4152            // Coordinates are in-bounds and shapes are consistent, so the ONLY
4153            // construction failure reason is the coalesced invariant.
4154            prop_assert_eq!(result.is_ok(), strictly_sorted_unique);
4155            let seed = det_seed(&[rows, cols, nnz]);
4156            let log = build_property_log(
4157                "prop_coalesced_coo_accepts_iff_strictly_sorted_unique",
4158                "strict",
4159                seed,
4160                seed,
4161                det_seed(&[usize::from(strictly_sorted_unique)]),
4162                "coalesced_coo_validation_characterized",
4163            );
4164            assert_log_contract(&log);
4165        }
4166    }
4167
4168    #[test]
4169    fn tensor_meta_with_integer_dtypes() {
4170        let meta_i64 = TensorMeta::from_shape(vec![2, 3], DType::I64, Device::Cpu);
4171        assert_eq!(meta_i64.dtype(), DType::I64);
4172        assert_eq!(meta_i64.numel(), 6);
4173        assert_eq!(meta_i64.strides(), &[3, 1]);
4174        assert!(meta_i64.is_contiguous());
4175
4176        let meta_i32 = TensorMeta::from_shape(vec![4], DType::I32, Device::Cpu);
4177        assert_eq!(meta_i32.dtype(), DType::I32);
4178        assert_eq!(meta_i32.numel(), 4);
4179    }
4180
4181    #[test]
4182    fn dense_i64_tensor_from_contiguous_values() {
4183        let dt = DenseI64Tensor::from_contiguous_values(vec![1, 2, 3, 4], vec![2, 2], Device::Cpu)
4184            .expect("create i64 tensor");
4185        assert_eq!(dt.meta().shape(), &[2, 2]);
4186        assert_eq!(dt.meta().dtype(), DType::I64);
4187        assert_eq!(dt.contiguous_values().expect("values"), &[1i64, 2, 3, 4]);
4188        assert!(dt.id() > 0);
4189        assert_eq!(dt.version(), 0);
4190    }
4191
4192    #[test]
4193    fn dense_i32_tensor_from_contiguous_values() {
4194        let dt = DenseI32Tensor::from_contiguous_values(vec![10, 20, 30], vec![3], Device::Cpu)
4195            .expect("create i32 tensor");
4196        assert_eq!(dt.meta().shape(), &[3]);
4197        assert_eq!(dt.meta().dtype(), DType::I32);
4198        assert_eq!(dt.contiguous_values().expect("values"), &[10i32, 20, 30]);
4199    }
4200
4201    #[test]
4202    fn dense_i64_tensor_rejects_wrong_dtype() {
4203        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu);
4204        let err = DenseI64Tensor::from_storage(meta, vec![1, 2])
4205            .expect_err("wrong dtype should be rejected");
4206        assert!(matches!(
4207            err,
4208            DenseTensorError::UnsupportedDType(DType::F64)
4209        ));
4210    }
4211
4212    #[test]
4213    fn dense_i32_tensor_rejects_wrong_dtype() {
4214        let meta = TensorMeta::from_shape(vec![2], DType::I64, Device::Cpu);
4215        let err = DenseI32Tensor::from_storage(meta, vec![1, 2])
4216            .expect_err("wrong dtype should be rejected");
4217        assert!(matches!(
4218            err,
4219            DenseTensorError::UnsupportedDType(DType::I64)
4220        ));
4221    }
4222
4223    #[test]
4224    fn dense_i64_tensor_rejects_insufficient_storage() {
4225        let meta = TensorMeta::from_shape(vec![5], DType::I64, Device::Cpu);
4226        let err = DenseI64Tensor::from_storage(meta, vec![1, 2])
4227            .expect_err("insufficient storage should fail");
4228        assert!(matches!(
4229            err,
4230            DenseTensorError::InsufficientStorage {
4231                needed: 5,
4232                actual: 2
4233            }
4234        ));
4235    }
4236
4237    #[test]
4238    fn dense_i64_tensor_rejects_storage_span_overflow() {
4239        let meta = TensorMeta::from_shape(vec![1], DType::I64, Device::Cpu)
4240            .with_storage_offset(usize::MAX);
4241        let err =
4242            DenseI64Tensor::from_storage(meta, vec![1]).expect_err("overflowing span must fail");
4243        assert!(matches!(
4244            err,
4245            DenseTensorError::StorageSpanOverflow {
4246                storage_offset,
4247                numel: 1
4248            } if storage_offset == usize::MAX
4249        ));
4250    }
4251
4252    #[test]
4253    fn dense_i64_tensor_negative_values() {
4254        let dt = DenseI64Tensor::from_contiguous_values(
4255            vec![-100, 0, i64::MAX, i64::MIN],
4256            vec![4],
4257            Device::Cpu,
4258        )
4259        .expect("create with extreme values");
4260        let vals = dt.contiguous_values().expect("values");
4261        assert_eq!(vals[0], -100);
4262        assert_eq!(vals[2], i64::MAX);
4263        assert_eq!(vals[3], i64::MIN);
4264    }
4265
4266    #[test]
4267    fn dense_i32_tensor_extreme_values() {
4268        let dt = DenseI32Tensor::from_contiguous_values(
4269            vec![i32::MAX, i32::MIN, 0, -1],
4270            vec![4],
4271            Device::Cpu,
4272        )
4273        .expect("create with extreme values");
4274        let vals = dt.contiguous_values().expect("values");
4275        assert_eq!(vals[0], i32::MAX);
4276        assert_eq!(vals[1], i32::MIN);
4277    }
4278
4279    #[test]
4280    fn dense_i64_tensor_scalar() {
4281        let dt = DenseI64Tensor::from_contiguous_values(vec![42], vec![], Device::Cpu)
4282            .expect("scalar i64 tensor");
4283        assert_eq!(dt.meta().shape(), &[] as &[usize]);
4284        assert_eq!(dt.meta().numel(), 1);
4285        assert_eq!(dt.contiguous_values().expect("values"), &[42]);
4286    }
4287
4288    #[test]
4289    fn dense_i64_tensor_empty() {
4290        let dt = DenseI64Tensor::from_contiguous_values(vec![], vec![0], Device::Cpu)
4291            .expect("empty i64 tensor");
4292        assert_eq!(dt.meta().numel(), 0);
4293        assert_eq!(dt.contiguous_values().expect("values"), &[] as &[i64]);
4294    }
4295
4296    #[test]
4297    fn dense_i64_tensor_storage_accessor() {
4298        let dt = DenseI64Tensor::from_contiguous_values(vec![5, 6, 7], vec![3], Device::Cpu)
4299            .expect("create i64 tensor");
4300        assert_eq!(dt.storage(), &[5i64, 6, 7]);
4301    }
4302
4303    #[test]
4304    fn dense_i32_tensor_storage_accessor() {
4305        let dt = DenseI32Tensor::from_contiguous_values(vec![8, 9], vec![2], Device::Cpu)
4306            .expect("create i32 tensor");
4307        assert_eq!(dt.storage(), &[8i32, 9]);
4308    }
4309
4310    #[test]
4311    fn dense_i32_tensor_rejects_storage_span_overflow() {
4312        let meta = TensorMeta::from_shape(vec![1], DType::I32, Device::Cpu)
4313            .with_storage_offset(usize::MAX);
4314        let err =
4315            DenseI32Tensor::from_storage(meta, vec![1]).expect_err("overflowing span must fail");
4316        assert!(matches!(
4317            err,
4318            DenseTensorError::StorageSpanOverflow {
4319                storage_offset,
4320                numel: 1
4321            } if storage_offset == usize::MAX
4322        ));
4323    }
4324
4325    // ---- Bool DType tests (bd-2do9.2) ----
4326
4327    #[test]
4328    fn bool_dtype_element_size_is_1() {
4329        assert_eq!(DType::Bool.element_size(), 1);
4330    }
4331
4332    #[test]
4333    fn bool_dtype_is_not_floating_point() {
4334        assert!(!DType::Bool.is_floating_point());
4335    }
4336
4337    #[test]
4338    fn bool_dtype_is_not_integer() {
4339        assert!(!DType::Bool.is_integer());
4340    }
4341
4342    #[test]
4343    fn bool_dtype_is_bool() {
4344        assert!(DType::Bool.is_bool());
4345        assert!(!DType::F64.is_bool());
4346        assert!(!DType::F32.is_bool());
4347        assert!(!DType::I64.is_bool());
4348        assert!(!DType::I32.is_bool());
4349    }
4350
4351    #[test]
4352    fn dense_bool_tensor_from_bools() {
4353        let t = DenseBoolTensor::from_bools(&[true, false, true, false], vec![4], Device::Cpu)
4354            .expect("create bool tensor");
4355        assert_eq!(t.meta().dtype(), DType::Bool);
4356        assert_eq!(t.meta().shape(), &[4]);
4357        assert_eq!(t.contiguous_values().unwrap(), &[1u8, 0, 1, 0]);
4358        assert_eq!(
4359            t.contiguous_bools().unwrap(),
4360            vec![true, false, true, false]
4361        );
4362    }
4363
4364    #[test]
4365    fn dense_bool_tensor_2d() {
4366        let t = DenseBoolTensor::from_bools(
4367            &[true, true, false, false, true, false],
4368            vec![2, 3],
4369            Device::Cpu,
4370        )
4371        .expect("create 2d bool tensor");
4372        assert_eq!(t.meta().shape(), &[2, 3]);
4373        assert_eq!(t.meta().numel(), 6);
4374        assert_eq!(t.contiguous_values().unwrap(), &[1u8, 1, 0, 0, 1, 0]);
4375    }
4376
4377    #[test]
4378    fn dense_bool_tensor_rejects_wrong_dtype() {
4379        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu);
4380        let err =
4381            DenseBoolTensor::from_storage(meta, vec![0, 1]).expect_err("wrong dtype should fail");
4382        assert!(matches!(
4383            err,
4384            DenseTensorError::UnsupportedDType(DType::F64)
4385        ));
4386    }
4387
4388    #[test]
4389    fn dense_bool_tensor_rejects_insufficient_storage() {
4390        let meta = TensorMeta::from_shape(vec![5], DType::Bool, Device::Cpu);
4391        let err = DenseBoolTensor::from_storage(meta, vec![0, 1])
4392            .expect_err("insufficient storage should fail");
4393        assert!(matches!(
4394            err,
4395            DenseTensorError::InsufficientStorage {
4396                needed: 5,
4397                actual: 2
4398            }
4399        ));
4400    }
4401
4402    #[test]
4403    fn dense_bool_tensor_rejects_storage_span_overflow() {
4404        let meta = TensorMeta::from_shape(vec![1], DType::Bool, Device::Cpu)
4405            .with_storage_offset(usize::MAX);
4406        let err =
4407            DenseBoolTensor::from_storage(meta, vec![1]).expect_err("overflowing span must fail");
4408        assert!(matches!(
4409            err,
4410            DenseTensorError::StorageSpanOverflow {
4411                storage_offset,
4412                numel: 1
4413            } if storage_offset == usize::MAX
4414        ));
4415    }
4416
4417    #[test]
4418    fn dense_i64_tensor_contiguous_values_rejects_storage_span_overflow() {
4419        let tensor = DenseI64Tensor {
4420            id: 1,
4421            storage_id: 1,
4422            meta: TensorMeta::from_shape(vec![1], DType::I64, Device::Cpu)
4423                .with_storage_offset(usize::MAX),
4424            storage: vec![1],
4425            version: 0,
4426        };
4427
4428        let err = tensor
4429            .contiguous_values()
4430            .expect_err("overflowing span must fail");
4431        assert!(matches!(
4432            err,
4433            DenseTensorError::StorageSpanOverflow {
4434                storage_offset,
4435                numel: 1
4436            } if storage_offset == usize::MAX
4437        ));
4438    }
4439
4440    #[test]
4441    fn dense_i32_tensor_contiguous_values_rejects_storage_span_overflow() {
4442        let tensor = DenseI32Tensor {
4443            id: 1,
4444            storage_id: 1,
4445            meta: TensorMeta::from_shape(vec![1], DType::I32, Device::Cpu)
4446                .with_storage_offset(usize::MAX),
4447            storage: vec![1],
4448            version: 0,
4449        };
4450
4451        let err = tensor
4452            .contiguous_values()
4453            .expect_err("overflowing span must fail");
4454        assert!(matches!(
4455            err,
4456            DenseTensorError::StorageSpanOverflow {
4457                storage_offset,
4458                numel: 1
4459            } if storage_offset == usize::MAX
4460        ));
4461    }
4462
4463    #[test]
4464    fn dense_bool_tensor_contiguous_values_rejects_storage_span_overflow() {
4465        let tensor = DenseBoolTensor {
4466            id: 1,
4467            storage_id: 1,
4468            meta: TensorMeta::from_shape(vec![1], DType::Bool, Device::Cpu)
4469                .with_storage_offset(usize::MAX),
4470            storage: vec![1],
4471            version: 0,
4472        };
4473
4474        let err = tensor
4475            .contiguous_values()
4476            .expect_err("overflowing span must fail");
4477        assert!(matches!(
4478            err,
4479            DenseTensorError::StorageSpanOverflow {
4480                storage_offset,
4481                numel: 1
4482            } if storage_offset == usize::MAX
4483        ));
4484    }
4485
4486    #[test]
4487    fn dense_bool_tensor_all_true() {
4488        let t = DenseBoolTensor::from_bools(&[true, true, true], vec![3], Device::Cpu)
4489            .expect("all-true tensor");
4490        assert!(t.contiguous_bools().unwrap().iter().all(|&b| b));
4491    }
4492
4493    #[test]
4494    fn dense_bool_tensor_all_false() {
4495        let t = DenseBoolTensor::from_bools(&[false, false, false], vec![3], Device::Cpu)
4496            .expect("all-false tensor");
4497        assert!(t.contiguous_bools().unwrap().iter().all(|&b| !b));
4498    }
4499
4500    #[test]
4501    fn dense_bool_tensor_empty() {
4502        let t = DenseBoolTensor::from_bools(&[], vec![0], Device::Cpu).expect("empty bool tensor");
4503        assert_eq!(t.meta().numel(), 0);
4504        assert_eq!(t.contiguous_values().unwrap(), &[] as &[u8]);
4505    }
4506
4507    #[test]
4508    fn dense_bool_tensor_scalar() {
4509        let t =
4510            DenseBoolTensor::from_bools(&[true], vec![], Device::Cpu).expect("scalar bool tensor");
4511        assert_eq!(t.meta().numel(), 1);
4512        assert_eq!(t.contiguous_bools().unwrap(), vec![true]);
4513    }
4514
4515    #[test]
4516    fn dense_bool_tensor_storage_accessor() {
4517        let t = DenseBoolTensor::from_bools(&[false, true], vec![2], Device::Cpu)
4518            .expect("create bool tensor");
4519        assert_eq!(t.storage(), &[0u8, 1]);
4520    }
4521
4522    #[test]
4523    fn dense_bool_tensor_has_unique_ids() {
4524        let t1 = DenseBoolTensor::from_bools(&[true], vec![1], Device::Cpu).expect("t1");
4525        let t2 = DenseBoolTensor::from_bools(&[false], vec![1], Device::Cpu).expect("t2");
4526        assert_ne!(t1.id(), t2.id());
4527        assert_ne!(t1.storage_id(), t2.storage_id());
4528    }
4529
4530    #[test]
4531    fn dense_bool_tensor_version_starts_at_zero() {
4532        let t =
4533            DenseBoolTensor::from_bools(&[true, false], vec![2], Device::Cpu).expect("bool tensor");
4534        assert_eq!(t.version(), 0);
4535    }
4536
4537    // ── bd-2do9.3: TensorStorage and F32 DenseTensor tests ──────────────
4538
4539    #[test]
4540    fn tensor_storage_f32_basic_ops() {
4541        let s = TensorStorage::F32(Arc::new(vec![1.0f32, 2.0, 3.0]));
4542        assert_eq!(s.len(), 3);
4543        assert!(!s.is_empty());
4544        assert_eq!(s.dtype(), DType::F32);
4545        assert!(s.as_f32().is_some());
4546        assert!(s.as_f64().is_none());
4547        assert_eq!(s.as_f32().unwrap(), &[1.0f32, 2.0, 3.0]);
4548    }
4549
4550    #[test]
4551    fn tensor_storage_f64_basic_ops() {
4552        let s = TensorStorage::F64(Arc::new(vec![1.0, 2.0]));
4553        assert_eq!(s.len(), 2);
4554        assert_eq!(s.dtype(), DType::F64);
4555        assert!(s.as_f64().is_some());
4556        assert!(s.as_f32().is_none());
4557    }
4558
4559    #[test]
4560    fn tensor_storage_f64_inline4_basic_ops() {
4561        let s = TensorStorage::F64Inline4([1.0, 2.0, 3.0, 4.0]);
4562        assert_eq!(s.len(), 4);
4563        assert_eq!(s.dtype(), DType::F64);
4564        assert_eq!(s.as_f64().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
4565        assert_eq!(s.to_f64_vec(), vec![1.0, 2.0, 3.0, 4.0]);
4566    }
4567
4568    #[test]
4569    fn tensor_storage_empty() {
4570        let s = TensorStorage::F32(Arc::new(Vec::new()));
4571        assert!(s.is_empty());
4572        assert_eq!(s.len(), 0);
4573    }
4574
4575    #[test]
4576    fn tensor_storage_to_f64_vec_from_f32() {
4577        let s = TensorStorage::F32(Arc::new(vec![1.5f32, 2.5]));
4578        let v = s.to_f64_vec();
4579        assert_eq!(v, vec![1.5f64, 2.5]);
4580    }
4581
4582    #[test]
4583    fn tensor_storage_to_f32_vec_from_f64() {
4584        let s = TensorStorage::F64(Arc::new(vec![1.5, 2.5]));
4585        let v = s.to_f32_vec();
4586        assert_eq!(v, vec![1.5f32, 2.5]);
4587    }
4588
4589    #[test]
4590    fn dense_tensor_f32_creation() {
4591        let dt =
4592            DenseTensor::from_contiguous_values_f32(vec![1.0f32, 2.0, 3.0], vec![3], Device::Cpu)
4593                .expect("create f32 dense tensor");
4594        assert_eq!(dt.meta().dtype(), DType::F32);
4595        assert_eq!(dt.contiguous_values_f32().unwrap(), &[1.0f32, 2.0, 3.0]);
4596        assert!(dt.contiguous_values().is_err()); // f64 access on f32 tensor fails
4597    }
4598
4599    #[test]
4600    fn dense_tensor_f32_contiguous_values_as_f64() {
4601        let dt =
4602            DenseTensor::from_contiguous_values_f32(vec![1.5f32, 2.5, 3.5], vec![3], Device::Cpu)
4603                .expect("create f32 dense tensor");
4604        let f64_vals = dt.contiguous_values_as_f64().unwrap();
4605        assert_eq!(f64_vals, vec![1.5f64, 2.5, 3.5]);
4606    }
4607
4608    #[test]
4609    fn dense_tensor_f64_contiguous_values_as_f64() {
4610        let dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0], vec![2], Device::Cpu)
4611            .expect("create f64 dense tensor");
4612        let f64_vals = dt.contiguous_values_as_f64().unwrap();
4613        assert_eq!(f64_vals, vec![1.0, 2.0]);
4614    }
4615
4616    #[test]
4617    fn dense_tensor_f64_inline4_storage_access_and_update() {
4618        let meta = TensorMeta::from_shape(vec![4], DType::F64, Device::Cpu);
4619        let tensor = DenseTensor::from_storage_f64_inline4(meta, [1.0, 2.0, 3.0, 4.0])
4620            .expect("create inline f64 tensor");
4621        assert_eq!(tensor.storage().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
4622        assert_eq!(
4623            tensor.typed_storage().as_f64().unwrap(),
4624            &[1.0, 2.0, 3.0, 4.0]
4625        );
4626
4627        let mut updated = tensor.clone();
4628        updated
4629            .update_contiguous_values(&[5.0, 6.0, 7.0, 8.0])
4630            .expect("inline update should work");
4631        assert_eq!(tensor.contiguous_values().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
4632        assert_eq!(updated.contiguous_values().unwrap(), &[5.0, 6.0, 7.0, 8.0]);
4633    }
4634
4635    #[test]
4636    fn dense_tensor_to_dtype_f64_to_f32() {
4637        let dt = DenseTensor::from_contiguous_values(vec![1.5, 2.5], vec![2], Device::Cpu)
4638            .expect("create f64 tensor");
4639        let f32_dt = dt.to_dtype(DType::F32).expect("cast to f32");
4640        assert_eq!(f32_dt.meta().dtype(), DType::F32);
4641        assert_eq!(f32_dt.contiguous_values_f32().unwrap(), &[1.5f32, 2.5]);
4642    }
4643
4644    #[test]
4645    fn dense_tensor_to_dtype_f32_to_f64() {
4646        let dt = DenseTensor::from_contiguous_values_f32(vec![1.5f32, 2.5], vec![2], Device::Cpu)
4647            .expect("create f32 tensor");
4648        let f64_dt = dt.to_dtype(DType::F64).expect("cast to f64");
4649        assert_eq!(f64_dt.meta().dtype(), DType::F64);
4650        assert_eq!(f64_dt.contiguous_values().unwrap(), &[1.5, 2.5]);
4651    }
4652
4653    #[test]
4654    fn dense_tensor_to_dtype_compacts_offset_quantized_view() {
4655        let meta = TensorMeta::quantized_from_shape_and_strides(
4656            vec![2],
4657            vec![1],
4658            1,
4659            DType::QInt8,
4660            Device::Cpu,
4661            0.5,
4662            -2,
4663        )
4664        .expect("offset quantized meta");
4665        let tensor =
4666            DenseTensor::from_storage_qint8(meta, vec![99, 0, 2]).expect("offset quantized tensor");
4667
4668        let cast = tensor.to_dtype(DType::F64).expect("cast offset qint8");
4669
4670        assert_eq!(cast.meta().dtype(), DType::F64);
4671        assert_eq!(cast.meta().storage_offset(), 0);
4672        assert_eq!(cast.typed_storage().as_f64().unwrap(), &[1.0, 2.0]);
4673        assert_eq!(cast.contiguous_values().unwrap(), &[1.0, 2.0]);
4674    }
4675
4676    #[test]
4677    fn dense_tensor_to_dtype_compacts_offset_quint8_view() {
4678        let meta = TensorMeta::quantized_from_shape_and_strides(
4679            vec![2],
4680            vec![1],
4681            2,
4682            DType::QUInt8,
4683            Device::Cpu,
4684            0.25,
4685            10,
4686        )
4687        .expect("offset quantized meta");
4688        let tensor = DenseTensor::from_storage_quint8(meta, vec![1, 3, 10, 14])
4689            .expect("offset quantized tensor");
4690
4691        let cast = tensor.to_dtype(DType::F32).expect("cast offset quint8");
4692
4693        assert_eq!(cast.meta().dtype(), DType::F32);
4694        assert_eq!(cast.meta().storage_offset(), 0);
4695        assert_eq!(cast.typed_storage().as_f32().unwrap(), &[0.0, 1.0]);
4696        assert_eq!(cast.contiguous_values_f32().unwrap(), &[0.0, 1.0]);
4697    }
4698
4699    #[test]
4700    fn dense_tensor_to_dtype_compacts_offset_float_view() {
4701        let meta = TensorMeta::from_shape(vec![2], DType::F32, Device::Cpu).with_storage_offset(2);
4702        let tensor = DenseTensor::from_storage_f32(meta, vec![99.0, 98.0, 1.5, 2.5])
4703            .expect("offset f32 tensor");
4704
4705        let cast = tensor.to_dtype(DType::F64).expect("cast offset f32");
4706
4707        assert_eq!(cast.meta().dtype(), DType::F64);
4708        assert_eq!(cast.meta().storage_offset(), 0);
4709        assert_eq!(cast.typed_storage().as_f64().unwrap(), &[1.5, 2.5]);
4710        assert_eq!(cast.contiguous_values().unwrap(), &[1.5, 2.5]);
4711    }
4712
4713    #[test]
4714    fn dense_tensor_to_dtype_compacts_offset_complex_view() {
4715        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu).with_storage_offset(1);
4716        let storage = TensorStorage::F64(Arc::new(vec![9.0, 1.5, 2.5, 8.0]));
4717        let tensor = DenseTensor::from_typed_storage(meta, storage).expect("offset f64 tensor");
4718
4719        let cast = tensor.to_dtype(DType::Complex128).expect("cast offset f64");
4720
4721        assert_eq!(cast.meta().dtype(), DType::Complex128);
4722        assert_eq!(cast.meta().storage_offset(), 0);
4723        match cast.typed_storage() {
4724            TensorStorage::Complex128(values) => {
4725                let values: Vec<(f64, f64)> =
4726                    values.iter().map(|value| (value.re, value.im)).collect();
4727                assert_eq!(values, vec![(1.5, 0.0), (2.5, 0.0)]);
4728            }
4729            other => panic!("expected Complex128 storage, got {other:?}"),
4730        }
4731    }
4732
4733    #[test]
4734    fn dense_tensor_to_dtype_preserves_complex64_imaginary_parts() {
4735        let meta =
4736            TensorMeta::from_shape(vec![2], DType::Complex64, Device::Cpu).with_storage_offset(1);
4737        let storage = TensorStorage::Complex64(Arc::new(vec![
4738            Complex64::new(9.0, 9.0),
4739            Complex64::new(1.25, -2.5),
4740            Complex64::new(-3.5, 4.75),
4741        ]));
4742        let tensor = DenseTensor::from_typed_storage(meta, storage).expect("offset c64 tensor");
4743
4744        let cast = tensor
4745            .to_dtype(DType::Complex128)
4746            .expect("cast c64 to c128");
4747
4748        assert_eq!(cast.meta().dtype(), DType::Complex128);
4749        assert_eq!(cast.meta().storage_offset(), 0);
4750        match cast.typed_storage() {
4751            TensorStorage::Complex128(values) => {
4752                let values: Vec<(f64, f64)> =
4753                    values.iter().map(|value| (value.re, value.im)).collect();
4754                assert_eq!(values, vec![(1.25, -2.5), (-3.5, 4.75)]);
4755            }
4756            other => panic!("expected Complex128 storage, got {other:?}"),
4757        }
4758    }
4759
4760    #[test]
4761    fn dense_tensor_to_dtype_preserves_complex128_imaginary_parts() {
4762        let meta =
4763            TensorMeta::from_shape(vec![2], DType::Complex128, Device::Cpu).with_storage_offset(1);
4764        let storage = TensorStorage::Complex128(Arc::new(vec![
4765            Complex128::new(9.0, 9.0),
4766            Complex128::new(1.25, -2.5),
4767            Complex128::new(-3.5, 4.75),
4768        ]));
4769        let tensor = DenseTensor::from_typed_storage(meta, storage).expect("offset c128 tensor");
4770
4771        let cast = tensor.to_dtype(DType::Complex64).expect("cast c128 to c64");
4772
4773        assert_eq!(cast.meta().dtype(), DType::Complex64);
4774        assert_eq!(cast.meta().storage_offset(), 0);
4775        match cast.typed_storage() {
4776            TensorStorage::Complex64(values) => {
4777                let values: Vec<(f32, f32)> =
4778                    values.iter().map(|value| (value.re, value.im)).collect();
4779                assert_eq!(values, vec![(1.25, -2.5), (-3.5, 4.75)]);
4780            }
4781            other => panic!("expected Complex64 storage, got {other:?}"),
4782        }
4783    }
4784
4785    #[test]
4786    fn dense_tensor_to_dtype_same_is_clone() {
4787        let dt = DenseTensor::from_contiguous_values(vec![1.0, 2.0], vec![2], Device::Cpu)
4788            .expect("create tensor");
4789        let same = dt.to_dtype(DType::F64).expect("same dtype");
4790        assert_eq!(same.contiguous_values().unwrap(), &[1.0, 2.0]);
4791    }
4792
4793    #[test]
4794    fn dense_tensor_to_dtype_rejects_non_float() {
4795        let dt = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu)
4796            .expect("create tensor");
4797        assert!(dt.to_dtype(DType::I64).is_err());
4798    }
4799
4800    #[test]
4801    fn dtype_promote_f32_f32() {
4802        assert_eq!(DType::F32.promote(DType::F32), Some(DType::F32));
4803    }
4804
4805    #[test]
4806    fn dtype_promote_f64_f64() {
4807        assert_eq!(DType::F64.promote(DType::F64), Some(DType::F64));
4808    }
4809
4810    #[test]
4811    fn dtype_promote_mixed() {
4812        assert_eq!(DType::F32.promote(DType::F64), Some(DType::F64));
4813        assert_eq!(DType::F64.promote(DType::F32), Some(DType::F64));
4814    }
4815
4816    #[test]
4817    fn dtype_promote_non_float_returns_none() {
4818        assert_eq!(DType::F32.promote(DType::I64), None);
4819        assert_eq!(DType::I64.promote(DType::F64), None);
4820    }
4821
4822    #[test]
4823    fn tensor_meta_with_dtype() {
4824        let meta = TensorMeta::from_shape(vec![2, 3], DType::F64, Device::Cpu);
4825        let meta_f32 = meta.clone().with_dtype(DType::F32);
4826        assert_eq!(meta_f32.dtype(), DType::F32);
4827        assert_eq!(meta_f32.shape(), meta.shape());
4828        assert_eq!(meta_f32.strides(), meta.strides());
4829    }
4830
4831    #[test]
4832    fn dense_tensor_typed_storage_accessor() {
4833        let dt = DenseTensor::from_contiguous_values_f32(vec![1.0f32, 2.0], vec![2], Device::Cpu)
4834            .expect("create f32 tensor");
4835        assert_eq!(dt.typed_storage().dtype(), DType::F32);
4836        assert_eq!(dt.typed_storage().as_f32().unwrap(), &[1.0f32, 2.0]);
4837    }
4838
4839    #[test]
4840    fn dense_tensor_f32_update_contiguous_values() {
4841        let mut dt =
4842            DenseTensor::from_contiguous_values_f32(vec![1.0f32, 2.0, 3.0], vec![3], Device::Cpu)
4843                .expect("create f32 tensor");
4844        dt.update_contiguous_values_f32(&[4.0f32, 5.0, 6.0])
4845            .expect("update should succeed");
4846        assert_eq!(dt.version(), 1);
4847        assert_eq!(dt.contiguous_values_f32().unwrap(), &[4.0f32, 5.0, 6.0]);
4848    }
4849
4850    #[test]
4851    fn dense_tensor_from_typed_storage_f32() {
4852        let meta = TensorMeta::from_shape(vec![2], DType::F32, Device::Cpu);
4853        let storage = TensorStorage::F32(Arc::new(vec![1.0f32, 2.0]));
4854        let dt = DenseTensor::from_typed_storage(meta, storage).expect("create from typed storage");
4855        assert_eq!(dt.meta().dtype(), DType::F32);
4856    }
4857
4858    #[test]
4859    fn dense_tensor_from_typed_storage_rejects_non_float() {
4860        let meta = TensorMeta::from_shape(vec![2], DType::I64, Device::Cpu);
4861        let storage = TensorStorage::F64(Arc::new(vec![1.0, 2.0]));
4862        let err = DenseTensor::from_typed_storage(meta, storage)
4863            .expect_err("non-float dtype should be rejected");
4864        assert!(matches!(
4865            err,
4866            DenseTensorError::UnsupportedDType(DType::I64)
4867        ));
4868    }
4869
4870    // ── promote_types tests ───────────────────────────────────────────
4871
4872    #[test]
4873    fn promote_types_same_dtype_is_identity() {
4874        for dtype in [
4875            DType::Bool,
4876            DType::I32,
4877            DType::I64,
4878            DType::F16,
4879            DType::BF16,
4880            DType::F32,
4881            DType::F64,
4882        ] {
4883            assert_eq!(dtype.promote_types(dtype), dtype);
4884        }
4885    }
4886
4887    #[test]
4888    fn promote_types_is_symmetric() {
4889        let dtypes = [
4890            DType::Bool,
4891            DType::I32,
4892            DType::I64,
4893            DType::F16,
4894            DType::BF16,
4895            DType::F32,
4896            DType::F64,
4897        ];
4898        for &a in &dtypes {
4899            for &b in &dtypes {
4900                assert_eq!(
4901                    a.promote_types(b),
4902                    b.promote_types(a),
4903                    "promote_types({a:?}, {b:?}) != promote_types({b:?}, {a:?})"
4904                );
4905            }
4906        }
4907    }
4908
4909    #[test]
4910    fn promote_types_bool_with_integers() {
4911        assert_eq!(DType::Bool.promote_types(DType::I32), DType::I32);
4912        assert_eq!(DType::Bool.promote_types(DType::I64), DType::I64);
4913    }
4914
4915    #[test]
4916    fn promote_types_bool_with_floats() {
4917        assert_eq!(DType::Bool.promote_types(DType::F32), DType::F32);
4918        assert_eq!(DType::Bool.promote_types(DType::F64), DType::F64);
4919    }
4920
4921    #[test]
4922    fn promote_types_int_with_int() {
4923        assert_eq!(DType::I32.promote_types(DType::I64), DType::I64);
4924    }
4925
4926    #[test]
4927    fn promote_types_int_with_float() {
4928        // Int + Float → Float (matching PyTorch: int64 + float32 → float32)
4929        assert_eq!(DType::I32.promote_types(DType::F32), DType::F32);
4930        assert_eq!(DType::I32.promote_types(DType::F64), DType::F64);
4931        assert_eq!(DType::I64.promote_types(DType::F32), DType::F32);
4932        assert_eq!(DType::I64.promote_types(DType::F64), DType::F64);
4933    }
4934
4935    #[test]
4936    fn promote_types_float_with_float() {
4937        assert_eq!(DType::F32.promote_types(DType::F64), DType::F64);
4938    }
4939
4940    #[test]
4941    fn promote_types_full_table() {
4942        // Exhaustive: every pair produces expected result
4943        let expected: &[(DType, DType, DType)] = &[
4944            (DType::Bool, DType::Bool, DType::Bool),
4945            (DType::Bool, DType::I32, DType::I32),
4946            (DType::Bool, DType::I64, DType::I64),
4947            (DType::Bool, DType::F16, DType::F16),
4948            (DType::Bool, DType::BF16, DType::BF16),
4949            (DType::Bool, DType::F32, DType::F32),
4950            (DType::Bool, DType::F64, DType::F64),
4951            (DType::I32, DType::I32, DType::I32),
4952            (DType::I32, DType::I64, DType::I64),
4953            (DType::I32, DType::F16, DType::F16),
4954            (DType::I32, DType::BF16, DType::BF16),
4955            (DType::I32, DType::F32, DType::F32),
4956            (DType::I32, DType::F64, DType::F64),
4957            (DType::I64, DType::I64, DType::I64),
4958            (DType::I64, DType::F16, DType::F16),
4959            (DType::I64, DType::BF16, DType::BF16),
4960            (DType::I64, DType::F32, DType::F32),
4961            (DType::I64, DType::F64, DType::F64),
4962            (DType::F16, DType::F16, DType::F16),
4963            (DType::F16, DType::BF16, DType::F32), // mixed half → F32
4964            (DType::F16, DType::F32, DType::F32),
4965            (DType::F16, DType::F64, DType::F64),
4966            (DType::BF16, DType::BF16, DType::BF16),
4967            (DType::BF16, DType::F32, DType::F32),
4968            (DType::BF16, DType::F64, DType::F64),
4969            (DType::F32, DType::F32, DType::F32),
4970            (DType::F32, DType::F64, DType::F64),
4971            (DType::F64, DType::F64, DType::F64),
4972        ];
4973        for &(a, b, result) in expected {
4974            assert_eq!(
4975                a.promote_types(b),
4976                result,
4977                "promote_types({a:?}, {b:?}) should be {result:?}"
4978            );
4979            assert_eq!(
4980                b.promote_types(a),
4981                result,
4982                "promote_types({b:?}, {a:?}) should be {result:?} (symmetry)"
4983            );
4984        }
4985    }
4986
4987    // ── F16 / BF16 tests ─────────────────────────────────────────────
4988
4989    #[test]
4990    fn dtype_f16_properties() {
4991        assert_eq!(DType::F16.element_size(), 2);
4992        assert!(DType::F16.is_floating_point());
4993        assert!(DType::F16.is_half());
4994        assert!(!DType::F16.is_integer());
4995        assert!(!DType::F16.is_bool());
4996    }
4997
4998    #[test]
4999    fn dtype_bf16_properties() {
5000        assert_eq!(DType::BF16.element_size(), 2);
5001        assert!(DType::BF16.is_floating_point());
5002        assert!(DType::BF16.is_half());
5003        assert!(!DType::BF16.is_integer());
5004        assert!(!DType::BF16.is_bool());
5005    }
5006
5007    #[test]
5008    fn f16_roundtrip_preserves_value() {
5009        let original: Vec<f32> = vec![1.0, 0.5, -3.25, 0.0, 100.0];
5010        let f16_vals: Vec<Float16> = original.iter().map(|&x| Float16::from_f32(x)).collect();
5011        let roundtrip: Vec<f32> = f16_vals.iter().map(|&x| x.to_f32()).collect();
5012        for (orig, rt) in original.iter().zip(roundtrip.iter()) {
5013            assert!(
5014                (orig - rt).abs() < 0.01,
5015                "f16 roundtrip failed: {orig} -> {rt}"
5016            );
5017        }
5018    }
5019
5020    #[test]
5021    fn bf16_roundtrip_preserves_value() {
5022        let original: Vec<f32> = vec![1.0, 0.5, -3.25, 0.0, 100.0];
5023        let bf16_vals: Vec<BFloat16> = original.iter().map(|&x| BFloat16::from_f32(x)).collect();
5024        let roundtrip: Vec<f32> = bf16_vals.iter().map(|&x| x.to_f32()).collect();
5025        for (orig, rt) in original.iter().zip(roundtrip.iter()) {
5026            assert!(
5027                (orig - rt).abs() < 1.0,
5028                "bf16 roundtrip failed: {orig} -> {rt}"
5029            );
5030        }
5031    }
5032
5033    #[test]
5034    fn f16_overflow_above_65504() {
5035        let big = Float16::from_f32(70000.0);
5036        assert!(
5037            big.to_f32().is_infinite(),
5038            "f16 should overflow to inf for values > 65504"
5039        );
5040    }
5041
5042    #[test]
5043    fn f16_inf_neg_inf() {
5044        let pos_inf = Float16::from_f32(f32::INFINITY);
5045        let neg_inf = Float16::from_f32(f32::NEG_INFINITY);
5046        assert!(pos_inf.to_f32().is_infinite() && pos_inf.to_f32() > 0.0);
5047        assert!(neg_inf.to_f32().is_infinite() && neg_inf.to_f32() < 0.0);
5048    }
5049
5050    #[test]
5051    fn f16_zero_and_neg_zero() {
5052        let zero = Float16::from_f32(0.0);
5053        let neg_zero = Float16::from_f32(-0.0);
5054        assert_eq!(zero.to_f32(), 0.0);
5055        assert_eq!(neg_zero.to_f32(), -0.0);
5056        // Both should compare equal as floats
5057        assert_eq!(zero.to_f32(), neg_zero.to_f32());
5058    }
5059
5060    #[test]
5061    fn bf16_same_exponent_range_as_f32() {
5062        // BF16 has same 8-bit exponent as f32, so same max magnitude
5063        let big = BFloat16::from_f32(1e38);
5064        assert!(
5065            big.to_f32().is_finite(),
5066            "bf16 should handle 1e38 (within f32 range)"
5067        );
5068    }
5069
5070    #[test]
5071    fn tensor_storage_f16_basic() {
5072        let vals: Vec<Float16> = vec![1.0f32, 2.0, 3.0]
5073            .into_iter()
5074            .map(Float16::from_f32)
5075            .collect();
5076        let s = TensorStorage::F16(Arc::new(vals));
5077        assert_eq!(s.len(), 3);
5078        assert_eq!(s.dtype(), DType::F16);
5079        assert!(!s.is_empty());
5080        // Conversion to f32
5081        let f32_vals = s.to_f32_vec();
5082        assert!((f32_vals[0] - 1.0).abs() < 0.01);
5083        assert!((f32_vals[1] - 2.0).abs() < 0.01);
5084        assert!((f32_vals[2] - 3.0).abs() < 0.01);
5085        // Conversion to f64
5086        let f64_vals = s.to_f64_vec();
5087        assert!((f64_vals[0] - 1.0).abs() < 0.01);
5088    }
5089
5090    #[test]
5091    fn tensor_storage_bf16_basic() {
5092        let vals: Vec<BFloat16> = vec![1.0f32, 2.0, 3.0]
5093            .into_iter()
5094            .map(BFloat16::from_f32)
5095            .collect();
5096        let s = TensorStorage::BF16(Arc::new(vals));
5097        assert_eq!(s.len(), 3);
5098        assert_eq!(s.dtype(), DType::BF16);
5099        let f32_vals = s.to_f32_vec();
5100        assert!((f32_vals[0] - 1.0).abs() < 0.01);
5101        assert!((f32_vals[1] - 2.0).abs() < 0.01);
5102    }
5103
5104    #[test]
5105    fn dense_tensor_f16_create_and_read() {
5106        let vals: Vec<Float16> = vec![1.0f32, 2.0, 3.0, 4.0]
5107            .into_iter()
5108            .map(Float16::from_f32)
5109            .collect();
5110        let dt = DenseTensor::from_contiguous_values_f16(vals, vec![2, 2], Device::Cpu).unwrap();
5111        assert_eq!(dt.meta().dtype(), DType::F16);
5112        assert_eq!(dt.meta().shape(), &[2, 2]);
5113        // contiguous_values_as_f64 should work
5114        let f64_vals = dt.contiguous_values_as_f64().unwrap();
5115        assert!((f64_vals[0] - 1.0).abs() < 0.01);
5116        assert!((f64_vals[3] - 4.0).abs() < 0.01);
5117    }
5118
5119    #[test]
5120    fn dense_tensor_bf16_create_and_read() {
5121        let vals: Vec<BFloat16> = vec![1.0f32, 2.0, 3.0]
5122            .into_iter()
5123            .map(BFloat16::from_f32)
5124            .collect();
5125        let dt = DenseTensor::from_contiguous_values_bf16(vals, vec![3], Device::Cpu).unwrap();
5126        assert_eq!(dt.meta().dtype(), DType::BF16);
5127        let f64_vals = dt.contiguous_values_as_f64().unwrap();
5128        assert!((f64_vals[0] - 1.0).abs() < 0.01);
5129    }
5130
5131    #[test]
5132    fn to_dtype_f32_to_f16() {
5133        let dt =
5134            DenseTensor::from_contiguous_values_f32(vec![1.0f32, 2.5, -3.0], vec![3], Device::Cpu)
5135                .unwrap();
5136        let f16_dt = dt.to_dtype(DType::F16).unwrap();
5137        assert_eq!(f16_dt.meta().dtype(), DType::F16);
5138        let vals = f16_dt.contiguous_values_as_f64().unwrap();
5139        assert!((vals[0] - 1.0).abs() < 0.01);
5140        assert!((vals[1] - 2.5).abs() < 0.01);
5141        assert!((vals[2] + 3.0).abs() < 0.01);
5142    }
5143
5144    #[test]
5145    fn to_dtype_f16_to_f32() {
5146        let vals: Vec<Float16> = vec![1.0f32, 2.5, -3.0]
5147            .into_iter()
5148            .map(Float16::from_f32)
5149            .collect();
5150        let dt = DenseTensor::from_contiguous_values_f16(vals, vec![3], Device::Cpu).unwrap();
5151        let f32_dt = dt.to_dtype(DType::F32).unwrap();
5152        assert_eq!(f32_dt.meta().dtype(), DType::F32);
5153        let vals = f32_dt.contiguous_values_f32().unwrap();
5154        assert!((vals[0] - 1.0).abs() < 0.01);
5155        assert!((vals[1] - 2.5).abs() < 0.01);
5156    }
5157
5158    #[test]
5159    fn to_dtype_f16_to_f64() {
5160        let vals: Vec<Float16> = vec![1.0f32, 2.0]
5161            .into_iter()
5162            .map(Float16::from_f32)
5163            .collect();
5164        let dt = DenseTensor::from_contiguous_values_f16(vals, vec![2], Device::Cpu).unwrap();
5165        let f64_dt = dt.to_dtype(DType::F64).unwrap();
5166        assert_eq!(f64_dt.meta().dtype(), DType::F64);
5167        let vals = f64_dt.contiguous_values().unwrap();
5168        assert!((vals[0] - 1.0).abs() < 0.01);
5169    }
5170
5171    #[test]
5172    fn to_dtype_bf16_roundtrip() {
5173        let dt =
5174            DenseTensor::from_contiguous_values_f32(vec![1.0f32, 0.5, -2.0], vec![3], Device::Cpu)
5175                .unwrap();
5176        let bf16_dt = dt.to_dtype(DType::BF16).unwrap();
5177        assert_eq!(bf16_dt.meta().dtype(), DType::BF16);
5178        let back_to_f32 = bf16_dt.to_dtype(DType::F32).unwrap();
5179        let vals = back_to_f32.contiguous_values_f32().unwrap();
5180        assert!((vals[0] - 1.0).abs() < 0.1);
5181        assert!((vals[1] - 0.5).abs() < 0.1);
5182        assert!((vals[2] + 2.0).abs() < 0.1);
5183    }
5184
5185    #[test]
5186    fn promote_f16_with_f32() {
5187        assert_eq!(DType::F16.promote(DType::F32), Some(DType::F32));
5188        assert_eq!(DType::F32.promote(DType::F16), Some(DType::F32));
5189    }
5190
5191    #[test]
5192    fn promote_f16_with_f64() {
5193        assert_eq!(DType::F16.promote(DType::F64), Some(DType::F64));
5194        assert_eq!(DType::F64.promote(DType::F16), Some(DType::F64));
5195    }
5196
5197    #[test]
5198    fn promote_f16_with_bf16() {
5199        assert_eq!(DType::F16.promote(DType::BF16), Some(DType::F32));
5200        assert_eq!(DType::BF16.promote(DType::F16), Some(DType::F32));
5201    }
5202
5203    #[test]
5204    fn promote_bf16_with_f32() {
5205        assert_eq!(DType::BF16.promote(DType::F32), Some(DType::F32));
5206        assert_eq!(DType::F32.promote(DType::BF16), Some(DType::F32));
5207    }
5208
5209    #[test]
5210    fn f16_dispatch_values_returns_error() {
5211        let vals: Vec<Float16> = vec![1.0f32, 2.0]
5212            .into_iter()
5213            .map(Float16::from_f32)
5214            .collect();
5215        let dt = DenseTensor::from_contiguous_values_f16(vals, vec![2], Device::Cpu).unwrap();
5216        assert!(dt.dispatch_values().is_err());
5217        assert!(dt.contiguous_values().is_err());
5218    }
5219
5220    #[test]
5221    fn f16_subnormal_handling() {
5222        // Smallest positive f16 subnormal: ~5.96e-8
5223        let tiny = Float16::from_f32(5.96e-8);
5224        let rt = tiny.to_f32();
5225        assert!(
5226            (0.0..1e-5).contains(&rt),
5227            "f16 subnormal should be small positive or zero"
5228        );
5229    }
5230
5231    // ── Complex dtype tests ─────────────────────────────────────────
5232
5233    #[test]
5234    fn complex_dtype_element_sizes() {
5235        assert_eq!(DType::Complex64.element_size(), 8);
5236        assert_eq!(DType::Complex128.element_size(), 16);
5237    }
5238
5239    #[test]
5240    fn complex_dtype_predicates() {
5241        assert!(DType::Complex64.is_complex());
5242        assert!(DType::Complex128.is_complex());
5243        assert!(!DType::F64.is_complex());
5244        assert!(!DType::Complex64.is_floating_point());
5245        assert!(!DType::Complex128.is_integer());
5246        assert!(!DType::Complex64.is_bool());
5247    }
5248
5249    #[test]
5250    fn complex_promote_types_hierarchy() {
5251        // Complex128 is the widest type
5252        assert_eq!(
5253            DType::Complex128.promote_types(DType::Complex64),
5254            DType::Complex128
5255        );
5256        assert_eq!(
5257            DType::Complex128.promote_types(DType::F64),
5258            DType::Complex128
5259        );
5260        assert_eq!(
5261            DType::Complex128.promote_types(DType::F32),
5262            DType::Complex128
5263        );
5264
5265        // Complex64 + F64 widens to Complex128 (f64 component)
5266        assert_eq!(
5267            DType::Complex64.promote_types(DType::F64),
5268            DType::Complex128
5269        );
5270        assert_eq!(
5271            DType::F64.promote_types(DType::Complex64),
5272            DType::Complex128
5273        );
5274
5275        // Complex64 + F32 stays Complex64
5276        assert_eq!(DType::Complex64.promote_types(DType::F32), DType::Complex64);
5277
5278        // Complex64 + integer → Complex64
5279        assert_eq!(DType::Complex64.promote_types(DType::I32), DType::Complex64);
5280        assert_eq!(
5281            DType::Complex64.promote_types(DType::Bool),
5282            DType::Complex64
5283        );
5284    }
5285
5286    #[test]
5287    fn complex_promote_float_function() {
5288        // promote() handles complex types
5289        assert_eq!(
5290            DType::Complex128.promote(DType::Complex64),
5291            Some(DType::Complex128)
5292        );
5293        assert_eq!(
5294            DType::Complex64.promote(DType::F64),
5295            Some(DType::Complex128)
5296        );
5297        assert_eq!(DType::Complex64.promote(DType::F32), Some(DType::Complex64));
5298        assert_eq!(
5299            DType::Complex128.promote(DType::F64),
5300            Some(DType::Complex128)
5301        );
5302    }
5303
5304    #[test]
5305    fn complex_storage_basic() {
5306        use super::Complex128;
5307
5308        let vals = vec![Complex128::new(1.0, 2.0), Complex128::new(3.0, 4.0)];
5309        let storage = TensorStorage::Complex128(Arc::new(vals));
5310        assert_eq!(storage.len(), 2);
5311        assert_eq!(storage.dtype(), DType::Complex128);
5312
5313        let slice = storage.as_complex128().unwrap();
5314        assert_eq!(slice[0].re, 1.0);
5315        assert_eq!(slice[0].im, 2.0);
5316
5317        // to_f64_vec extracts real parts
5318        let f64s = storage.to_f64_vec();
5319        assert_eq!(f64s, vec![1.0, 3.0]);
5320    }
5321
5322    #[test]
5323    fn complex64_storage_basic() {
5324        use super::Complex64;
5325
5326        let vals = vec![Complex64::new(1.0, -1.0), Complex64::new(0.0, 5.0)];
5327        let storage = TensorStorage::Complex64(Arc::new(vals));
5328        assert_eq!(storage.len(), 2);
5329        assert_eq!(storage.dtype(), DType::Complex64);
5330
5331        let slice = storage.as_complex64().unwrap();
5332        assert_eq!(slice[1].im, 5.0);
5333    }
5334
5335    #[test]
5336    fn complex_dense_tensor_creation() {
5337        use super::Complex128;
5338
5339        let vals = vec![
5340            Complex128::new(1.0, 0.0),
5341            Complex128::new(0.0, 1.0),
5342            Complex128::new(-1.0, 0.0),
5343        ];
5344        let meta = TensorMeta::from_shape(vec![3], DType::Complex128, Device::Cpu);
5345        let storage = TensorStorage::Complex128(Arc::new(vals));
5346        let dt = DenseTensor::from_typed_storage(meta, storage);
5347        assert!(dt.is_ok(), "complex tensor creation should succeed");
5348
5349        let t = dt.unwrap();
5350        assert_eq!(t.meta().dtype(), DType::Complex128);
5351        assert_eq!(t.meta().shape(), &[3]);
5352    }
5353
5354    #[test]
5355    fn complex_to_dtype_from_real() {
5356        // Create an f64 tensor, cast to Complex128
5357        let meta = TensorMeta::from_shape(vec![2], DType::F64, Device::Cpu);
5358        let storage = TensorStorage::F64(Arc::new(vec![3.0, 4.0]));
5359        let dt = DenseTensor::from_typed_storage(meta, storage).unwrap();
5360
5361        let complex = dt.to_dtype(DType::Complex128).unwrap();
5362        assert_eq!(complex.meta().dtype(), DType::Complex128);
5363
5364        let c_slice = complex.typed_storage().as_complex128().unwrap();
5365        assert_eq!(c_slice[0].re, 3.0);
5366        assert_eq!(c_slice[0].im, 0.0);
5367        assert_eq!(c_slice[1].re, 4.0);
5368        assert_eq!(c_slice[1].im, 0.0);
5369    }
5370
5371    // ── Sparse Tensor Tests ────────────────────────────────────────────────
5372
5373    #[test]
5374    fn sparse_coo_creation() {
5375        // Create a 3x4 sparse matrix with 2 non-zero elements
5376        // indices: [[0, 2], [1, 3]] (row 0 col 1, row 2 col 3)
5377        // values: [1.0, 2.0]
5378        let indices =
5379            DenseI64Tensor::from_contiguous_values(vec![0, 2, 1, 3], vec![2, 2], Device::Cpu)
5380                .unwrap();
5381
5382        let values =
5383            DenseTensor::from_contiguous_values(vec![1.0, 2.0], vec![2], Device::Cpu).unwrap();
5384
5385        let sparse = SparseCOOTensor::new(indices, values, vec![3, 4], true).unwrap();
5386
5387        assert_eq!(sparse.dense_shape(), &[3, 4]);
5388        assert_eq!(sparse.sparse_dim(), 2);
5389        assert_eq!(sparse.nnz(), 2);
5390        assert!(sparse.is_coalesced());
5391        assert_eq!(sparse.dtype(), DType::F64);
5392    }
5393
5394    #[test]
5395    fn sparse_coo_to_dense_roundtrip() {
5396        // Create a 2x3 sparse matrix
5397        // [[1, 0, 2],
5398        //  [0, 3, 0]]
5399        let coords = vec![vec![0, 0], vec![0, 2], vec![1, 1]];
5400        let values = vec![1.0, 2.0, 3.0];
5401
5402        let sparse =
5403            SparseCOOTensor::from_coords(&coords, values, vec![2, 3], DType::F64, Device::Cpu)
5404                .unwrap();
5405
5406        let dense = sparse.to_dense().unwrap();
5407
5408        let expected = vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0];
5409        let actual = dense.contiguous_values().unwrap();
5410        assert_eq!(actual, expected.as_slice());
5411    }
5412
5413    #[test]
5414    fn sparse_coo_respects_index_tensor_storage_offset() {
5415        let indices_meta =
5416            TensorMeta::from_shape(vec![2, 1], DType::I64, Device::Cpu).with_storage_offset(2);
5417        let indices = DenseI64Tensor::from_storage(indices_meta, vec![99, 99, 1, 2]).unwrap();
5418        let values = DenseTensor::from_contiguous_values(vec![7.0], vec![1], Device::Cpu).unwrap();
5419
5420        let sparse = SparseCOOTensor::new(indices, values, vec![3, 4], false).unwrap();
5421        let dense = sparse.to_dense().unwrap();
5422
5423        let mut expected = vec![0.0; 12];
5424        expected[6] = 7.0;
5425        assert_eq!(dense.contiguous_values().unwrap(), expected.as_slice());
5426    }
5427
5428    #[test]
5429    fn sparse_coo_to_dense_sums_duplicate_indices() {
5430        // Duplicate entry at (0, 1) should accumulate.
5431        let coords = vec![vec![0, 1], vec![0, 1]];
5432        let values = vec![1.0, 2.5];
5433
5434        let sparse =
5435            SparseCOOTensor::from_coords(&coords, values, vec![2, 2], DType::F64, Device::Cpu)
5436                .unwrap();
5437
5438        let dense = sparse.to_dense().unwrap();
5439        let expected = vec![0.0, 3.5, 0.0, 0.0];
5440        let actual = dense.contiguous_values().unwrap();
5441        assert_eq!(actual, expected.as_slice());
5442    }
5443
5444    #[test]
5445    fn sparse_coo_to_dense_preserves_complex_imaginary_values() {
5446        let indices =
5447            DenseI64Tensor::from_contiguous_values(vec![0, 0, 1, 1, 1, 0], vec![2, 3], Device::Cpu)
5448                .unwrap();
5449        let values_meta = TensorMeta::from_shape(vec![3], DType::Complex64, Device::Cpu);
5450        let values_storage = TensorStorage::Complex64(Arc::new(vec![
5451            Complex64::new(1.0, 2.0),
5452            Complex64::new(3.0, -0.5),
5453            Complex64::new(-4.0, 5.0),
5454        ]));
5455        let values = DenseTensor::from_typed_storage(values_meta, values_storage).unwrap();
5456        let sparse = SparseCOOTensor::new(indices, values, vec![2, 2], false).unwrap();
5457
5458        let dense = sparse.to_dense().unwrap();
5459
5460        assert_eq!(dense.meta().dtype(), DType::Complex64);
5461        match dense.typed_storage() {
5462            TensorStorage::Complex64(values) => {
5463                let values: Vec<(f32, f32)> =
5464                    values.iter().map(|value| (value.re, value.im)).collect();
5465                assert_eq!(
5466                    values,
5467                    vec![(0.0, 0.0), (4.0, 1.5), (-4.0, 5.0), (0.0, 0.0)]
5468                );
5469            }
5470            other => panic!("expected Complex64 dense storage, got {other:?}"),
5471        }
5472    }
5473
5474    #[test]
5475    fn sparse_coo_coalesced_rejects_duplicate_indices() {
5476        let indices =
5477            DenseI64Tensor::from_contiguous_values(vec![0, 0, 1, 1], vec![2, 2], Device::Cpu)
5478                .unwrap();
5479        let values =
5480            DenseTensor::from_contiguous_values(vec![1.0, 2.5], vec![2], Device::Cpu).unwrap();
5481
5482        let result = SparseCOOTensor::new(indices, values, vec![2, 2], true);
5483
5484        assert!(matches!(
5485            result,
5486            Err(SparseTensorError::DuplicateCooIndex { position: 1, .. })
5487        ));
5488    }
5489
5490    #[test]
5491    fn sparse_coo_coalesced_rejects_unsorted_indices() {
5492        // Coordinates are (1, 0), then (0, 1), which is decreasing in
5493        // lexicographic COO order and therefore cannot be marked coalesced.
5494        let indices =
5495            DenseI64Tensor::from_contiguous_values(vec![1, 0, 0, 1], vec![2, 2], Device::Cpu)
5496                .unwrap();
5497        let values =
5498            DenseTensor::from_contiguous_values(vec![1.0, 2.5], vec![2], Device::Cpu).unwrap();
5499
5500        let result = SparseCOOTensor::new(indices, values, vec![2, 2], true);
5501
5502        assert!(matches!(
5503            result,
5504            Err(SparseTensorError::UnsortedCooIndex { position: 1, .. })
5505        ));
5506    }
5507
5508    #[test]
5509    fn sparse_coo_to_dense_sums_duplicate_dense_blocks() {
5510        // Duplicate entry at (0, 1, :) should accumulate elementwise.
5511        let indices =
5512            DenseI64Tensor::from_contiguous_values(vec![0, 0, 1, 1], vec![2, 2], Device::Cpu)
5513                .unwrap();
5514        let values =
5515            DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2], Device::Cpu)
5516                .unwrap();
5517
5518        let sparse = SparseCOOTensor::new(indices, values, vec![2, 2, 2], false).unwrap();
5519        let dense = sparse.to_dense().unwrap();
5520
5521        let expected = vec![0.0, 0.0, 4.0, 6.0, 0.0, 0.0, 0.0, 0.0];
5522        let actual = dense.contiguous_values().unwrap();
5523        assert_eq!(actual, expected.as_slice());
5524    }
5525
5526    #[test]
5527    fn sparse_coo_from_coords_supports_dense_value_blocks() {
5528        let coords = vec![vec![0, 1], vec![1, 0]];
5529        let values = vec![1.0, 2.0, 3.0, 4.0];
5530
5531        let sparse =
5532            SparseCOOTensor::from_coords(&coords, values, vec![2, 2, 2], DType::F64, Device::Cpu)
5533                .unwrap();
5534
5535        assert_eq!(sparse.sparse_dim(), 2);
5536        assert_eq!(sparse.values().meta().shape(), &[2, 2]);
5537
5538        let dense = sparse.to_dense().unwrap();
5539        let expected = vec![0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 0.0, 0.0];
5540        assert_eq!(dense.contiguous_values().unwrap(), expected.as_slice());
5541    }
5542
5543    #[test]
5544    fn sparse_coo_to_dense_preserves_complex_dense_value_blocks() {
5545        let indices =
5546            DenseI64Tensor::from_contiguous_values(vec![0, 1, 1, 0], vec![2, 2], Device::Cpu)
5547                .unwrap();
5548        let values_meta = TensorMeta::from_shape(vec![2, 2], DType::Complex128, Device::Cpu);
5549        let values_storage = TensorStorage::Complex128(Arc::new(vec![
5550            Complex128::new(1.0, 2.0),
5551            Complex128::new(3.0, 4.0),
5552            Complex128::new(-5.0, 6.0),
5553            Complex128::new(7.0, -8.0),
5554        ]));
5555        let values = DenseTensor::from_typed_storage(values_meta, values_storage).unwrap();
5556        let sparse = SparseCOOTensor::new(indices, values, vec![2, 2, 2], false).unwrap();
5557
5558        let dense = sparse.to_dense().unwrap();
5559
5560        assert_eq!(dense.meta().dtype(), DType::Complex128);
5561        match dense.typed_storage() {
5562            TensorStorage::Complex128(values) => {
5563                let values: Vec<(f64, f64)> =
5564                    values.iter().map(|value| (value.re, value.im)).collect();
5565                assert_eq!(
5566                    values,
5567                    vec![
5568                        (0.0, 0.0),
5569                        (0.0, 0.0),
5570                        (1.0, 2.0),
5571                        (3.0, 4.0),
5572                        (-5.0, 6.0),
5573                        (7.0, -8.0),
5574                        (0.0, 0.0),
5575                        (0.0, 0.0),
5576                    ]
5577                );
5578            }
5579            other => panic!("expected Complex128 dense storage, got {other:?}"),
5580        }
5581    }
5582
5583    #[test]
5584    fn sparse_coo_from_coords_supports_rank2_dense_blocks() {
5585        let coords = vec![vec![0, 1], vec![1, 0]];
5586        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
5587
5588        let sparse = SparseCOOTensor::from_coords(
5589            &coords,
5590            values,
5591            vec![2, 2, 2, 2],
5592            DType::F64,
5593            Device::Cpu,
5594        )
5595        .unwrap();
5596
5597        assert_eq!(sparse.sparse_dim(), 2);
5598        assert_eq!(sparse.values().meta().shape(), &[2, 2, 2]);
5599
5600        let dense = sparse.to_dense().unwrap();
5601        let expected = vec![
5602            0.0, 0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 0.0, 0.0, 0.0, 0.0,
5603        ];
5604        assert_eq!(dense.contiguous_values().unwrap(), expected.as_slice());
5605    }
5606
5607    #[test]
5608    fn sparse_coo_from_coords_nonempty_preserves_requested_f32_dtype() {
5609        let coords = vec![vec![0, 1], vec![1, 0]];
5610        let sparse = SparseCOOTensor::from_coords(
5611            &coords,
5612            vec![1.5, 2.5],
5613            vec![2, 2],
5614            DType::F32,
5615            Device::Cpu,
5616        )
5617        .unwrap();
5618
5619        assert_eq!(sparse.dtype(), DType::F32);
5620        match sparse.values().typed_storage() {
5621            TensorStorage::F32(values) => assert_eq!(values.as_slice(), &[1.5f32, 2.5]),
5622            other => panic!("expected F32 values storage, got {other:?}"),
5623        }
5624
5625        let dense = sparse.to_dense().unwrap();
5626        assert_eq!(dense.meta().dtype(), DType::F32);
5627        assert_eq!(
5628            dense.contiguous_values_f32().unwrap(),
5629            &[0.0, 1.5, 2.5, 0.0]
5630        );
5631    }
5632
5633    #[test]
5634    fn sparse_coo_from_coords_nonempty_preserves_requested_bf16_dtype() {
5635        let coords = vec![vec![0, 1], vec![1, 0]];
5636        let sparse = SparseCOOTensor::from_coords(
5637            &coords,
5638            vec![1.5, 2.5],
5639            vec![2, 2],
5640            DType::BF16,
5641            Device::Cpu,
5642        )
5643        .unwrap();
5644
5645        assert_eq!(sparse.dtype(), DType::BF16);
5646        match sparse.values().typed_storage() {
5647            TensorStorage::BF16(values) => {
5648                let vals: Vec<f32> = values.iter().map(|value| value.to_f32()).collect();
5649                assert_eq!(vals, vec![1.5, 2.5]);
5650            }
5651            other => panic!("expected BF16 values storage, got {other:?}"),
5652        }
5653
5654        let dense = sparse.to_dense().unwrap();
5655        assert_eq!(dense.meta().dtype(), DType::BF16);
5656        assert_eq!(
5657            dense.contiguous_values_as_f64().unwrap(),
5658            &[0.0, 1.5, 2.5, 0.0]
5659        );
5660    }
5661
5662    #[test]
5663    fn sparse_coo_from_coords_nonempty_preserves_requested_f16_dtype() {
5664        let coords = vec![vec![0, 1], vec![1, 0]];
5665        let sparse = SparseCOOTensor::from_coords(
5666            &coords,
5667            vec![1.5, 2.5],
5668            vec![2, 2],
5669            DType::F16,
5670            Device::Cpu,
5671        )
5672        .unwrap();
5673
5674        assert_eq!(sparse.dtype(), DType::F16);
5675        match sparse.values().typed_storage() {
5676            TensorStorage::F16(values) => {
5677                let vals: Vec<f32> = values.iter().map(|value| value.to_f32()).collect();
5678                assert_eq!(vals, vec![1.5, 2.5]);
5679            }
5680            other => panic!("expected F16 values storage, got {other:?}"),
5681        }
5682
5683        let dense = sparse.to_dense().unwrap();
5684        assert_eq!(dense.meta().dtype(), DType::F16);
5685        assert_eq!(
5686            dense.contiguous_values_as_f64().unwrap(),
5687            &[0.0, 1.5, 2.5, 0.0]
5688        );
5689    }
5690
5691    #[test]
5692    fn sparse_coo_from_coords_nonempty_preserves_requested_complex64_dtype() {
5693        let coords = vec![vec![0, 1], vec![1, 0]];
5694        let sparse = SparseCOOTensor::from_coords(
5695            &coords,
5696            vec![1.5, 2.5],
5697            vec![2, 2],
5698            DType::Complex64,
5699            Device::Cpu,
5700        )
5701        .unwrap();
5702
5703        assert_eq!(sparse.dtype(), DType::Complex64);
5704        match sparse.values().typed_storage() {
5705            TensorStorage::Complex64(values) => {
5706                let vals: Vec<(f32, f32)> =
5707                    values.iter().map(|value| (value.re, value.im)).collect();
5708                assert_eq!(vals, vec![(1.5, 0.0), (2.5, 0.0)]);
5709            }
5710            other => panic!("expected Complex64 values storage, got {other:?}"),
5711        }
5712
5713        let dense = sparse.to_dense().unwrap();
5714        assert_eq!(dense.meta().dtype(), DType::Complex64);
5715        assert_eq!(
5716            dense.contiguous_values_as_f64().unwrap(),
5717            &[0.0, 1.5, 2.5, 0.0]
5718        );
5719    }
5720
5721    #[test]
5722    fn sparse_coo_from_coords_nonempty_preserves_requested_complex128_dtype() {
5723        let coords = vec![vec![0, 1], vec![1, 0]];
5724        let sparse = SparseCOOTensor::from_coords(
5725            &coords,
5726            vec![1.5, 2.5],
5727            vec![2, 2],
5728            DType::Complex128,
5729            Device::Cpu,
5730        )
5731        .unwrap();
5732
5733        assert_eq!(sparse.dtype(), DType::Complex128);
5734        match sparse.values().typed_storage() {
5735            TensorStorage::Complex128(values) => {
5736                let vals: Vec<(f64, f64)> =
5737                    values.iter().map(|value| (value.re, value.im)).collect();
5738                assert_eq!(vals, vec![(1.5, 0.0), (2.5, 0.0)]);
5739            }
5740            other => panic!("expected Complex128 values storage, got {other:?}"),
5741        }
5742
5743        let dense = sparse.to_dense().unwrap();
5744        assert_eq!(dense.meta().dtype(), DType::Complex128);
5745        assert_eq!(
5746            dense.contiguous_values_as_f64().unwrap(),
5747            &[0.0, 1.5, 2.5, 0.0]
5748        );
5749    }
5750
5751    #[test]
5752    fn sparse_coo_empty() {
5753        let sparse =
5754            SparseCOOTensor::from_coords(&[], vec![], vec![5, 5], DType::F64, Device::Cpu).unwrap();
5755
5756        assert_eq!(sparse.nnz(), 0);
5757        assert_eq!(sparse.dense_shape(), &[5, 5]);
5758
5759        let dense = sparse.to_dense().unwrap();
5760        let values = dense.contiguous_values().unwrap();
5761        assert!(values.iter().all(|&v| v == 0.0));
5762    }
5763
5764    #[test]
5765    fn sparse_coo_empty_preserves_requested_dtype() {
5766        let sparse =
5767            SparseCOOTensor::from_coords(&[], vec![], vec![2, 3], DType::F32, Device::Cpu).unwrap();
5768
5769        assert_eq!(sparse.dtype(), DType::F32);
5770        assert!(matches!(
5771            sparse.values().typed_storage(),
5772            TensorStorage::F32(values) if values.is_empty()
5773        ));
5774
5775        let dense = sparse.to_dense().unwrap();
5776        assert_eq!(dense.meta().dtype(), DType::F32);
5777        assert_eq!(dense.contiguous_values_f32().unwrap(), &[0.0; 6]);
5778    }
5779
5780    #[test]
5781    fn sparse_coo_empty_preserves_requested_bf16_dtype() {
5782        let sparse =
5783            SparseCOOTensor::from_coords(&[], vec![], vec![2, 2], DType::BF16, Device::Cpu)
5784                .unwrap();
5785
5786        assert_eq!(sparse.dtype(), DType::BF16);
5787        assert!(matches!(
5788            sparse.values().typed_storage(),
5789            TensorStorage::BF16(values) if values.is_empty()
5790        ));
5791
5792        let dense = sparse.to_dense().unwrap();
5793        assert_eq!(dense.meta().dtype(), DType::BF16);
5794        assert_eq!(dense.contiguous_values_as_f64().unwrap(), &[0.0; 4]);
5795    }
5796
5797    #[test]
5798    fn sparse_coo_empty_preserves_requested_f16_dtype() {
5799        let sparse =
5800            SparseCOOTensor::from_coords(&[], vec![], vec![2, 2], DType::F16, Device::Cpu).unwrap();
5801
5802        assert_eq!(sparse.dtype(), DType::F16);
5803        assert!(matches!(
5804            sparse.values().typed_storage(),
5805            TensorStorage::F16(values) if values.is_empty()
5806        ));
5807
5808        let dense = sparse.to_dense().unwrap();
5809        assert_eq!(dense.meta().dtype(), DType::F16);
5810        assert_eq!(dense.contiguous_values_as_f64().unwrap(), &[0.0; 4]);
5811    }
5812
5813    #[test]
5814    fn sparse_coo_empty_preserves_requested_complex64_dtype() {
5815        let sparse =
5816            SparseCOOTensor::from_coords(&[], vec![], vec![2, 2], DType::Complex64, Device::Cpu)
5817                .unwrap();
5818
5819        assert_eq!(sparse.dtype(), DType::Complex64);
5820        assert!(matches!(
5821            sparse.values().typed_storage(),
5822            TensorStorage::Complex64(values) if values.is_empty()
5823        ));
5824
5825        let dense = sparse.to_dense().unwrap();
5826        assert_eq!(dense.meta().dtype(), DType::Complex64);
5827        assert_eq!(dense.contiguous_values_as_f64().unwrap(), &[0.0; 4]);
5828    }
5829
5830    #[test]
5831    fn sparse_coo_empty_preserves_requested_complex128_dtype() {
5832        let sparse =
5833            SparseCOOTensor::from_coords(&[], vec![], vec![2, 2], DType::Complex128, Device::Cpu)
5834                .unwrap();
5835
5836        assert_eq!(sparse.dtype(), DType::Complex128);
5837        assert!(matches!(
5838            sparse.values().typed_storage(),
5839            TensorStorage::Complex128(values) if values.is_empty()
5840        ));
5841
5842        let dense = sparse.to_dense().unwrap();
5843        assert_eq!(dense.meta().dtype(), DType::Complex128);
5844        assert_eq!(dense.contiguous_values_as_f64().unwrap(), &[0.0; 4]);
5845    }
5846
5847    #[test]
5848    fn sparse_coo_index_out_of_bounds() {
5849        // Index [5, 0] is out of bounds for shape [3, 4] (row 5 >= 3)
5850        // indices shape [2, 1]: 2 sparse dims, 1 non-zero
5851        let indices =
5852            DenseI64Tensor::from_contiguous_values(vec![5, 0], vec![2, 1], Device::Cpu).unwrap();
5853        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
5854
5855        let result = SparseCOOTensor::new(indices, values, vec![3, 4], true);
5856        assert!(matches!(
5857            result,
5858            Err(SparseTensorError::IndexOutOfBounds { .. })
5859        ));
5860    }
5861
5862    #[test]
5863    fn sparse_coo_uncoalesced_out_of_bounds_rejected() {
5864        let indices =
5865            DenseI64Tensor::from_contiguous_values(vec![5, 0], vec![2, 1], Device::Cpu).unwrap();
5866        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
5867
5868        let result = SparseCOOTensor::new(indices, values, vec![3, 4], false);
5869        assert!(matches!(
5870            result,
5871            Err(SparseTensorError::IndexOutOfBounds { .. })
5872        ));
5873    }
5874
5875    #[test]
5876    fn sparse_coo_device_mismatch() {
5877        let indices =
5878            DenseI64Tensor::from_contiguous_values(vec![0, 1], vec![2, 1], Device::Cpu).unwrap();
5879        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cuda).unwrap();
5880
5881        let result = SparseCOOTensor::new(indices, values, vec![2, 2], true);
5882        assert!(matches!(
5883            result,
5884            Err(SparseTensorError::DeviceMismatch { .. })
5885        ));
5886    }
5887
5888    #[test]
5889    fn sparse_csr_creation() {
5890        // Create a 3x4 sparse matrix:
5891        // [[1, 0, 2, 0],
5892        //  [0, 0, 0, 3],
5893        //  [4, 0, 0, 0]]
5894        // crow_indices: [0, 2, 3, 4] (row 0 has 2 elems, row 1 has 1, row 2 has 1)
5895        // col_indices: [0, 2, 3, 0]
5896        // values: [1, 2, 3, 4]
5897        let crow =
5898            DenseI64Tensor::from_contiguous_values(vec![0, 2, 3, 4], vec![4], Device::Cpu).unwrap();
5899        let col =
5900            DenseI64Tensor::from_contiguous_values(vec![0, 2, 3, 0], vec![4], Device::Cpu).unwrap();
5901        let values =
5902            DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0, 4.0], vec![4], Device::Cpu)
5903                .unwrap();
5904
5905        let csr = SparseCSRTensor::new(crow, col, values, [3, 4]).unwrap();
5906
5907        assert_eq!(csr.shape(), [3, 4]);
5908        assert_eq!(csr.nrows(), 3);
5909        assert_eq!(csr.ncols(), 4);
5910        assert_eq!(csr.nnz(), 4);
5911    }
5912
5913    #[test]
5914    fn sparse_csr_to_dense() {
5915        // Same matrix as above
5916        let crow =
5917            DenseI64Tensor::from_contiguous_values(vec![0, 2, 3, 4], vec![4], Device::Cpu).unwrap();
5918        let col =
5919            DenseI64Tensor::from_contiguous_values(vec![0, 2, 3, 0], vec![4], Device::Cpu).unwrap();
5920        let values =
5921            DenseTensor::from_contiguous_values(vec![1.0, 2.0, 3.0, 4.0], vec![4], Device::Cpu)
5922                .unwrap();
5923
5924        let csr = SparseCSRTensor::new(crow, col, values, [3, 4]).unwrap();
5925        let dense = csr.to_dense().unwrap();
5926
5927        let expected = vec![1.0, 0.0, 2.0, 0.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0, 0.0, 0.0];
5928        let actual = dense.contiguous_values().unwrap();
5929        assert_eq!(actual, expected.as_slice());
5930    }
5931
5932    #[test]
5933    fn sparse_csr_to_dense_preserves_complex_imaginary_values() {
5934        let crow =
5935            DenseI64Tensor::from_contiguous_values(vec![0, 2, 3], vec![3], Device::Cpu).unwrap();
5936        let col =
5937            DenseI64Tensor::from_contiguous_values(vec![0, 2, 1], vec![3], Device::Cpu).unwrap();
5938        let values_meta = TensorMeta::from_shape(vec![3], DType::Complex128, Device::Cpu);
5939        let values_storage = TensorStorage::Complex128(Arc::new(vec![
5940            Complex128::new(1.0, 2.0),
5941            Complex128::new(3.0, -4.0),
5942            Complex128::new(-5.0, 6.0),
5943        ]));
5944        let values = DenseTensor::from_typed_storage(values_meta, values_storage).unwrap();
5945        let csr = SparseCSRTensor::new(crow, col, values, [2, 3]).unwrap();
5946
5947        let dense = csr.to_dense().unwrap();
5948
5949        assert_eq!(dense.meta().dtype(), DType::Complex128);
5950        match dense.typed_storage() {
5951            TensorStorage::Complex128(values) => {
5952                let values: Vec<(f64, f64)> =
5953                    values.iter().map(|value| (value.re, value.im)).collect();
5954                assert_eq!(
5955                    values,
5956                    vec![
5957                        (1.0, 2.0),
5958                        (0.0, 0.0),
5959                        (3.0, -4.0),
5960                        (0.0, 0.0),
5961                        (-5.0, 6.0),
5962                        (0.0, 0.0),
5963                    ]
5964                );
5965            }
5966            other => panic!("expected Complex128 dense storage, got {other:?}"),
5967        }
5968    }
5969
5970    #[test]
5971    fn sparse_csr_respects_index_tensor_storage_offsets() {
5972        let crow_meta =
5973            TensorMeta::from_shape(vec![3], DType::I64, Device::Cpu).with_storage_offset(1);
5974        let crow = DenseI64Tensor::from_storage(crow_meta, vec![99, 0, 1, 1]).unwrap();
5975        let col_meta =
5976            TensorMeta::from_shape(vec![1], DType::I64, Device::Cpu).with_storage_offset(1);
5977        let col = DenseI64Tensor::from_storage(col_meta, vec![99, 2]).unwrap();
5978        let values = DenseTensor::from_contiguous_values(vec![7.0], vec![1], Device::Cpu).unwrap();
5979
5980        let csr = SparseCSRTensor::new(crow, col, values, [2, 4]).unwrap();
5981        let dense = csr.to_dense().unwrap();
5982
5983        let mut expected = vec![0.0; 8];
5984        expected[2] = 7.0;
5985        assert_eq!(dense.contiguous_values().unwrap(), expected.as_slice());
5986    }
5987
5988    #[test]
5989    fn sparse_csr_col_out_of_bounds() {
5990        // Column index 5 is out of bounds for 4 columns
5991        let crow =
5992            DenseI64Tensor::from_contiguous_values(vec![0, 1], vec![2], Device::Cpu).unwrap();
5993        let col = DenseI64Tensor::from_contiguous_values(vec![5], vec![1], Device::Cpu).unwrap();
5994        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
5995
5996        let result = SparseCSRTensor::new(crow, col, values, [1, 4]);
5997        assert!(matches!(
5998            result,
5999            Err(SparseTensorError::ColIndexOutOfBounds { .. })
6000        ));
6001    }
6002
6003    #[test]
6004    fn sparse_csr_rejects_duplicate_row_columns() {
6005        let crow =
6006            DenseI64Tensor::from_contiguous_values(vec![0, 2], vec![2], Device::Cpu).unwrap();
6007        let col = DenseI64Tensor::from_contiguous_values(vec![0, 0], vec![2], Device::Cpu).unwrap();
6008        let values =
6009            DenseTensor::from_contiguous_values(vec![2.0, 3.0], vec![2], Device::Cpu).unwrap();
6010
6011        let result = SparseCSRTensor::new(crow, col, values, [1, 2]);
6012        assert!(matches!(
6013            result,
6014            Err(SparseTensorError::DuplicateCsrColumn { row: 0, col: 0 })
6015        ));
6016    }
6017
6018    #[test]
6019    fn sparse_csr_device_mismatch() {
6020        let crow =
6021            DenseI64Tensor::from_contiguous_values(vec![0, 1], vec![2], Device::Cuda).unwrap();
6022        let col = DenseI64Tensor::from_contiguous_values(vec![0], vec![1], Device::Cpu).unwrap();
6023        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
6024
6025        let result = SparseCSRTensor::new(crow, col, values, [1, 2]);
6026        assert!(matches!(
6027            result,
6028            Err(SparseTensorError::DeviceMismatch { .. })
6029        ));
6030    }
6031
6032    #[test]
6033    fn sparse_csr_non_monotonic_crow() {
6034        // crow_indices [0, 2, 1] is not monotonic at row 1: 2 > 1
6035        // nnz = crow[nrows] = crow[2] = 1, so col and values need 1 element
6036        let crow =
6037            DenseI64Tensor::from_contiguous_values(vec![0, 2, 1], vec![3], Device::Cpu).unwrap();
6038        let col = DenseI64Tensor::from_contiguous_values(vec![0], vec![1], Device::Cpu).unwrap();
6039        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
6040
6041        let result = SparseCSRTensor::new(crow, col, values, [2, 3]);
6042        assert!(matches!(
6043            result,
6044            Err(SparseTensorError::NonMonotonicCrowIndices { .. })
6045        ));
6046    }
6047
6048    #[test]
6049    fn sparse_csr_rejects_nonzero_crow_start() {
6050        let crow =
6051            DenseI64Tensor::from_contiguous_values(vec![1, 1], vec![2], Device::Cpu).unwrap();
6052        let col = DenseI64Tensor::from_contiguous_values(vec![0], vec![1], Device::Cpu).unwrap();
6053        let values = DenseTensor::from_contiguous_values(vec![1.0], vec![1], Device::Cpu).unwrap();
6054
6055        let result = SparseCSRTensor::new(crow, col, values, [1, 3]);
6056        assert!(matches!(
6057            result,
6058            Err(SparseTensorError::InvalidCrowIndexValue { .. })
6059        ));
6060    }
6061
6062    #[test]
6063    fn sparse_csr_rejects_negative_crow_tail() {
6064        let crow =
6065            DenseI64Tensor::from_contiguous_values(vec![0, -1], vec![2], Device::Cpu).unwrap();
6066        let col = DenseI64Tensor::from_contiguous_values(vec![], vec![0], Device::Cpu).unwrap();
6067        let values = DenseTensor::from_contiguous_values(vec![], vec![0], Device::Cpu).unwrap();
6068
6069        let result = SparseCSRTensor::new(crow, col, values, [1, 3]);
6070        assert!(matches!(
6071            result,
6072            Err(SparseTensorError::InvalidCrowIndexValue { .. })
6073        ));
6074    }
6075}