tensor4all_tensorbackend/
any_scalar.rs

1use std::cmp::Ordering;
2use std::fmt;
3use std::ops::{Add, Div, Mul, Neg, Sub};
4
5use anyhow::{anyhow, ensure, Result};
6use num_complex::{Complex32, Complex64};
7use num_traits::{One, Zero};
8use tenferro::{DType, Tensor as NativeTensor, TensorScalar};
9
10use crate::storage::{Storage, SumFromStorage};
11use crate::tensor_element::TensorElement;
12
13#[derive(Clone, Copy, Debug, PartialEq)]
14enum ScalarValue {
15    F32(f32),
16    F64(f64),
17    I32(i32),
18    I64(i64),
19    Bool(bool),
20    C32(Complex32),
21    C64(Complex64),
22}
23
24impl ScalarValue {
25    fn real(self) -> f64 {
26        match self {
27            Self::F32(value) => value as f64,
28            Self::F64(value) => value,
29            Self::I32(value) => value as f64,
30            Self::I64(value) => value as f64,
31            Self::Bool(value) => {
32                if value {
33                    1.0
34                } else {
35                    0.0
36                }
37            }
38            Self::C32(value) => value.re as f64,
39            Self::C64(value) => value.re,
40        }
41    }
42
43    fn imag(self) -> f64 {
44        match self {
45            Self::F32(_) | Self::F64(_) | Self::I32(_) | Self::I64(_) | Self::Bool(_) => 0.0,
46            Self::C32(value) => value.im as f64,
47            Self::C64(value) => value.im,
48        }
49    }
50
51    fn abs(self) -> f64 {
52        match self {
53            Self::F32(value) => value.abs() as f64,
54            Self::F64(value) => value.abs(),
55            Self::I32(value) => value.abs() as f64,
56            Self::I64(value) => value.abs() as f64,
57            Self::Bool(value) => {
58                if value {
59                    1.0
60                } else {
61                    0.0
62                }
63            }
64            Self::C32(value) => {
65                if value.re.is_nan() || value.im.is_nan() {
66                    f64::NAN
67                } else {
68                    value.norm() as f64
69                }
70            }
71            Self::C64(value) => {
72                if value.re.is_nan() || value.im.is_nan() {
73                    f64::NAN
74                } else {
75                    value.norm()
76                }
77            }
78        }
79    }
80
81    fn is_complex(self) -> bool {
82        matches!(self, Self::C32(_) | Self::C64(_))
83    }
84
85    fn is_zero(self) -> bool {
86        match self {
87            Self::F32(value) => value == 0.0,
88            Self::F64(value) => value == 0.0,
89            Self::I32(value) => value == 0,
90            Self::I64(value) => value == 0,
91            Self::Bool(value) => !value,
92            Self::C32(value) => value == Complex32::new(0.0, 0.0),
93            Self::C64(value) => value == Complex64::new(0.0, 0.0),
94        }
95    }
96
97    fn into_complex(self) -> Complex64 {
98        match self {
99            Self::F32(value) => Complex64::new(value as f64, 0.0),
100            Self::F64(value) => Complex64::new(value, 0.0),
101            Self::I32(value) => Complex64::new(value as f64, 0.0),
102            Self::I64(value) => Complex64::new(value as f64, 0.0),
103            Self::Bool(value) => Complex64::new(if value { 1.0 } else { 0.0 }, 0.0),
104            Self::C32(value) => Complex64::new(value.re as f64, value.im as f64),
105            Self::C64(value) => value,
106        }
107    }
108}
109
110fn scalar_value_from_storage(storage: &Storage) -> ScalarValue {
111    if storage.is_f64() {
112        ScalarValue::F64(f64::sum_from_storage(storage))
113    } else {
114        ScalarValue::C64(Complex64::sum_from_storage(storage))
115    }
116}
117
118fn scalar_value_from_native(native: &NativeTensor) -> Result<ScalarValue> {
119    ensure!(
120        native.shape().is_empty(),
121        "expected rank-0 scalar tensor, got shape {:?}",
122        native.shape()
123    );
124
125    match native.dtype() {
126        DType::F32 => native
127            .as_slice::<f32>()
128            .map_err(anyhow::Error::new)?
129            .first()
130            .copied()
131            .map(ScalarValue::F32)
132            .ok_or_else(|| anyhow!("failed to read f32 scalar tensor value")),
133        DType::F64 => native
134            .as_slice::<f64>()
135            .map_err(anyhow::Error::new)?
136            .first()
137            .copied()
138            .map(ScalarValue::F64)
139            .ok_or_else(|| anyhow!("failed to read f64 scalar tensor value")),
140        DType::I32 => native
141            .as_slice::<i32>()
142            .map_err(anyhow::Error::new)?
143            .first()
144            .copied()
145            .map(ScalarValue::I32)
146            .ok_or_else(|| anyhow!("failed to read i32 scalar tensor value")),
147        DType::I64 => native
148            .as_slice::<i64>()
149            .map_err(anyhow::Error::new)?
150            .first()
151            .copied()
152            .map(ScalarValue::I64)
153            .ok_or_else(|| anyhow!("failed to read i64 scalar tensor value")),
154        DType::Bool => native
155            .as_slice::<bool>()
156            .map_err(anyhow::Error::new)?
157            .first()
158            .copied()
159            .map(ScalarValue::Bool)
160            .ok_or_else(|| anyhow!("failed to read bool scalar tensor value")),
161        DType::C32 => native
162            .as_slice::<Complex32>()
163            .map_err(anyhow::Error::new)?
164            .first()
165            .copied()
166            .map(ScalarValue::C32)
167            .ok_or_else(|| anyhow!("failed to read c32 scalar tensor value")),
168        DType::C64 => native
169            .as_slice::<Complex64>()
170            .map_err(anyhow::Error::new)?
171            .first()
172            .copied()
173            .map(ScalarValue::C64)
174            .ok_or_else(|| anyhow!("failed to read c64 scalar tensor value")),
175    }
176}
177
178trait ScalarTensorElement: TensorElement {
179    fn scalar_value(value: Self) -> ScalarValue;
180}
181
182fn scalar_native<T: TensorScalar>(value: T) -> NativeTensor {
183    crate::require_invariant(
184        NativeTensor::from_vec_col_major(vec![], vec![value]),
185        "rank-0 scalar construction failed",
186    )
187}
188
189impl ScalarTensorElement for f32 {
190    fn scalar_value(value: Self) -> ScalarValue {
191        ScalarValue::F32(value)
192    }
193}
194
195impl ScalarTensorElement for f64 {
196    fn scalar_value(value: Self) -> ScalarValue {
197        ScalarValue::F64(value)
198    }
199}
200
201impl ScalarTensorElement for Complex32 {
202    fn scalar_value(value: Self) -> ScalarValue {
203        ScalarValue::C32(value)
204    }
205}
206
207impl ScalarTensorElement for Complex64 {
208    fn scalar_value(value: Self) -> ScalarValue {
209        ScalarValue::C64(value)
210    }
211}
212
213pub(crate) fn promote_scalar_native(native: &NativeTensor, target: DType) -> Result<NativeTensor> {
214    let promoted = match (scalar_value_from_native(native)?, target) {
215        (ScalarValue::F32(value), DType::F32) => BackendScalar::from_value(value),
216        (ScalarValue::F32(value), DType::F64) => BackendScalar::from_value(value as f64),
217        (ScalarValue::F32(value), DType::C32) => {
218            BackendScalar::from_value(Complex32::new(value, 0.0))
219        }
220        (ScalarValue::F32(value), DType::C64) => {
221            BackendScalar::from_value(Complex64::new(value as f64, 0.0))
222        }
223        (ScalarValue::F32(_), DType::I64) => {
224            return Err(anyhow!(
225                "cannot promote f32 scalar to i64 without truncation"
226            ));
227        }
228        (ScalarValue::F32(_), DType::I32 | DType::Bool) => {
229            return Err(anyhow!(
230                "cannot promote f32 scalar to integer/bool without truncation"
231            ));
232        }
233        (ScalarValue::F64(value), DType::F32) => BackendScalar::from_value(value as f32),
234        (ScalarValue::F64(value), DType::F64) => BackendScalar::from_value(value),
235        (ScalarValue::F64(_), DType::I64) => {
236            return Err(anyhow!(
237                "cannot promote f64 scalar to i64 without truncation"
238            ));
239        }
240        (ScalarValue::F64(_), DType::I32 | DType::Bool) => {
241            return Err(anyhow!(
242                "cannot promote f64 scalar to integer/bool without truncation"
243            ));
244        }
245        (ScalarValue::F64(value), DType::C32) => {
246            BackendScalar::from_value(Complex32::new(value as f32, 0.0))
247        }
248        (ScalarValue::F64(value), DType::C64) => {
249            BackendScalar::from_value(Complex64::new(value, 0.0))
250        }
251        (ScalarValue::I32(value), DType::F32) => BackendScalar::from_value(value as f32),
252        (ScalarValue::I32(value), DType::F64) => BackendScalar::from_value(value as f64),
253        (ScalarValue::I32(value), DType::I32) => return Ok(scalar_native(value)),
254        (ScalarValue::I32(value), DType::I64) => BackendScalar::from_i64(value as i64),
255        (ScalarValue::I32(value), DType::C32) => {
256            BackendScalar::from_value(Complex32::new(value as f32, 0.0))
257        }
258        (ScalarValue::I32(value), DType::C64) => {
259            BackendScalar::from_value(Complex64::new(value as f64, 0.0))
260        }
261        (ScalarValue::I32(_), DType::Bool) => {
262            return Err(anyhow!("cannot promote i32 scalar to bool"));
263        }
264        (ScalarValue::I64(value), DType::F32) => BackendScalar::from_value(value as f32),
265        (ScalarValue::I64(value), DType::F64) => BackendScalar::from_value(value as f64),
266        (ScalarValue::I64(_), DType::I32 | DType::Bool) => {
267            return Err(anyhow!("cannot promote i64 scalar to i32/bool"));
268        }
269        (ScalarValue::I64(value), DType::I64) => BackendScalar::from_i64(value),
270        (ScalarValue::I64(value), DType::C32) => {
271            BackendScalar::from_value(Complex32::new(value as f32, 0.0))
272        }
273        (ScalarValue::I64(value), DType::C64) => {
274            BackendScalar::from_value(Complex64::new(value as f64, 0.0))
275        }
276        (ScalarValue::Bool(value), DType::F32) => {
277            BackendScalar::from_value(if value { 1.0_f32 } else { 0.0_f32 })
278        }
279        (ScalarValue::Bool(value), DType::F64) => {
280            BackendScalar::from_value(if value { 1.0 } else { 0.0 })
281        }
282        (ScalarValue::Bool(value), DType::I32) => {
283            return Ok(scalar_native(if value { 1 } else { 0 }));
284        }
285        (ScalarValue::Bool(value), DType::I64) => {
286            BackendScalar::from_i64(if value { 1 } else { 0 })
287        }
288        (ScalarValue::Bool(value), DType::Bool) => return Ok(scalar_native(value)),
289        (ScalarValue::Bool(value), DType::C32) => {
290            BackendScalar::from_value(Complex32::new(if value { 1.0 } else { 0.0 }, 0.0))
291        }
292        (ScalarValue::Bool(value), DType::C64) => {
293            BackendScalar::from_value(Complex64::new(if value { 1.0 } else { 0.0 }, 0.0))
294        }
295        (ScalarValue::C32(value), DType::F32) => BackendScalar::from_value(value.re),
296        (ScalarValue::C32(value), DType::F64) => BackendScalar::from_value(value.re as f64),
297        (ScalarValue::C32(_), DType::I32 | DType::I64 | DType::Bool) => {
298            return Err(anyhow!("cannot promote c32 scalar to integer/bool"));
299        }
300        (ScalarValue::C32(value), DType::C32) => BackendScalar::from_value(value),
301        (ScalarValue::C32(value), DType::C64) => {
302            BackendScalar::from_value(Complex64::new(value.re as f64, value.im as f64))
303        }
304        (ScalarValue::C64(value), DType::F32) => BackendScalar::from_value(value.re as f32),
305        (ScalarValue::C64(value), DType::F64) => BackendScalar::from_value(value.re),
306        (ScalarValue::C64(_), DType::I32 | DType::I64 | DType::Bool) => {
307            return Err(anyhow!("cannot promote c64 scalar to integer/bool"));
308        }
309        (ScalarValue::C64(value), DType::C32) => {
310            BackendScalar::from_value(Complex32::new(value.re as f32, value.im as f32))
311        }
312        (ScalarValue::C64(value), DType::C64) => BackendScalar::from_value(value),
313    };
314    Ok(promoted.native)
315}
316
317/// Dynamic scalar used across tensor4all backends.
318///
319/// This is a tensor4all-owned rank-0 wrapper over tenferro's dynamic tensor.
320///
321/// # Examples
322///
323/// ```
324/// use tensor4all_tensorbackend::BackendScalar;
325///
326/// // Real scalar
327/// let a = BackendScalar::new_real(3.14);
328/// assert!((a.real() - 3.14).abs() < 1e-10);
329/// assert_eq!(a.imag(), 0.0);
330/// assert!(a.is_real());
331/// assert!(!a.is_complex());
332///
333/// // Complex scalar
334/// let b = BackendScalar::new_complex(1.0, 2.0);
335/// assert!((b.real() - 1.0).abs() < 1e-10);
336/// assert!((b.imag() - 2.0).abs() < 1e-10);
337/// assert!(b.is_complex());
338///
339/// // Arithmetic
340/// let c = BackendScalar::new_real(2.0);
341/// let d = a + c;
342/// assert!((d.real() - 5.14).abs() < 1e-10);
343/// ```
344pub struct BackendScalar {
345    native: NativeTensor,
346    value: ScalarValue,
347}
348
349impl BackendScalar {
350    fn wrap_native(native: NativeTensor) -> Result<Self> {
351        if native.shape().is_empty() {
352            let value = scalar_value_from_native(&native)?;
353            Ok(Self { native, value })
354        } else {
355            Err(anyhow!(
356                "BackendScalar requires a rank-0 tensor, got shape {:?}",
357                native.shape()
358            ))
359        }
360    }
361
362    fn value(&self) -> ScalarValue {
363        self.value
364    }
365
366    fn from_i64(value: i64) -> Self {
367        Self {
368            native: scalar_native(value),
369            value: ScalarValue::I64(value),
370        }
371    }
372
373    fn from_i32(value: i32) -> Self {
374        Self {
375            native: scalar_native(value),
376            value: ScalarValue::I32(value),
377        }
378    }
379
380    fn from_bool(value: bool) -> Self {
381        Self {
382            native: scalar_native(value),
383            value: ScalarValue::Bool(value),
384        }
385    }
386
387    pub(crate) fn from_native(value: NativeTensor) -> Result<Self> {
388        Self::wrap_native(value)
389    }
390
391    pub(crate) fn as_native(&self) -> &NativeTensor {
392        &self.native
393    }
394
395    /// Creates a scalar from any supported public tensor element type.
396    ///
397    /// # Examples
398    ///
399    /// ```
400    /// use tensor4all_tensorbackend::BackendScalar;
401    ///
402    /// let s = BackendScalar::from_value(3.14_f64);
403    /// assert!((s.real() - 3.14).abs() < 1e-10);
404    ///
405    /// use num_complex::Complex64;
406    /// let z = BackendScalar::from_value(Complex64::new(1.0, 2.0));
407    /// assert!(z.is_complex());
408    /// assert!((z.real() - 1.0).abs() < 1e-10);
409    /// assert!((z.imag() - 2.0).abs() < 1e-10);
410    /// ```
411    #[allow(private_bounds)]
412    pub fn from_value<T: ScalarTensorElement>(value: T) -> Self {
413        Self {
414            native: scalar_native(value),
415            value: T::scalar_value(value),
416        }
417    }
418
419    /// Creates a real scalar from an `f64` value.
420    ///
421    /// # Examples
422    ///
423    /// ```
424    /// use tensor4all_tensorbackend::BackendScalar;
425    ///
426    /// let s = BackendScalar::from_real(2.5);
427    /// assert!((s.real() - 2.5).abs() < 1e-10);
428    /// assert!(s.is_real());
429    /// ```
430    pub fn from_real(x: f64) -> Self {
431        Self::from_value(x)
432    }
433
434    /// Creates a complex scalar from real and imaginary parts.
435    ///
436    /// # Examples
437    ///
438    /// ```
439    /// use tensor4all_tensorbackend::BackendScalar;
440    ///
441    /// let s = BackendScalar::from_complex(1.0, -1.0);
442    /// assert!((s.real() - 1.0).abs() < 1e-10);
443    /// assert!((s.imag() - (-1.0)).abs() < 1e-10);
444    /// assert!(s.is_complex());
445    /// ```
446    pub fn from_complex(re: f64, im: f64) -> Self {
447        Self::from_value(Complex64::new(re, im))
448    }
449
450    /// Backward-compatible constructor for a real scalar.
451    ///
452    /// # Examples
453    ///
454    /// ```
455    /// use tensor4all_tensorbackend::BackendScalar;
456    ///
457    /// let s = BackendScalar::new_real(42.0);
458    /// assert!((s.real() - 42.0).abs() < 1e-10);
459    /// ```
460    pub fn new_real(x: f64) -> Self {
461        Self::from_real(x)
462    }
463
464    /// Backward-compatible constructor for a complex scalar.
465    ///
466    /// # Examples
467    ///
468    /// ```
469    /// use tensor4all_tensorbackend::BackendScalar;
470    ///
471    /// let s = BackendScalar::new_complex(3.0, 4.0);
472    /// assert!((s.abs() - 5.0).abs() < 1e-10); // |3 + 4i| = 5
473    /// ```
474    pub fn new_complex(re: f64, im: f64) -> Self {
475        Self::from_complex(re, im)
476    }
477
478    /// Returns the scalar's plain rank-0 tensor value.
479    ///
480    /// # Examples
481    ///
482    /// ```
483    /// use tensor4all_tensorbackend::BackendScalar;
484    ///
485    /// let s = BackendScalar::new_real(5.0);
486    /// let p = s.primal();
487    /// assert!((p.real() - 5.0).abs() < 1e-10);
488    /// ```
489    pub fn primal(&self) -> Self {
490        self.clone()
491    }
492
493    /// Returns the real part while intentionally dropping AD metadata.
494    ///
495    /// # Examples
496    ///
497    /// ```
498    /// use tensor4all_tensorbackend::BackendScalar;
499    ///
500    /// let s = BackendScalar::new_complex(3.0, 4.0);
501    /// assert!((s.real() - 3.0).abs() < 1e-10);
502    /// ```
503    pub fn real(&self) -> f64 {
504        self.value().real()
505    }
506
507    /// Returns the imaginary part while intentionally dropping AD metadata.
508    ///
509    /// Returns `0.0` for real scalars.
510    ///
511    /// # Examples
512    ///
513    /// ```
514    /// use tensor4all_tensorbackend::BackendScalar;
515    ///
516    /// let s = BackendScalar::new_complex(3.0, 4.0);
517    /// assert!((s.imag() - 4.0).abs() < 1e-10);
518    ///
519    /// let r = BackendScalar::new_real(5.0);
520    /// assert_eq!(r.imag(), 0.0);
521    /// ```
522    pub fn imag(&self) -> f64 {
523        self.value().imag()
524    }
525
526    /// Returns the absolute value while intentionally dropping AD metadata.
527    ///
528    /// For complex scalars, returns the complex modulus (`sqrt(re^2 + im^2)`).
529    ///
530    /// # Examples
531    ///
532    /// ```
533    /// use tensor4all_tensorbackend::BackendScalar;
534    ///
535    /// let s = BackendScalar::new_complex(3.0, 4.0);
536    /// assert!((s.abs() - 5.0).abs() < 1e-10);
537    ///
538    /// let r = BackendScalar::new_real(-7.0);
539    /// assert!((r.abs() - 7.0).abs() < 1e-10);
540    /// ```
541    pub fn abs(&self) -> f64 {
542        self.value().abs()
543    }
544
545    /// Returns true when the scalar is complex.
546    ///
547    /// # Examples
548    ///
549    /// ```
550    /// use tensor4all_tensorbackend::BackendScalar;
551    ///
552    /// assert!(BackendScalar::new_complex(1.0, 0.0).is_complex());
553    /// assert!(!BackendScalar::new_real(1.0).is_complex());
554    /// ```
555    pub fn is_complex(&self) -> bool {
556        self.value().is_complex()
557    }
558
559    /// Returns true when the scalar is real.
560    ///
561    /// # Examples
562    ///
563    /// ```
564    /// use tensor4all_tensorbackend::BackendScalar;
565    ///
566    /// assert!(BackendScalar::new_real(1.0).is_real());
567    /// assert!(!BackendScalar::new_complex(1.0, 2.0).is_real());
568    /// ```
569    pub fn is_real(&self) -> bool {
570        !self.is_complex()
571    }
572
573    /// Returns true when the scalar is zero.
574    ///
575    /// # Examples
576    ///
577    /// ```
578    /// use tensor4all_tensorbackend::BackendScalar;
579    ///
580    /// assert!(BackendScalar::new_real(0.0).is_zero());
581    /// assert!(!BackendScalar::new_real(1.0).is_zero());
582    /// ```
583    pub fn is_zero(&self) -> bool {
584        self.value().is_zero()
585    }
586
587    /// Returns the underlying value as `f64` when real.
588    ///
589    /// Returns `None` for complex scalars.
590    ///
591    /// # Examples
592    ///
593    /// ```
594    /// use tensor4all_tensorbackend::BackendScalar;
595    ///
596    /// let r = BackendScalar::new_real(2.5);
597    /// assert_eq!(r.as_f64(), Some(2.5));
598    ///
599    /// let c = BackendScalar::new_complex(1.0, 1.0);
600    /// assert_eq!(c.as_f64(), None);
601    /// ```
602    pub fn as_f64(&self) -> Option<f64> {
603        match self.value() {
604            ScalarValue::F32(value) => Some(value as f64),
605            ScalarValue::F64(value) => Some(value),
606            ScalarValue::I32(value) => Some(value as f64),
607            ScalarValue::I64(value) => Some(value as f64),
608            ScalarValue::Bool(value) => Some(if value { 1.0 } else { 0.0 }),
609            ScalarValue::C32(_) | ScalarValue::C64(_) => None,
610        }
611    }
612
613    /// Returns the underlying value as `Complex64` when complex.
614    ///
615    /// Returns `None` for real scalars.
616    ///
617    /// # Examples
618    ///
619    /// ```
620    /// use tensor4all_tensorbackend::BackendScalar;
621    /// use num_complex::Complex64;
622    ///
623    /// let c = BackendScalar::new_complex(1.0, 2.0);
624    /// assert_eq!(c.as_c64(), Some(Complex64::new(1.0, 2.0)));
625    ///
626    /// let r = BackendScalar::new_real(1.0);
627    /// assert_eq!(r.as_c64(), None);
628    /// ```
629    pub fn as_c64(&self) -> Option<Complex64> {
630        match self.value() {
631            ScalarValue::F32(_)
632            | ScalarValue::F64(_)
633            | ScalarValue::I32(_)
634            | ScalarValue::I64(_)
635            | ScalarValue::Bool(_) => None,
636            ScalarValue::C32(value) => Some(Complex64::new(value.re as f64, value.im as f64)),
637            ScalarValue::C64(value) => Some(value),
638        }
639    }
640
641    /// Returns the complex conjugate.
642    ///
643    /// For real scalars, returns a copy (conjugate of a real number is itself).
644    ///
645    /// # Examples
646    ///
647    /// ```
648    /// use tensor4all_tensorbackend::BackendScalar;
649    ///
650    /// let c = BackendScalar::new_complex(3.0, 4.0);
651    /// let cc = c.conj();
652    /// assert!((cc.real() - 3.0).abs() < 1e-10);
653    /// assert!((cc.imag() - (-4.0)).abs() < 1e-10);
654    ///
655    /// let r = BackendScalar::new_real(5.0);
656    /// assert!((r.conj().real() - 5.0).abs() < 1e-10);
657    /// ```
658    pub fn conj(&self) -> Self {
659        match self.value() {
660            ScalarValue::F32(value) => Self::from_value(value),
661            ScalarValue::F64(value) => Self::from_value(value),
662            ScalarValue::I32(value) => Self::from_i32(value),
663            ScalarValue::I64(value) => Self::from_i64(value),
664            ScalarValue::Bool(value) => Self::from_bool(value),
665            ScalarValue::C32(value) => Self::from_value(value.conj()),
666            ScalarValue::C64(value) => Self::from_value(value.conj()),
667        }
668    }
669
670    /// Returns the real part as a scalar, preserving scalar semantics.
671    ///
672    /// Unlike [`real`](Self::real), this returns an `BackendScalar` rather than raw `f64`.
673    ///
674    /// # Examples
675    ///
676    /// ```
677    /// use tensor4all_tensorbackend::BackendScalar;
678    ///
679    /// let c = BackendScalar::new_complex(3.0, 4.0);
680    /// let re = c.real_part();
681    /// assert!(re.is_real());
682    /// assert!((re.real() - 3.0).abs() < 1e-10);
683    /// ```
684    pub fn real_part(&self) -> Self {
685        Self::from_real(self.real())
686    }
687
688    /// Returns the imaginary part as a scalar, preserving scalar semantics.
689    ///
690    /// Unlike [`imag`](Self::imag), this returns an `BackendScalar` rather than raw `f64`.
691    /// The result is always a real scalar.
692    ///
693    /// # Examples
694    ///
695    /// ```
696    /// use tensor4all_tensorbackend::BackendScalar;
697    ///
698    /// let c = BackendScalar::new_complex(3.0, 4.0);
699    /// let im = c.imag_part();
700    /// assert!(im.is_real());
701    /// assert!((im.real() - 4.0).abs() < 1e-10);
702    /// ```
703    pub fn imag_part(&self) -> Self {
704        Self::from_real(self.imag())
705    }
706
707    /// Compose a complex scalar from real-valued parts.
708    ///
709    /// # Errors
710    ///
711    /// Returns an error when either input is not a real scalar (a dtype
712    /// mismatch failure).
713    ///
714    /// # Examples
715    ///
716    /// ```
717    /// use tensor4all_tensorbackend::BackendScalar;
718    ///
719    /// let re = BackendScalar::new_real(3.0);
720    /// let im = BackendScalar::new_real(4.0);
721    /// let c = BackendScalar::compose_complex(re, im).unwrap();
722    /// assert!(c.is_complex());
723    /// assert!((c.real() - 3.0).abs() < 1e-10);
724    /// assert!((c.imag() - 4.0).abs() < 1e-10);
725    /// ```
726    pub fn compose_complex(real: Self, imag: Self) -> Result<Self> {
727        if !real.is_real() || !imag.is_real() {
728            return Err(anyhow!(
729                "compose_complex requires real-valued inputs, got real={:?}, imag={:?}",
730                real.native.dtype(),
731                imag.native.dtype()
732            ));
733        }
734        Ok(Self::from_complex(real.real(), imag.real()))
735    }
736
737    /// Square root, preserving AD metadata.
738    ///
739    /// Automatically promotes to complex if the value is negative.
740    ///
741    /// # Examples
742    ///
743    /// ```
744    /// use tensor4all_tensorbackend::BackendScalar;
745    ///
746    /// let s = BackendScalar::new_real(9.0);
747    /// assert!((s.sqrt().real() - 3.0).abs() < 1e-10);
748    /// ```
749    pub fn sqrt(&self) -> Self {
750        if self.is_complex() || self.real() < 0.0 {
751            let value = self.value().into_complex().sqrt();
752            if value.im == 0.0 {
753                Self::from_real(value.re)
754            } else {
755                Self::from_value(value)
756            }
757        } else {
758            Self::from_real(self.real().sqrt())
759        }
760    }
761
762    /// Real exponent power, preserving AD metadata.
763    ///
764    /// Automatically promotes to complex when the base is negative and the
765    /// exponent is non-integer.
766    ///
767    /// # Examples
768    ///
769    /// ```
770    /// use tensor4all_tensorbackend::BackendScalar;
771    ///
772    /// let s = BackendScalar::new_real(2.0);
773    /// assert!((s.powf(3.0).real() - 8.0).abs() < 1e-10);
774    /// ```
775    pub fn powf(&self, exponent: f64) -> Self {
776        let needs_complex_promotion =
777            self.is_complex() || (self.real() < 0.0 && exponent.fract() != 0.0);
778        if needs_complex_promotion {
779            let value = self.value().into_complex().powf(exponent);
780            if value.im == 0.0 {
781                Self::from_real(value.re)
782            } else {
783                Self::from_value(value)
784            }
785        } else {
786            Self::from_real(self.real().powf(exponent))
787        }
788    }
789
790    /// Integer exponent power, preserving AD metadata.
791    ///
792    /// # Examples
793    ///
794    /// ```
795    /// use tensor4all_tensorbackend::BackendScalar;
796    ///
797    /// let s = BackendScalar::new_real(3.0);
798    /// assert!((s.powi(2).real() - 9.0).abs() < 1e-10);
799    /// ```
800    pub fn powi(&self, exponent: i32) -> Self {
801        self.powf(exponent as f64)
802    }
803}
804
805impl SumFromStorage for BackendScalar {
806    fn sum_from_storage(storage: &Storage) -> Self {
807        match scalar_value_from_storage(storage) {
808            ScalarValue::F32(value) => Self::from_value(value),
809            ScalarValue::F64(value) => Self::from_value(value),
810            ScalarValue::I32(value) => Self::from_i32(value),
811            ScalarValue::I64(value) => Self::from_i64(value),
812            ScalarValue::Bool(value) => Self::from_bool(value),
813            ScalarValue::C32(value) => Self::from_value(value),
814            ScalarValue::C64(value) => Self::from_value(value),
815        }
816    }
817}
818
819impl From<f32> for BackendScalar {
820    fn from(value: f32) -> Self {
821        Self::from_value(value)
822    }
823}
824
825impl From<f64> for BackendScalar {
826    fn from(value: f64) -> Self {
827        Self::from_value(value)
828    }
829}
830
831impl From<Complex32> for BackendScalar {
832    fn from(value: Complex32) -> Self {
833        Self::from_value(value)
834    }
835}
836
837impl From<Complex64> for BackendScalar {
838    fn from(value: Complex64) -> Self {
839        Self::from_value(value)
840    }
841}
842
843impl TryFrom<BackendScalar> for f64 {
844    type Error = &'static str;
845
846    fn try_from(value: BackendScalar) -> std::result::Result<Self, Self::Error> {
847        match value.value() {
848            ScalarValue::F32(real) => Ok(real as f64),
849            ScalarValue::F64(real) => Ok(real),
850            ScalarValue::I32(real) => Ok(real as f64),
851            ScalarValue::I64(real) => Ok(real as f64),
852            ScalarValue::Bool(real) => Ok(if real { 1.0 } else { 0.0 }),
853            ScalarValue::C32(_) | ScalarValue::C64(_) => {
854                Err("cannot convert complex scalar to f64")
855            }
856        }
857    }
858}
859
860impl From<BackendScalar> for Complex64 {
861    fn from(value: BackendScalar) -> Self {
862        value.value().into_complex()
863    }
864}
865
866impl Add for BackendScalar {
867    type Output = Self;
868
869    fn add(self, rhs: Self) -> Self::Output {
870        match (self.value(), rhs.value()) {
871            (ScalarValue::F32(lhs), ScalarValue::F32(rhs)) => Self::from_value(lhs + rhs),
872            (lhs, rhs) if lhs.is_complex() || rhs.is_complex() => {
873                Self::from_value(lhs.into_complex() + rhs.into_complex())
874            }
875            (lhs, rhs) => Self::from_real(lhs.real() + rhs.real()),
876        }
877    }
878}
879
880impl Sub for BackendScalar {
881    type Output = Self;
882
883    fn sub(self, rhs: Self) -> Self::Output {
884        self + (-rhs)
885    }
886}
887
888impl Mul for BackendScalar {
889    type Output = Self;
890
891    fn mul(self, rhs: Self) -> Self::Output {
892        match (self.value(), rhs.value()) {
893            (ScalarValue::F32(lhs), ScalarValue::F32(rhs)) => Self::from_value(lhs * rhs),
894            (lhs, rhs) if lhs.is_complex() || rhs.is_complex() => {
895                Self::from_value(lhs.into_complex() * rhs.into_complex())
896            }
897            (lhs, rhs) => Self::from_real(lhs.real() * rhs.real()),
898        }
899    }
900}
901
902impl Div for BackendScalar {
903    type Output = Self;
904
905    fn div(self, rhs: Self) -> Self::Output {
906        match (self.value(), rhs.value()) {
907            (ScalarValue::F32(lhs), ScalarValue::F32(rhs)) => Self::from_value(lhs / rhs),
908            (lhs, rhs) if lhs.is_complex() || rhs.is_complex() => {
909                Self::from_value(lhs.into_complex() / rhs.into_complex())
910            }
911            (lhs, rhs) => Self::from_real(lhs.real() / rhs.real()),
912        }
913    }
914}
915
916impl Neg for BackendScalar {
917    type Output = Self;
918
919    fn neg(self) -> Self::Output {
920        match self.value() {
921            ScalarValue::F32(value) => Self::from_value(-value),
922            ScalarValue::F64(value) => Self::from_value(-value),
923            ScalarValue::I32(value) => Self::from_i32(-value),
924            ScalarValue::I64(value) => Self::from_i64(-value),
925            ScalarValue::Bool(value) => Self::from_real(if value { -1.0 } else { 0.0 }),
926            ScalarValue::C32(value) => Self::from_value(-value),
927            ScalarValue::C64(value) => Self::from_value(-value),
928        }
929    }
930}
931
932impl Mul<BackendScalar> for f64 {
933    type Output = BackendScalar;
934
935    fn mul(self, rhs: BackendScalar) -> Self::Output {
936        BackendScalar::from_real(self) * rhs
937    }
938}
939
940impl Mul<BackendScalar> for Complex64 {
941    type Output = BackendScalar;
942
943    fn mul(self, rhs: BackendScalar) -> Self::Output {
944        BackendScalar::from(self) * rhs
945    }
946}
947
948impl Div<BackendScalar> for Complex64 {
949    type Output = BackendScalar;
950
951    fn div(self, rhs: BackendScalar) -> Self::Output {
952        BackendScalar::from(self) / rhs
953    }
954}
955
956impl Default for BackendScalar {
957    fn default() -> Self {
958        Self::zero()
959    }
960}
961
962impl Zero for BackendScalar {
963    fn zero() -> Self {
964        Self::from_real(0.0)
965    }
966
967    fn is_zero(&self) -> bool {
968        BackendScalar::is_zero(self)
969    }
970}
971
972impl One for BackendScalar {
973    fn one() -> Self {
974        Self::from_real(1.0)
975    }
976}
977
978impl PartialEq for BackendScalar {
979    fn eq(&self, other: &Self) -> bool {
980        self.native.dtype() == other.native.dtype() && self.value() == other.value()
981    }
982}
983
984impl PartialOrd for BackendScalar {
985    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
986        match (self.value(), other.value()) {
987            (ScalarValue::F32(lhs), ScalarValue::F32(rhs)) => lhs.partial_cmp(&rhs),
988            (ScalarValue::F32(lhs), ScalarValue::F64(rhs)) => (lhs as f64).partial_cmp(&rhs),
989            (ScalarValue::F32(lhs), ScalarValue::I32(rhs)) => {
990                (lhs as f64).partial_cmp(&(rhs as f64))
991            }
992            (ScalarValue::F32(lhs), ScalarValue::I64(rhs)) => {
993                (lhs as f64).partial_cmp(&(rhs as f64))
994            }
995            (ScalarValue::F32(lhs), ScalarValue::Bool(rhs)) => {
996                (lhs as f64).partial_cmp(&(if rhs { 1.0 } else { 0.0 }))
997            }
998            (ScalarValue::F64(lhs), ScalarValue::F32(rhs)) => lhs.partial_cmp(&(rhs as f64)),
999            (ScalarValue::F64(lhs), ScalarValue::F64(rhs)) => lhs.partial_cmp(&rhs),
1000            (ScalarValue::F64(lhs), ScalarValue::I32(rhs)) => lhs.partial_cmp(&(rhs as f64)),
1001            (ScalarValue::F64(lhs), ScalarValue::I64(rhs)) => lhs.partial_cmp(&(rhs as f64)),
1002            (ScalarValue::F64(lhs), ScalarValue::Bool(rhs)) => {
1003                lhs.partial_cmp(&(if rhs { 1.0 } else { 0.0 }))
1004            }
1005            (ScalarValue::I32(lhs), ScalarValue::F32(rhs)) => {
1006                (lhs as f64).partial_cmp(&(rhs as f64))
1007            }
1008            (ScalarValue::I32(lhs), ScalarValue::F64(rhs)) => (lhs as f64).partial_cmp(&rhs),
1009            (ScalarValue::I32(lhs), ScalarValue::I32(rhs)) => lhs.partial_cmp(&rhs),
1010            (ScalarValue::I32(lhs), ScalarValue::I64(rhs)) => {
1011                (lhs as f64).partial_cmp(&(rhs as f64))
1012            }
1013            (ScalarValue::I32(lhs), ScalarValue::Bool(rhs)) => {
1014                (lhs as f64).partial_cmp(&(if rhs { 1.0 } else { 0.0 }))
1015            }
1016            (ScalarValue::I64(lhs), ScalarValue::F32(rhs)) => {
1017                (lhs as f64).partial_cmp(&(rhs as f64))
1018            }
1019            (ScalarValue::I64(lhs), ScalarValue::F64(rhs)) => (lhs as f64).partial_cmp(&rhs),
1020            (ScalarValue::I64(lhs), ScalarValue::I32(rhs)) => {
1021                (lhs as f64).partial_cmp(&(rhs as f64))
1022            }
1023            (ScalarValue::I64(lhs), ScalarValue::I64(rhs)) => lhs.partial_cmp(&rhs),
1024            (ScalarValue::I64(lhs), ScalarValue::Bool(rhs)) => {
1025                (lhs as f64).partial_cmp(&(if rhs { 1.0 } else { 0.0 }))
1026            }
1027            (ScalarValue::Bool(lhs), ScalarValue::F32(rhs)) => {
1028                (if lhs { 1.0 } else { 0.0 }).partial_cmp(&(rhs as f64))
1029            }
1030            (ScalarValue::Bool(lhs), ScalarValue::F64(rhs)) => {
1031                (if lhs { 1.0 } else { 0.0 }).partial_cmp(&rhs)
1032            }
1033            (ScalarValue::Bool(lhs), ScalarValue::I32(rhs)) => {
1034                (if lhs { 1.0 } else { 0.0 }).partial_cmp(&(rhs as f64))
1035            }
1036            (ScalarValue::Bool(lhs), ScalarValue::I64(rhs)) => {
1037                (if lhs { 1.0 } else { 0.0 }).partial_cmp(&(rhs as f64))
1038            }
1039            (ScalarValue::Bool(lhs), ScalarValue::Bool(rhs)) => lhs.partial_cmp(&rhs),
1040            _ => None,
1041        }
1042    }
1043}
1044
1045impl fmt::Display for BackendScalar {
1046    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1047        match self.value() {
1048            ScalarValue::F32(value) => value.fmt(f),
1049            ScalarValue::F64(value) => value.fmt(f),
1050            ScalarValue::I32(value) => value.fmt(f),
1051            ScalarValue::I64(value) => value.fmt(f),
1052            ScalarValue::Bool(value) => value.fmt(f),
1053            ScalarValue::C32(value) => value.fmt(f),
1054            ScalarValue::C64(value) => value.fmt(f),
1055        }
1056    }
1057}
1058
1059impl fmt::Debug for BackendScalar {
1060    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1061        f.debug_struct("BackendScalar")
1062            .field("dtype", &self.native.dtype())
1063            .field("value", &self.value())
1064            .finish()
1065    }
1066}
1067
1068impl Clone for BackendScalar {
1069    fn clone(&self) -> Self {
1070        match self.value {
1071            ScalarValue::F32(value) => Self::from_value(value),
1072            ScalarValue::F64(value) => Self::from_value(value),
1073            ScalarValue::I32(value) => Self::from_i32(value),
1074            ScalarValue::I64(value) => Self::from_i64(value),
1075            ScalarValue::Bool(value) => Self::from_bool(value),
1076            ScalarValue::C32(value) => Self::from_value(value),
1077            ScalarValue::C64(value) => Self::from_value(value),
1078        }
1079    }
1080}
1081
1082#[cfg(test)]
1083mod tests;
tensor4all_tensorbackend/any_scalar.rs

tensor4all_tensorbackend/
any_scalar.rs
