add sign, sgn

CHANGELOG.md

@@ -23,6 +23,8 @@
 - [#854](https://github.com/helmholtz-analytics/heat/pull/854) New Feature: `moveaxis`
 ### Random
 - [#858](https://github.com/helmholtz-analytics/heat/pull/858) New Feature: `standard_normal`, `normal`
+### Rounding
+- [#827](https://github.com/helmholtz-analytics/heat/pull/827) New feature: `sign`, `sgn`
 
 # v1.1.1
 - [#864](https://github.com/helmholtz-analytics/heat/pull/864) Dependencies: constrain `torchvision` version range to match supported `pytorch` version range.

heat/core/rounding.py

@@ -12,7 +12,19 @@
 from . import sanitation
 from . import types
 
-__all__ = ["abs", "absolute", "ceil", "clip", "fabs", "floor", "modf", "round", "trunc"]
+__all__ = [
+    "abs",
+    "absolute",
+    "ceil",
+    "clip",
+    "fabs",
+    "floor",
+    "modf",
+    "round",
+    "sgn",
+    "sign",
+    "trunc",
+]
 
 
 def abs(
@@ -328,6 +340,90 @@ def round(
 DNDarray.round.__doc__ = round.__doc__
 
 
+def sgn(x: DNDarray, out: Optional[DNDarray] = None) -> DNDarray:
+    """
+    Returns an indication of the sign of a number, element-wise. The definition for complex values is equivalent to :math:`x / |x|`.
+
+    Parameters
+    ----------
+    x : DNDarray
+        Input array
+    out : DNDarray, optional
+        A location in which to store the results.
+
+    See Also
+    --------
+    :func:`sign`
+        Equivalent function on non-complex arrays. The definition for complex values is equivalent to :math:`x / \\sqrt{x \\cdot x}`
+
+    Examples
+    --------
+    >>> a = ht.array([-1, -0.5, 0, 0.5, 1])
+    >>> ht.sign(a)
+    DNDarray([-1., -1.,  0.,  1.,  1.], dtype=ht.float32, device=cpu:0, split=None)
+    >>> ht.sgn(ht.array([5-2j, 3+4j]))
+    DNDarray([(0.9284766912460327-0.3713906705379486j),  (0.6000000238418579+0.800000011920929j)], dtype=ht.complex64, device=cpu:0, split=None)
+    """
+    return _operations.__local_op(torch.sgn, x, out)
+
+
+def sign(x: DNDarray, out: Optional[DNDarray] = None) -> DNDarray:
+    """
+    Returns an indication of the sign of a number, element-wise. The definition for complex values is equivalent to :math:`x / \\sqrt{x \\cdot x}`.
+
+    Parameters
+    ----------
+    x : DNDarray
+        Input array
+    out : DNDarray, optional
+        A location in which to store the results.
+
+    See Also
+    --------
+    :func:`sgn`
+        Equivalent function on non-complex arrays. The definition for complex values is equivalent to :math:`x / |x|`.
+
+    Examples
+    --------
+    >>> a = ht.array([-1, -0.5, 0, 0.5, 1])
+    >>> ht.sign(a)
+    DNDarray([-1., -1.,  0.,  1.,  1.], dtype=ht.float32, device=cpu:0, split=None)
+    >>> ht.sign(ht.array([5-2j, 3+4j]))
+    DNDarray([(1+0j), (1+0j)], dtype=ht.complex64, device=cpu:0, split=None)
+    """
+    # special case for complex values
+    if types.heat_type_is_complexfloating(x.dtype):
+        sanitation.sanitize_in(x)
+        if out is not None:
+            sanitation.sanitize_out(out, x.shape, x.split, x.device)
+            out.larray.copy_(x.larray)
+            data = out.larray
+        else:
+            data = torch.clone(x.larray)
+        # NOTE remove when min version >= 1.9
+        if "1.7" in torch.__version__ or "1.8" in torch.__version__:
+            pos = data != 0
+        else:  # pragma: no cover
+            indices = torch.nonzero(data)
+            pos = torch.split(indices, 1, 1)
+        data[pos] = x.larray[pos] / torch.sqrt(torch.square(x.larray[pos]))
+
+        if out is not None:
+            out.__dtype = types.heat_type_of(data)
+            return out
+        return DNDarray(
+            data,
+            gshape=x.shape,
+            dtype=types.heat_type_of(data),
+            split=x.split,
+            device=x.device,
+            comm=x.comm,
+            balanced=x.balanced,
+        )
+
+    return _operations.__local_op(torch.sign, x, out)
+
+
 def trunc(x: DNDarray, out: Optional[DNDarray] = None) -> DNDarray:
     """
     Return the trunc of the input, element-wise.

heat/core/tests/test_rounding.py

@@ -301,6 +301,78 @@ def test_round(self):
         self.assertEqual(float64_round_distrbd.dtype, ht.float64)
         self.assert_array_equal(float64_round_distrbd, comparison)
 
+    def test_sgn(self):
+        # floats
+        a = ht.array([-1, -0.5, 0, 0.5, 1])
+        signed = ht.sgn(a)
+        comparison = ht.array([-1.0, -1, 0, 1, 1])
+
+        self.assertEqual(signed.dtype, comparison.dtype)
+        self.assertEqual(signed.shape, comparison.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertTrue(ht.equal(signed, comparison))
+
+        # complex
+        a = ht.array([[1 - 2j, -0.5 + 1j], [0 - 3j, 4 + 6j]], split=0)
+        signed = ht.sgn(a)
+        comparison = torch.sgn(torch.tensor([[1 - 2j, -0.5 + 1j], [0 - 3j, 4 + 6j]]))
+        comparison = comparison.to(a.device.torch_device)
+
+        self.assertEqual(signed.dtype, ht.heat_type_of(comparison))
+        self.assertEqual(signed.shape, a.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertTrue(ht.equal(signed, ht.array(comparison, split=0)))
+
+    def test_sign(self):
+        # floats 1d
+        a = ht.array([-1, -0.5, 0, 0.5, 1])
+        signed = ht.sign(a)
+        comparison = ht.array([-1.0, -1, 0, 1, 1])
+
+        self.assertEqual(signed.dtype, comparison.dtype)
+        self.assertEqual(signed.shape, comparison.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertEqual(signed.split, a.split)
+        self.assertTrue(ht.equal(signed, comparison))
+
+        # complex + 2d + split
+        a = ht.array([[1 - 2j, -0.5 + 1j], [0, 4 + 6j]], split=0)
+        signed = ht.sign(a)
+        comparison = ht.array([[1 + 0j, -1 + 0j], [0 + 0j, 1 + 0j]], split=0)
+
+        self.assertEqual(signed.dtype, comparison.dtype)
+        self.assertEqual(signed.shape, comparison.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertEqual(signed.split, a.split)
+        self.assertTrue(ht.allclose(signed.real, comparison.real))
+        self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
+
+        # complex + split + out
+        a = ht.array([[1 - 2j, -0.5 + 1j], [0, 4 + 6j]], split=1)
+        b = ht.empty_like(a)
+        signed = ht.sign(a, b)
+        comparison = ht.array([[1 + 0j, -1 + 0j], [0 + 0j, 1 + 0j]], split=1)
+
+        self.assertIs(b, signed)
+        self.assertEqual(signed.dtype, comparison.dtype)
+        self.assertEqual(signed.shape, comparison.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertEqual(signed.split, a.split)
+        self.assertTrue(ht.allclose(signed.real, comparison.real))
+        self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
+
+        # zeros + 3d + complex + split
+        a = ht.zeros((4, 4, 4), dtype=ht.complex128, split=2)
+        signed = ht.sign(a)
+        comparison = ht.zeros((4, 4, 4), dtype=ht.complex128, split=2)
+
+        self.assertEqual(signed.dtype, comparison.dtype)
+        self.assertEqual(signed.shape, comparison.shape)
+        self.assertEqual(signed.device, a.device)
+        self.assertEqual(signed.split, a.split)
+        self.assertTrue(ht.allclose(signed.real, comparison.real))
+        self.assertTrue(ht.allclose(signed.imag, comparison.imag, atol=2e-5))
+
     def test_trunc(self):
         base_array = np.random.randn(20)