1#[cfg(test)]
2use std::cell::Cell;
3use std::cmp::Ordering;
4use std::fmt;
5use std::ops::{Add, Div, Mul, Neg, Sub};
6use std::sync::Arc;
7
8use anyhow::{anyhow, Result};
9use num_complex::{Complex32, Complex64};
10use num_traits::{One, Zero};
11use tenferro::DType;
12use tensor4all_tensorbackend::BackendScalar;
13
14use crate::defaults::eager_session::{
15 with_eager_session_typed, with_engine, with_execution_engine,
16};
17use crate::defaults::idx_tensor::IdxTensor;
18use crate::TensorElement;
19use tensor4all_tensorbackend::{Storage, SumFromStorage};
20
21#[derive(Clone, Copy, Debug, PartialEq)]
22pub(crate) enum ScalarValue {
23 F32(f32),
24 F64(f64),
25 C32(Complex32),
26 C64(Complex64),
27}
28
29impl From<f64> for ScalarValue {
30 fn from(value: f64) -> Self {
31 Self::F64(value)
32 }
33}
34
35impl From<f32> for ScalarValue {
36 fn from(value: f32) -> Self {
37 Self::F32(value)
38 }
39}
40
41impl From<Complex64> for ScalarValue {
42 fn from(value: Complex64) -> Self {
43 Self::C64(value)
44 }
45}
46
47impl From<Complex32> for ScalarValue {
48 fn from(value: Complex32) -> Self {
49 Self::C32(value)
50 }
51}
52
53impl ScalarValue {
54 fn real(self) -> f64 {
55 match self {
56 Self::F32(value) => value as f64,
57 Self::F64(value) => value,
58 Self::C32(value) => value.re as f64,
59 Self::C64(value) => value.re,
60 }
61 }
62
63 fn imag(self) -> f64 {
64 match self {
65 Self::F32(_) | Self::F64(_) => 0.0,
66 Self::C32(value) => value.im as f64,
67 Self::C64(value) => value.im,
68 }
69 }
70
71 fn abs(self) -> f64 {
72 match self {
73 Self::F32(value) => value.abs() as f64,
74 Self::F64(value) => value.abs(),
75 Self::C32(value) => {
76 if value.re.is_nan() || value.im.is_nan() {
77 f64::NAN
78 } else {
79 f64::from(value.re).hypot(f64::from(value.im))
80 }
81 }
82 Self::C64(value) => {
83 if value.re.is_nan() || value.im.is_nan() {
84 f64::NAN
85 } else {
86 value.re.hypot(value.im)
87 }
88 }
89 }
90 }
91
92 fn is_complex(self) -> bool {
93 matches!(self, Self::C32(_) | Self::C64(_))
94 }
95
96 fn is_zero(self) -> bool {
97 match self {
98 Self::F32(value) => value == 0.0,
99 Self::F64(value) => value == 0.0,
100 Self::C32(value) => value == Complex32::new(0.0, 0.0),
101 Self::C64(value) => value == Complex64::new(0.0, 0.0),
102 }
103 }
104
105 fn into_complex(self) -> Complex64 {
106 match self {
107 Self::F32(value) => Complex64::new(value as f64, 0.0),
108 Self::F64(value) => Complex64::new(value, 0.0),
109 Self::C32(value) => Complex64::new(value.re as f64, value.im as f64),
110 Self::C64(value) => value,
111 }
112 }
113}
114
115trait ScalarTensorElement: TensorElement {
116 fn scalar_value(value: Self) -> ScalarValue;
117}
118
119impl ScalarTensorElement for f32 {
120 fn scalar_value(value: Self) -> ScalarValue {
121 ScalarValue::F32(value)
122 }
123}
124
125impl ScalarTensorElement for f64 {
126 fn scalar_value(value: Self) -> ScalarValue {
127 ScalarValue::F64(value)
128 }
129}
130
131impl ScalarTensorElement for Complex32 {
132 fn scalar_value(value: Self) -> ScalarValue {
133 ScalarValue::C32(value)
134 }
135}
136
137impl ScalarTensorElement for Complex64 {
138 fn scalar_value(value: Self) -> ScalarValue {
139 ScalarValue::C64(value)
140 }
141}
142
143#[cfg(test)]
144thread_local! {
145 static FORCE_ANY_SCALAR_TENSOR_INITIALIZATION_FAILURE: Cell<bool> = const { Cell::new(false) };
146}
147
148#[derive(Debug, Clone, thiserror::Error)]
149enum AnyScalarTensorError {
150 #[error("AnyScalar tensor initialization failed: {source}")]
151 Initialization {
152 #[source]
153 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
154 },
155 #[error("AnyScalar::{op} failed: {source}")]
156 Operation {
157 op: &'static str,
158 #[source]
159 source: Arc<dyn std::error::Error + Send + Sync + 'static>,
160 },
161}
162
163fn initialize_tensor<T: ScalarTensorElement>(
164 value: T,
165) -> std::result::Result<IdxTensor, AnyScalarTensorError> {
166 #[cfg(test)]
167 if FORCE_ANY_SCALAR_TENSOR_INITIALIZATION_FAILURE.with(Cell::get) {
168 return Err(AnyScalarTensorError::Initialization {
169 source: Arc::new(std::io::Error::other(
170 "forced AnyScalar eager initialization failure",
171 )),
172 });
173 }
174
175 IdxTensor::scalar(value).map_err(|source| AnyScalarTensorError::Initialization {
176 source: Arc::from(anyhow::Error::new(source).into_boxed_dyn_error()),
177 })
178}
179
180#[derive(Debug, thiserror::Error)]
198#[error("AnyScalar eager-tensor operation failed: {source}")]
199pub struct AnyScalarError {
200 #[source]
203 pub source: anyhow::Error,
204}
205
206impl From<anyhow::Error> for AnyScalarError {
207 fn from(source: anyhow::Error) -> Self {
208 Self { source }
209 }
210}
211
212fn operation_error<E>(op: &'static str, source: E) -> anyhow::Error
213where
214 E: std::error::Error + Send + Sync + 'static,
215{
216 anyhow::Error::new(AnyScalarTensorError::Operation {
217 op,
218 source: Arc::new(source),
219 })
220}
221
222fn operation_error_from_anyhow(op: &'static str, source: anyhow::Error) -> anyhow::Error {
223 match source.downcast::<AnyScalarTensorError>() {
224 Ok(source) => anyhow::Error::new(source),
225 Err(source) => anyhow::Error::new(AnyScalarTensorError::Operation {
226 op,
227 source: Arc::from(source.into_boxed_dyn_error()),
228 }),
229 }
230}
231
232#[derive(Clone)]
240pub struct AnyScalar {
241 tensor: std::result::Result<IdxTensor, AnyScalarTensorError>,
242 value: ScalarValue,
243 tracks_grad: bool,
244}
245
246impl AnyScalar {
247 fn wrap_tensor(tensor: IdxTensor) -> Result<Self> {
248 let dims = tensor.dims();
249 anyhow::ensure!(
250 dims.is_empty(),
251 "AnyScalar requires a rank-0 tensor, got dims {:?}",
252 dims
253 );
254 let value = Self::scalar_value_from_tensor(&tensor)?;
255 let tracks_grad = tensor.tracks_grad();
256 Ok(Self {
257 tensor: Ok(tensor),
258 value,
259 tracks_grad,
260 })
261 }
262
263 fn from_tensor_result(tensor: Result<IdxTensor>, op: &'static str) -> Result<Self> {
264 let tensor = tensor.map_err(|error| operation_error_from_anyhow(op, error))?;
265 Self::wrap_tensor(tensor).map_err(|error| operation_error_from_anyhow(op, error))
266 }
267
268 fn fallback_result(
269 result: Result<Self>,
270 op: &'static str,
271 fallback: impl FnOnce() -> ScalarValue,
272 tracks_grad: bool,
273 ) -> Self {
274 match result {
275 Ok(result) => result,
276 Err(error) => {
277 let error = match error.downcast::<AnyScalarTensorError>() {
278 Ok(error) => error,
279 Err(error) => AnyScalarTensorError::Operation {
280 op,
281 source: Arc::from(error.into_boxed_dyn_error()),
282 },
283 };
284 Self {
285 tensor: Err(error),
286 value: fallback(),
287 tracks_grad,
288 }
289 }
290 }
291 }
292
293 fn scalar_value_from_backend(value: BackendScalar) -> ScalarValue {
294 value
295 .as_c64()
296 .map(ScalarValue::C64)
297 .unwrap_or_else(|| ScalarValue::F64(value.real()))
298 }
299
300 fn zero_like(&self) -> Self {
301 match self.value() {
302 ScalarValue::F32(_) => Self::from_value(0.0_f32),
303 ScalarValue::F64(_) => Self::from_value(0.0_f64),
304 ScalarValue::C32(_) => Self::from_value(Complex32::new(0.0, 0.0)),
305 ScalarValue::C64(_) => Self::from_value(Complex64::new(0.0, 0.0)),
306 }
307 }
308
309 fn one_like(&self) -> Self {
310 match self.value() {
311 ScalarValue::F32(_) => Self::from_value(1.0_f32),
312 ScalarValue::F64(_) => Self::from_value(1.0_f64),
313 ScalarValue::C32(_) => Self::from_value(Complex32::new(1.0, 0.0)),
314 ScalarValue::C64(_) => Self::from_value(Complex64::new(1.0, 0.0)),
315 }
316 }
317
318 fn from_eager_binary<E>(
319 lhs: &Self,
320 rhs: &Self,
321 op: &'static str,
322 f: impl FnOnce(
323 &mut tenferro_ad::EagerSession<'_>,
324 &tenferro_ad::EagerTensor,
325 &tenferro_ad::EagerTensor,
326 ) -> std::result::Result<tenferro_ad::EagerTensor, E>,
327 ) -> Result<Self>
328 where
329 E: From<tenferro_ad::Error> + std::error::Error + Send + Sync + 'static,
330 {
331 let lhs_inner = lhs.as_tensor()?.as_inner()?;
332 let rhs_inner = rhs.as_tensor()?.as_inner()?;
333 let result = with_eager_session_typed(lhs_inner.runtime(), None, |session| {
334 f(session, lhs_inner, rhs_inner)
335 })
336 .map_err(|error| operation_error(op, error))?
337 .map_err(|error| operation_error(op, error))?;
338 Self::from_tensor_result(IdxTensor::from_inner(vec![], result), op)
339 }
340
341 fn from_eager_unary<E>(
342 input: &Self,
343 op: &'static str,
344 f: impl FnOnce(
345 &mut tenferro_ad::EagerSession<'_>,
346 &tenferro_ad::EagerTensor,
347 ) -> std::result::Result<tenferro_ad::EagerTensor, E>,
348 ) -> Result<Self>
349 where
350 E: From<tenferro_ad::Error> + std::error::Error + Send + Sync + 'static,
351 {
352 let inner = input.as_tensor()?.as_inner()?;
353 let result = with_eager_session_typed(inner.runtime(), None, |session| f(session, inner))
354 .map_err(|error| operation_error(op, error))?
355 .map_err(|error| operation_error(op, error))?;
356 Self::from_tensor_result(IdxTensor::from_inner(vec![], result), op)
357 }
358
359 fn scalar_value_from_tensor(tensor: &IdxTensor) -> Result<ScalarValue> {
360 let inner = tensor.as_inner()?;
361 match inner.dtype() {
362 DType::F32 => inner
363 .value()?
364 .as_slice::<f32>()?
365 .first()
366 .copied()
367 .map(ScalarValue::F32)
368 .ok_or_else(|| anyhow!("rank-0 f32 scalar tensor is empty")),
369 DType::F64 => inner
370 .value()?
371 .as_slice::<f64>()?
372 .first()
373 .copied()
374 .map(ScalarValue::F64)
375 .ok_or_else(|| anyhow!("rank-0 f64 scalar tensor is empty")),
376 DType::C32 => inner
377 .value()?
378 .as_slice::<Complex32>()?
379 .first()
380 .copied()
381 .map(ScalarValue::C32)
382 .ok_or_else(|| anyhow!("rank-0 c32 scalar tensor is empty")),
383 DType::C64 => inner
384 .value()?
385 .as_slice::<Complex64>()?
386 .first()
387 .copied()
388 .map(ScalarValue::C64)
389 .ok_or_else(|| anyhow!("rank-0 c64 scalar tensor is empty")),
390 dtype => Err(anyhow!("unsupported scalar tensor dtype {dtype:?}")),
391 }
392 }
393
394 fn value(&self) -> ScalarValue {
395 self.value
396 }
397
398 fn from_backend_scalar(value: BackendScalar) -> Self {
399 match Self::scalar_value_from_backend(value) {
400 ScalarValue::F32(value) => Self::from_value(value),
401 ScalarValue::F64(value) => Self::from_value(value),
402 ScalarValue::C32(value) => Self::from_value(value),
403 ScalarValue::C64(value) => Self::from_value(value),
404 }
405 }
406
407 pub(crate) fn from_tensor(tensor: IdxTensor) -> Result<Self> {
408 Self::wrap_tensor(tensor)
409 }
410
411 pub(crate) fn from_published<T: Into<ScalarValue>>(value: T, tensor: IdxTensor) -> Self {
418 Self {
419 tensor: Ok(tensor),
420 value: value.into(),
421 tracks_grad: false,
422 }
423 }
424
425 pub(crate) fn as_tensor(&self) -> Result<&IdxTensor> {
426 self.tensor
427 .as_ref()
428 .map_err(|error| anyhow::Error::new(error.clone()))
429 }
430
431 #[allow(private_bounds)]
461 pub fn from_value<T: ScalarTensorElement>(value: T) -> Self {
462 Self {
463 tensor: initialize_tensor(value),
464 value: T::scalar_value(value),
465 tracks_grad: false,
466 }
467 }
468
469 pub fn from_real(x: f64) -> Self {
491 Self::from_value(x)
492 }
493
494 pub fn from_complex(re: f64, im: f64) -> Self {
517 Self::from_value(Complex64::new(re, im))
518 }
519
520 pub fn new_real(x: f64) -> Self {
542 Self::from_real(x)
543 }
544
545 pub fn new_complex(re: f64, im: f64) -> Self {
568 Self::from_complex(re, im)
569 }
570
571 pub fn primal(&self) -> std::result::Result<Self, AnyScalarError> {
594 self.detach()
595 }
596
597 pub fn enable_grad(self) -> std::result::Result<Self, AnyScalarError> {
618 let tensor = self.tensor.map_err(anyhow::Error::new)?;
619 Self::from_tensor(tensor.enable_grad().map_err(anyhow::Error::from)?)
620 .map_err(AnyScalarError::from)
621 }
622
623 pub fn tracks_grad(&self) -> bool {
639 self.tracks_grad || self.tensor.as_ref().is_ok_and(IdxTensor::tracks_grad)
640 }
641
642 pub fn grad(&self) -> std::result::Result<Option<Self>, AnyScalarError> {
669 self.as_tensor()?
670 .grad()
671 .map_err(anyhow::Error::from)
672 .and_then(|maybe_grad| maybe_grad.map(Self::from_tensor).transpose())
673 .map_err(AnyScalarError::from)
674 }
675
676 pub fn clear_grad(&self) -> std::result::Result<(), AnyScalarError> {
702 self.as_tensor()?
703 .clear_grad()
704 .map_err(anyhow::Error::from)
705 .map_err(AnyScalarError::from)
706 }
707
708 pub fn backward(&self) -> std::result::Result<(), AnyScalarError> {
733 let inner = self.as_tensor()?.as_inner()?;
734 with_execution_engine(inner.runtime(), None, |engine| inner.backward(engine))
735 .map(|_| ())
736 .map_err(|error| AnyScalarError::from(anyhow::Error::new(error)))
737 }
738
739 pub fn detach(&self) -> std::result::Result<Self, AnyScalarError> {
764 let tensor = self.as_tensor()?;
765 let inner = tensor.as_inner()?;
766 let detached = with_engine(inner.runtime(), None, |engine| inner.detach(engine))
767 .map_err(|error| AnyScalarError::from(anyhow::Error::new(error)))?;
768 Self::from_tensor(IdxTensor::from_inner(vec![], detached)?).map_err(AnyScalarError::from)
769 }
770
771 pub fn real(&self) -> f64 {
787 self.value().real()
788 }
789
790 pub fn imag(&self) -> f64 {
805 self.value().imag()
806 }
807
808 pub fn abs(&self) -> f64 {
824 self.value().abs()
825 }
826
827 pub fn is_complex(&self) -> bool {
842 self.value().is_complex()
843 }
844
845 pub fn is_real(&self) -> bool {
860 !self.is_complex()
861 }
862
863 pub fn is_zero(&self) -> bool {
878 self.value().is_zero()
879 }
880
881 pub fn as_f64(&self) -> Option<f64> {
897 match self.value() {
898 ScalarValue::F32(value) => Some(value as f64),
899 ScalarValue::F64(value) => Some(value),
900 ScalarValue::C32(_) | ScalarValue::C64(_) => None,
901 }
902 }
903
904 pub fn as_c64(&self) -> Option<Complex64> {
921 match self.value() {
922 ScalarValue::F32(_) | ScalarValue::F64(_) => None,
923 ScalarValue::C32(value) => Some(Complex64::new(value.re as f64, value.im as f64)),
924 ScalarValue::C64(value) => Some(value),
925 }
926 }
927
928 pub fn try_conj(&self) -> std::result::Result<Self, AnyScalarError> {
950 self.as_tensor()?;
951 if !self.tracks_grad() {
952 return Ok(Self::from_backend_scalar(self.to_backend_scalar().conj()));
953 }
954 Self::from_eager_unary(self, "conj", |session, tensor| session.conj(tensor))
955 .map_err(AnyScalarError::from)
956 }
957
958 pub fn conj(&self) -> Self {
960 Self::fallback_result(
961 self.try_conj().map_err(|error| error.source),
962 "conj",
963 || Self::scalar_value_from_backend(self.to_backend_scalar().conj()),
964 self.tracks_grad(),
965 )
966 }
967
968 pub fn real_part(&self) -> Self {
984 Self::fallback_result(
985 self.try_real_part(),
986 "real_part",
987 || Self::from_real(self.real()).value(),
988 self.tracks_grad(),
989 )
990 }
991
992 pub fn imag_part(&self) -> Self {
1008 Self::fallback_result(
1009 self.try_imag_part(),
1010 "imag_part",
1011 || Self::from_real(self.imag()).value(),
1012 self.tracks_grad(),
1013 )
1014 }
1015
1016 pub fn compose_complex(real: Self, imag: Self) -> std::result::Result<Self, AnyScalarError> {
1047 if !real.is_real() || !imag.is_real() {
1048 return Err(anyhow!("compose_complex requires real-valued inputs").into());
1049 }
1050 let imag_term = imag.try_mul(&Self::new_complex(0.0, 1.0))?;
1051 real.try_add(&imag_term).map_err(AnyScalarError::from)
1052 }
1053
1054 pub fn sqrt(&self) -> Self {
1071 Self::fallback_result(
1072 self.try_sqrt(),
1073 "sqrt",
1074 || Self::scalar_value_from_backend(self.to_backend_scalar().sqrt()),
1075 self.tracks_grad(),
1076 )
1077 }
1078
1079 pub fn powf(&self, exponent: f64) -> Self {
1098 Self::fallback_result(
1099 self.try_powf(exponent),
1100 "powf",
1101 || Self::scalar_value_from_backend(self.to_backend_scalar().powf(exponent)),
1102 self.tracks_grad(),
1103 )
1104 }
1105
1106 pub fn powi(&self, exponent: i32) -> Self {
1127 Self::fallback_result(
1128 self.try_powi(exponent),
1129 "powi",
1130 || Self::scalar_value_from_backend(self.to_backend_scalar().powi(exponent)),
1131 self.tracks_grad(),
1132 )
1133 }
1134
1135 pub(crate) fn to_backend_scalar(&self) -> BackendScalar {
1136 match self.value() {
1137 ScalarValue::F32(value) => BackendScalar::from_value(value),
1138 ScalarValue::F64(value) => BackendScalar::from_value(value),
1139 ScalarValue::C32(value) => BackendScalar::from_value(value),
1140 ScalarValue::C64(value) => BackendScalar::from_value(value),
1141 }
1142 }
1143
1144 pub(crate) fn try_add(&self, rhs: &Self) -> Result<Self> {
1145 self.as_tensor()?;
1146 rhs.as_tensor()?;
1147 if !self.tracks_grad() && !rhs.tracks_grad() {
1148 return Ok(Self::from_backend_scalar(
1149 self.to_backend_scalar() + rhs.to_backend_scalar(),
1150 ));
1151 }
1152 Self::from_eager_binary(self, rhs, "add", |session, lhs, rhs| session.add(lhs, rhs))
1153 }
1154
1155 pub(crate) fn try_mul(&self, rhs: &Self) -> Result<Self> {
1156 self.as_tensor()?;
1157 rhs.as_tensor()?;
1158 if !self.tracks_grad() && !rhs.tracks_grad() {
1159 return Ok(Self::from_backend_scalar(
1160 self.to_backend_scalar() * rhs.to_backend_scalar(),
1161 ));
1162 }
1163 Self::from_eager_binary(self, rhs, "mul", |session, lhs, rhs| session.mul(lhs, rhs))
1164 }
1165
1166 pub(crate) fn try_div(&self, rhs: &Self) -> Result<Self> {
1167 self.as_tensor()?;
1168 rhs.as_tensor()?;
1169 if !self.tracks_grad() && !rhs.tracks_grad() {
1170 return Ok(Self::from_backend_scalar(
1171 self.to_backend_scalar() / rhs.to_backend_scalar(),
1172 ));
1173 }
1174 Self::from_eager_binary(self, rhs, "div", |session, lhs, rhs| session.div(lhs, rhs))
1175 }
1176
1177 pub(crate) fn try_neg(&self) -> Result<Self> {
1178 self.as_tensor()?;
1179 if !self.tracks_grad() {
1180 return Ok(Self::from_backend_scalar(-self.to_backend_scalar()));
1181 }
1182 Self::from_eager_unary(self, "neg", |session, tensor| session.neg(tensor))
1183 }
1184
1185 fn try_real_part(&self) -> Result<Self> {
1186 self.as_tensor()?;
1187 if !self.tracks_grad() {
1188 return Ok(Self::from_real(self.real()));
1189 }
1190 if self.is_complex() {
1191 Self::from_eager_unary(self, "real_part", |session, tensor| {
1192 session.cast(tensor, DType::F64)
1193 })
1194 } else {
1195 self.try_mul(&Self::new_real(1.0))
1196 }
1197 }
1198
1199 fn try_imag_part(&self) -> Result<Self> {
1200 self.as_tensor()?;
1201 if !self.tracks_grad() {
1202 return Ok(Self::from_real(self.imag()));
1203 }
1204 if self.is_complex() {
1205 let factor = Self::new_complex(0.0, -1.0);
1206 let imaginary =
1207 Self::from_eager_binary(self, &factor, "imag_part", |session, value, factor| {
1208 session.mul(value, factor)
1209 })?;
1210 Self::from_eager_unary(&imaginary, "imag_part", |session, tensor| {
1211 session.cast(tensor, DType::F64)
1212 })
1213 } else {
1214 self.try_mul(&Self::new_real(0.0))
1215 }
1216 }
1217
1218 fn try_sqrt(&self) -> Result<Self> {
1219 self.as_tensor()?;
1220 if !self.tracks_grad() {
1221 return Ok(Self::from_backend_scalar(self.to_backend_scalar().sqrt()));
1222 }
1223 if self.is_real() && self.real() < 0.0 {
1224 let magnitude_input =
1225 Self::from_eager_unary(self, "sqrt", |session, tensor| session.neg(tensor))?;
1226 let magnitude = magnitude_input.try_sqrt()?;
1227 let factor = Self::new_complex(0.0, 1.0);
1228 return Self::from_eager_binary(
1229 &magnitude,
1230 &factor,
1231 "sqrt",
1232 |session, value, factor| session.mul(value, factor),
1233 );
1234 }
1235 Self::from_eager_unary(self, "sqrt", |session, tensor| session.sqrt(tensor))
1236 }
1237
1238 fn try_powf(&self, exponent: f64) -> Result<Self> {
1239 self.as_tensor()?;
1240 if !self.tracks_grad() {
1241 return Ok(Self::from_backend_scalar(
1242 self.to_backend_scalar().powf(exponent),
1243 ));
1244 }
1245 if self.is_real() && self.real() < 0.0 && exponent.fract() != 0.0 {
1246 let magnitude_input =
1247 Self::from_eager_unary(self, "powf", |session, tensor| session.neg(tensor))?;
1248 let magnitude = magnitude_input.try_powf(exponent)?;
1249 let phase = std::f64::consts::PI * exponent;
1250 let factor = Self::new_complex(phase.cos(), phase.sin());
1251 return Self::from_eager_binary(
1252 &magnitude,
1253 &factor,
1254 "powf",
1255 |session, value, factor| session.mul(value, factor),
1256 );
1257 }
1258 let exponent = if self.is_complex() {
1259 Self::new_complex(exponent, 0.0)
1260 } else {
1261 Self::new_real(exponent)
1262 };
1263 Self::from_eager_binary(self, &exponent, "powf", |session, base, exponent| {
1264 session.pow(base, exponent)
1265 })
1266 }
1267
1268 fn try_powi(&self, exponent: i32) -> Result<Self> {
1269 self.as_tensor()?;
1270 if exponent == 0 {
1271 if self.tracks_grad() {
1272 let zeroed = self.try_mul(&self.zero_like())?;
1276 return zeroed.try_add(&self.one_like());
1277 }
1278 return Ok(Self::one());
1279 }
1280 if self.tracks_grad() {
1281 return self.try_powf(exponent as f64);
1282 }
1283 Ok(Self::from_backend_scalar(
1284 self.to_backend_scalar().powi(exponent),
1285 ))
1286 }
1287}
1288
1289impl SumFromStorage for AnyScalar {
1290 fn sum_from_storage(storage: &Storage) -> Self {
1291 Self::from_backend_scalar(BackendScalar::sum_from_storage(storage))
1292 }
1293}
1294
1295impl From<f32> for AnyScalar {
1296 fn from(value: f32) -> Self {
1297 Self::from_value(value)
1298 }
1299}
1300
1301impl From<f64> for AnyScalar {
1302 fn from(value: f64) -> Self {
1303 Self::from_value(value)
1304 }
1305}
1306
1307impl From<Complex32> for AnyScalar {
1308 fn from(value: Complex32) -> Self {
1309 Self::from_value(value)
1310 }
1311}
1312
1313impl From<Complex64> for AnyScalar {
1314 fn from(value: Complex64) -> Self {
1315 Self::from_value(value)
1316 }
1317}
1318
1319impl TryFrom<AnyScalar> for f64 {
1320 type Error = &'static str;
1321
1322 fn try_from(value: AnyScalar) -> std::result::Result<Self, Self::Error> {
1323 value.as_f64().ok_or("cannot convert complex scalar to f64")
1324 }
1325}
1326
1327impl From<AnyScalar> for Complex64 {
1328 fn from(value: AnyScalar) -> Self {
1329 value.value().into_complex()
1330 }
1331}
1332
1333impl Add<&AnyScalar> for &AnyScalar {
1334 type Output = AnyScalar;
1335
1336 fn add(self, rhs: &AnyScalar) -> Self::Output {
1337 AnyScalar::fallback_result(
1338 self.try_add(rhs),
1339 "add",
1340 || {
1341 AnyScalar::scalar_value_from_backend(
1342 self.to_backend_scalar() + rhs.to_backend_scalar(),
1343 )
1344 },
1345 self.tracks_grad() || rhs.tracks_grad(),
1346 )
1347 }
1348}
1349
1350impl Add<AnyScalar> for AnyScalar {
1351 type Output = AnyScalar;
1352
1353 fn add(self, rhs: AnyScalar) -> Self::Output {
1354 Add::add(&self, &rhs)
1355 }
1356}
1357
1358impl Add<AnyScalar> for &AnyScalar {
1359 type Output = AnyScalar;
1360
1361 fn add(self, rhs: AnyScalar) -> Self::Output {
1362 Add::add(self, &rhs)
1363 }
1364}
1365
1366impl Add<&AnyScalar> for AnyScalar {
1367 type Output = AnyScalar;
1368
1369 fn add(self, rhs: &AnyScalar) -> Self::Output {
1370 Add::add(&self, rhs)
1371 }
1372}
1373
1374impl Sub<&AnyScalar> for &AnyScalar {
1375 type Output = AnyScalar;
1376
1377 fn sub(self, rhs: &AnyScalar) -> Self::Output {
1378 Add::add(self, &Neg::neg(rhs))
1379 }
1380}
1381
1382impl Sub<AnyScalar> for AnyScalar {
1383 type Output = AnyScalar;
1384
1385 fn sub(self, rhs: AnyScalar) -> Self::Output {
1386 Sub::sub(&self, &rhs)
1387 }
1388}
1389
1390impl Sub<AnyScalar> for &AnyScalar {
1391 type Output = AnyScalar;
1392
1393 fn sub(self, rhs: AnyScalar) -> Self::Output {
1394 Sub::sub(self, &rhs)
1395 }
1396}
1397
1398impl Sub<&AnyScalar> for AnyScalar {
1399 type Output = AnyScalar;
1400
1401 fn sub(self, rhs: &AnyScalar) -> Self::Output {
1402 Sub::sub(&self, rhs)
1403 }
1404}
1405
1406impl Mul<&AnyScalar> for &AnyScalar {
1407 type Output = AnyScalar;
1408
1409 fn mul(self, rhs: &AnyScalar) -> Self::Output {
1410 AnyScalar::fallback_result(
1411 self.try_mul(rhs),
1412 "mul",
1413 || {
1414 AnyScalar::scalar_value_from_backend(
1415 self.to_backend_scalar() * rhs.to_backend_scalar(),
1416 )
1417 },
1418 self.tracks_grad() || rhs.tracks_grad(),
1419 )
1420 }
1421}
1422
1423impl Mul<AnyScalar> for AnyScalar {
1424 type Output = AnyScalar;
1425
1426 fn mul(self, rhs: AnyScalar) -> Self::Output {
1427 Mul::mul(&self, &rhs)
1428 }
1429}
1430
1431impl Mul<AnyScalar> for &AnyScalar {
1432 type Output = AnyScalar;
1433
1434 fn mul(self, rhs: AnyScalar) -> Self::Output {
1435 Mul::mul(self, &rhs)
1436 }
1437}
1438
1439impl Mul<&AnyScalar> for AnyScalar {
1440 type Output = AnyScalar;
1441
1442 fn mul(self, rhs: &AnyScalar) -> Self::Output {
1443 Mul::mul(&self, rhs)
1444 }
1445}
1446
1447impl Div<&AnyScalar> for &AnyScalar {
1448 type Output = AnyScalar;
1449
1450 fn div(self, rhs: &AnyScalar) -> Self::Output {
1451 AnyScalar::fallback_result(
1452 self.try_div(rhs),
1453 "div",
1454 || {
1455 AnyScalar::scalar_value_from_backend(
1456 self.to_backend_scalar() / rhs.to_backend_scalar(),
1457 )
1458 },
1459 self.tracks_grad() || rhs.tracks_grad(),
1460 )
1461 }
1462}
1463
1464impl Div<AnyScalar> for AnyScalar {
1465 type Output = AnyScalar;
1466
1467 fn div(self, rhs: AnyScalar) -> Self::Output {
1468 Div::div(&self, &rhs)
1469 }
1470}
1471
1472impl Div<AnyScalar> for &AnyScalar {
1473 type Output = AnyScalar;
1474
1475 fn div(self, rhs: AnyScalar) -> Self::Output {
1476 Div::div(self, &rhs)
1477 }
1478}
1479
1480impl Div<&AnyScalar> for AnyScalar {
1481 type Output = AnyScalar;
1482
1483 fn div(self, rhs: &AnyScalar) -> Self::Output {
1484 Div::div(&self, rhs)
1485 }
1486}
1487
1488impl Neg for &AnyScalar {
1489 type Output = AnyScalar;
1490
1491 fn neg(self) -> Self::Output {
1492 AnyScalar::fallback_result(
1493 self.try_neg(),
1494 "neg",
1495 || AnyScalar::scalar_value_from_backend(-self.to_backend_scalar()),
1496 self.tracks_grad(),
1497 )
1498 }
1499}
1500
1501impl Neg for AnyScalar {
1502 type Output = AnyScalar;
1503
1504 fn neg(self) -> Self::Output {
1505 Neg::neg(&self)
1506 }
1507}
1508
1509impl Mul<AnyScalar> for f64 {
1510 type Output = AnyScalar;
1511
1512 fn mul(self, rhs: AnyScalar) -> Self::Output {
1513 AnyScalar::from_real(self) * rhs
1514 }
1515}
1516
1517impl Mul<AnyScalar> for Complex64 {
1518 type Output = AnyScalar;
1519
1520 fn mul(self, rhs: AnyScalar) -> Self::Output {
1521 AnyScalar::from(self) * rhs
1522 }
1523}
1524
1525impl Div<AnyScalar> for Complex64 {
1526 type Output = AnyScalar;
1527
1528 fn div(self, rhs: AnyScalar) -> Self::Output {
1529 AnyScalar::from(self) / rhs
1530 }
1531}
1532
1533impl Default for AnyScalar {
1534 fn default() -> Self {
1535 Self::zero()
1536 }
1537}
1538
1539impl Zero for AnyScalar {
1540 fn zero() -> Self {
1541 Self::from_real(0.0)
1542 }
1543
1544 fn is_zero(&self) -> bool {
1545 AnyScalar::is_zero(self)
1546 }
1547}
1548
1549impl One for AnyScalar {
1550 fn one() -> Self {
1551 Self::from_real(1.0)
1552 }
1553}
1554
1555impl PartialEq for AnyScalar {
1556 fn eq(&self, other: &Self) -> bool {
1557 self.value() == other.value()
1558 }
1559}
1560
1561impl PartialOrd for AnyScalar {
1562 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1563 match (self.value(), other.value()) {
1564 (ScalarValue::F32(lhs), ScalarValue::F32(rhs)) => lhs.partial_cmp(&rhs),
1565 (ScalarValue::F32(lhs), ScalarValue::F64(rhs)) => (lhs as f64).partial_cmp(&rhs),
1566 (ScalarValue::F64(lhs), ScalarValue::F32(rhs)) => lhs.partial_cmp(&(rhs as f64)),
1567 (ScalarValue::F64(lhs), ScalarValue::F64(rhs)) => lhs.partial_cmp(&rhs),
1568 _ => None,
1569 }
1570 }
1571}
1572
1573impl fmt::Display for AnyScalar {
1574 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1575 match self.value() {
1576 ScalarValue::F32(value) => value.fmt(f),
1577 ScalarValue::F64(value) => value.fmt(f),
1578 ScalarValue::C32(value) => value.fmt(f),
1579 ScalarValue::C64(value) => value.fmt(f),
1580 }
1581 }
1582}
1583
1584impl fmt::Debug for AnyScalar {
1585 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1586 let dtype = match self.value {
1587 ScalarValue::F32(_) => "f32",
1588 ScalarValue::F64(_) => "f64",
1589 ScalarValue::C32(_) => "c32",
1590 ScalarValue::C64(_) => "c64",
1591 };
1592 f.debug_struct("AnyScalar")
1593 .field("dtype", &dtype)
1594 .field("value", &self.value())
1595 .field("tracks_grad", &self.tracks_grad())
1596 .finish()
1597 }
1598}
1599
1600#[cfg(test)]
1601mod tests;