rstsr-dtype-traits: implement ext_sign for integers, floats and complex

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<T, D> OpSignAPI<T, D> for DeviceCpuSerial
 where
-    T: Clone + ComplexFloat + Div<T::Real, Output = T>,
+    T: ExtNum,
     D: DimAPI,
 {
     type TOut = T;
 
     fn op_muta_refb(&self, a: &mut Vec<MaybeUninit<T>>, la: &Layout<D>, b: &Vec<T>, lb: &Layout<D>) -> 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<MaybeUninit<T>>, la: &Layout<D>) -> 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());
         })
     }
 }

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<T, D> OpSignAPI<T, D> for DeviceRayonAutoImpl
 where
-    T: Clone + Send + Sync + ComplexFloat + Div<T::Real, Output = T>,
+    T: ExtNum + Send + Sync,
     D: DimAPI,
 {
     type TOut = T;
 
     fn op_muta_refb(&self, a: &mut Vec<MaybeUninit<T>>, la: &Layout<D>, b: &Vec<T>, lb: &Layout<D>) -> 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<MaybeUninit<T>>, la: &Layout<D>) -> 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());
         })
     }
 }

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.
 
 */
 

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<u8, _> = 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<f64, _> = 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)]

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 {