diff --git a/rstsr-core/src/device_cpu_serial/operators/op_binary_common.rs b/rstsr-core/src/device_cpu_serial/operators/op_binary_common.rs index bc3a305c..5bf2e0d3 100644 --- a/rstsr-core/src/device_cpu_serial/operators/op_binary_common.rs +++ b/rstsr-core/src/device_cpu_serial/operators/op_binary_common.rs @@ -1,7 +1,6 @@ use crate::prelude_dev::*; -use core::ops::Div; use num::complex::ComplexFloat; -use num::{Float, Signed, Zero}; +use num::{Float, Signed}; use rstsr_dtype_traits::{DTypeIntoFloatAPI, ExtNum}; // TODO: log1p @@ -216,33 +215,20 @@ where impl OpSignAPI for DeviceCpuSerial where - T: Clone + ComplexFloat + Div, + T: ExtNum, D: DimAPI, { type TOut = T; fn op_muta_refb(&self, a: &mut Vec>, la: &Layout, b: &Vec, lb: &Layout) -> Result<()> { self.op_muta_refb_func(a, la, b, lb, &mut |a, b| { - // sign(x) = x / |x|, but 0 / 0 yields NaN, so a zero magnitude maps to - // 0 (preserving the sign of -0.0 and complex zero, matching NumPy). - let abs = b.abs(); - if abs.is_zero() { - a.write(*b); - } else { - a.write(*b / abs); - } + a.write(b.clone().ext_sign()); }) } fn op_muta(&self, a: &mut Vec>, la: &Layout) -> Result<()> { self.op_muta_func(a, la, &mut |a| unsafe { - let b = a.assume_init_read(); - let abs = b.abs(); - if abs.is_zero() { - a.write(b); - } else { - a.write(b / abs); - } + a.write(a.assume_init_read().ext_sign()); }) } } diff --git a/rstsr-core/src/feature_rayon/auto_impl/op_binary_common.rs b/rstsr-core/src/feature_rayon/auto_impl/op_binary_common.rs index 943f5c06..5712062e 100644 --- a/rstsr-core/src/feature_rayon/auto_impl/op_binary_common.rs +++ b/rstsr-core/src/feature_rayon/auto_impl/op_binary_common.rs @@ -1,7 +1,6 @@ use crate::prelude_dev::*; -use core::ops::Div; use num::complex::ComplexFloat; -use num::{Float, Signed, Zero}; +use num::{Float, Signed}; use rstsr_dtype_traits::{DTypeIntoFloatAPI, ExtNum}; // TODO: log1p @@ -219,33 +218,20 @@ where impl OpSignAPI for DeviceRayonAutoImpl where - T: Clone + Send + Sync + ComplexFloat + Div, + T: ExtNum + Send + Sync, D: DimAPI, { type TOut = T; fn op_muta_refb(&self, a: &mut Vec>, la: &Layout, b: &Vec, lb: &Layout) -> Result<()> { self.op_muta_refb_func(a, la, b, lb, &mut |a, b| { - // sign(x) = x / |x|, but 0 / 0 yields NaN, so a zero magnitude maps to - // 0 (preserving the sign of -0.0 and complex zero, matching NumPy). - let abs = b.abs(); - if abs.is_zero() { - a.write(*b); - } else { - a.write(*b / abs); - } + a.write(b.clone().ext_sign()); }) } fn op_muta(&self, a: &mut Vec>, la: &Layout) -> Result<()> { self.op_muta_func(a, la, &mut |a| unsafe { - let b = a.assume_init_read(); - let abs = b.abs(); - if abs.is_zero() { - a.write(b); - } else { - a.write(b / abs); - } + a.write(a.assume_init_read().ext_sign()); }) } } diff --git a/rstsr-core/src/tensor/operators/op_unary_common.rs b/rstsr-core/src/tensor/operators/op_unary_common.rs index b0c3186e..80c3ff04 100644 --- a/rstsr-core/src/tensor/operators/op_unary_common.rs +++ b/rstsr-core/src/tensor/operators/op_unary_common.rs @@ -10,9 +10,7 @@ Most unary functions are of the same type. However, there are some exceptions, a - `Imag, Real, Abs`: - complex: generalized, not for inplace. - real: specialized, for inplace. -- `Sign`: - - complex: generalized, for inplace. - - real: specialized, for inplace. +- `Sign`: any `ExtNum` dtype (integers, floats, complex), same-type output, for inplace. */ diff --git a/rstsr-core/tests/core_func/math/test_unary_math.rs b/rstsr-core/tests/core_func/math/test_unary_math.rs index 3ebb1b1c..068bc3fb 100644 --- a/rstsr-core/tests/core_func/math/test_unary_math.rs +++ b/rstsr-core/tests/core_func/math/test_unary_math.rs @@ -31,16 +31,43 @@ mod custom_math_basic { let s = rt::tensor_from_nested!([0.0, 1.0, 4.0, 9.0], &device); assert_equal(rt::sqrt(&s), rt::tensor_from_nested!([0.0, 1.0, 2.0, 3.0], &device), None); - // np.sign([-2, 0, 3]) == [-1, 0, 1]; rstsr sign requires a Float input. + // np.sign([-2, 0, 3]) == [-1, 0, 1]; sign accepts any ExtNum dtype (floats and integers). let g = rt::tensor_from_nested!([-2.0, 0.0, 3.0], &device); assert_equal(rt::sign(&g), rt::tensor_from_nested!([-1.0, 0.0, 1.0], &device), None); + // np.sign(np.array([-2, 0, 3])) == array([-1, 0, 1]); integer dtypes stay integers + let gi = rt::tensor_from_nested!([-2, 0, 3], &device); + assert_eq!(rt::sign(&gi).to_vec(), vec![-1, 0, 1]); + + // np.sign(np.array([0, 5], dtype=np.uint8)) == array([0, 1], dtype=uint8) + let gu: Tensor = rt::tensor_from_nested!([0, 5], &device); + assert_eq!(rt::sign(&gu).to_vec(), vec![0u8, 1u8]); + // np.floor / np.ceil / np.trunc on [-1.5, 0.5, 2.4] let f = rt::tensor_from_nested!([-1.5, 0.5, 2.4], &device); assert_equal(rt::floor(&f), rt::tensor_from_nested!([-2.0, 0.0, 2.0], &device), None); assert_equal(rt::ceil(&f), rt::tensor_from_nested!([-1.0, 1.0, 3.0], &device), None); assert_equal(rt::trunc(&f), rt::tensor_from_nested!([-1.0, 0.0, 2.0], &device), None); } + + #[test] + fn test_sign_special_values() { + crate::specify_test!("test_sign_special_values"); + + let mut device = TESTCFG.device.clone(); + device.set_default_order(RowMajor); + + // np.sign([nan, inf, -inf, -0.0, 0.0]) == [nan, 1.0, -1.0, 0.0, 0.0]; + // infinities map to ±1.0, and both signed zeros map to +0.0. + let a: Tensor = rt::asarray((vec![f64::NAN, f64::INFINITY, f64::NEG_INFINITY, -0.0, 0.0], &device)); + let s = rt::sign(&a).to_vec(); + assert!(s[0].is_nan()); + assert_eq!(s[1], 1.0); + assert_eq!(s[2], -1.0); + assert_eq!(s[3], 0.0); + assert!(!s[3].is_sign_negative()); + assert_eq!(s[4], 0.0); + } } #[cfg(test)] diff --git a/rstsr-dtype-traits/src/ext_num.rs b/rstsr-dtype-traits/src/ext_num.rs index f453f83d..f2cbeed9 100644 --- a/rstsr-dtype-traits/src/ext_num.rs +++ b/rstsr-dtype-traits/src/ext_num.rs @@ -25,6 +25,19 @@ pub trait ExtNum: Clone { /* #endregion */ + /* #region sign */ + + /// Computes the sign of the number, with the same output type as the input. + /// + /// This follows NumPy's `np.sign` conventions: + /// - signed integers: `-1`, `0` or `1` (no overflow at the type's minimum); + /// - unsigned integers: `0` or `1`; + /// - floats: `-1.0`, `0.0` or `1.0`; infinities map to `±1.0`, NaN maps to NaN; + /// - complex: `z / |z|`, and complex zero maps to zero. + fn ext_sign(self) -> Self; + + /* #endregion */ + /* #region real-imag */ /// Returns the real part of the number. @@ -65,6 +78,17 @@ impl ExtNum for T { } /* #endregion */ + /* #region sign */ + #[inline] + fn ext_sign(self) -> Self { + if self == 0 { + 0 + } else { + 1 + } + } + /* #endregion */ + /* #region real-imag */ #[inline] fn ext_real(self) -> Self { @@ -96,6 +120,20 @@ impl ExtNum for T { } /* #endregion */ + /* #region sign */ + #[inline] + fn ext_sign(self) -> Self { + // comparison form avoids overflow at the type's minimum (e.g. i8::MIN) + if self > 0 { + 1 + } else if self < 0 { + -1 + } else { + 0 + } + } + /* #endregion */ + /* #region real-imag */ #[inline] fn ext_real(self) -> Self { @@ -124,6 +162,22 @@ impl ExtNum for T { } /* #endregion */ + /* #region sign */ + #[inline] + fn ext_sign(self) -> Self { + // NaN maps to NaN; infinities map to ±1; both zeros map to 0 (NumPy convention) + if self.is_nan() { + self + } else if self > 0.0 { + 1.0 + } else if self < 0.0 { + -1.0 + } else { + 0.0 + } + } + /* #endregion */ + /* #region real-imag */ #[inline] fn ext_real(self) -> Self { @@ -160,6 +214,22 @@ impl ExtNum for T { } /* #endregion */ + /* #region sign */ + #[inline] + fn ext_sign(self) -> Self { + // NaN maps to NaN; infinities map to ±1; both zeros map to 0 (NumPy convention) + if self.is_nan() { + self + } else if self > Self::ZERO { + Self::ONE + } else if self < Self::ZERO { + Self::NEG_ONE + } else { + Self::ZERO + } + } + /* #endregion */ + /* #region real-imag */ #[inline] fn ext_real(self) -> Self { @@ -196,6 +266,20 @@ impl ExtNum for T { } /* #endregion */ + /* #region sign */ + #[inline] + fn ext_sign(self) -> Self { + // sign(z) = z / |z|, but 0 / 0 yields NaN, so a zero magnitude maps to + // 0 (matching NumPy) + let abs = self.norm(); + if abs == 0.0 { + Self::ZERO + } else { + self / abs + } + } + /* #endregion */ + /* #region real-imag */ #[inline] fn ext_real(self) -> Self::AbsOut {