add cross

CHANGELOG.md

@@ -19,9 +19,11 @@
 ### DNDarray
 - [#856](https://github.com/helmholtz-analytics/heat/pull/856) New `DNDarray` method `__torch_proxy__`
 
+# Feature additions
 ### Linear Algebra
 - [#840](https://github.com/helmholtz-analytics/heat/pull/840) New feature: `vecdot()`
 - [#846](https://github.com/helmholtz-analytics/heat/pull/846) New features `norm`, `vector_norm`, `matrix_norm`
+- [#850](https://github.com/helmholtz-analytics/heat/pull/850) New Feature `cross`
 ### Logical
 - [#862](https://github.com/helmholtz-analytics/heat/pull/862) New feature `signbit`
 ### Manipulations

heat/core/linalg/basics.py

@@ -21,9 +21,11 @@
 from .. import rounding
 from .. import sanitation
 from .. import statistics
+from .. import stride_tricks
 from .. import types
 
 __all__ = [
+    "cross",
     "dot",
     "matmul",
     "matrix_norm",
@@ -40,6 +42,77 @@
 ]
 
 
+def cross(x1: DNDarray, x2: DNDarray, axis: int = -1) -> DNDarray:
+    """
+    Returns the cross product.
+
+    Parameters
+    ----------
+    x1 : DNDarray
+        First input array.
+    x2 : DNDarray
+        Second input array. Must have the same shape as 'x1'.
+    axis : int
+        Axis that defines the vectors for which to compute the cross product. Default: -1
+
+    Raises
+    ------
+    ValueError
+        If the two input arrays don't match in shape, split, device, or comm. If the vectors are along the split axis.
+    TypeError
+        If 'axis' is not an integer.
+
+    Examples
+    --------
+    >>> a = ht.eye(3)
+    >>> b = ht.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]])
+    >>> cross = ht.cross(a, b)
+    DNDarray([[0., 0., 1.],
+              [1., 0., 0.],
+              [0., 1., 0.]], dtype=ht.float32, device=cpu:0, split=None)
+    """
+    sanitation.sanitize_in(x1)
+    sanitation.sanitize_in(x2)
+
+    if x1.gshape != x2.gshape:
+        raise ValueError(
+            "'x1' and 'x2' must have the same shape, {} != {}".format(x1.gshape, x2.gshape)
+        )
+    if x1.split != x2.split:
+        raise ValueError(
+            "'x1' and 'x2' must have the same split, {} != {}".format(x1.split, x2.split)
+        )
+    if x1.device != x2.device:
+        raise ValueError(
+            "'x1' and 'x2' must have the same device type, {} != {}".format(x1.device, x2.device)
+        )
+    if x1.comm != x2.comm:  # pragma: no cover
+        raise ValueError("'x1' and 'x2' must have the same comm, {} != {}".format(x1.comm, x2.comm))
+
+    if not isinstance(axis, int):
+        try:
+            axis = int(axis)
+        except Exception:
+            raise TypeError("'axis' must be an integer.")
+
+    axis = stride_tricks.sanitize_axis(x1.shape, axis)
+
+    if x1.split == axis:
+        raise ValueError(
+            "The computation of the cross product with vectors along the split axis is not supported."
+        )
+    else:
+        x1.balance_()
+        x2.balance_()
+
+    promoted = torch.promote_types(x1.larray.dtype, x2.larray.dtype)
+
+    ret = torch.cross(x1.larray.type(promoted), x2.larray.type(promoted), dim=axis)
+    ret = DNDarray(ret, x1.gshape, types.heat_type_of(ret), x1.split, x1.device, x1.comm, True)
+
+    return ret
+
+
 def dot(a: DNDarray, b: DNDarray, out: Optional[DNDarray] = None) -> Union[DNDarray, float]:
     """
     Returns the dot product of two ``DNDarrays``.

heat/core/linalg/tests/test_basics.py

@@ -1,3 +1,4 @@
+from typing import Type
 import torch
 import os
 import unittest
@@ -8,6 +9,54 @@
 
 
 class TestLinalgBasics(TestCase):
+    def test_cross(self):
+        a = ht.eye(3)
+        b = ht.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]])
+
+        # different types
+        cross = ht.cross(a, b)
+        self.assertEqual(cross.shape, a.shape)
+        self.assertEqual(cross.dtype, a.dtype)
+        self.assertEqual(cross.split, a.split)
+        self.assertEqual(cross.comm, a.comm)
+        self.assertEqual(cross.device, a.device)
+        self.assertTrue(ht.equal(cross, ht.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])))
+
+        # axis
+        a = ht.eye(3, split=0)
+        b = ht.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=ht.float, split=0)
+
+        cross = ht.cross(a, b)
+        self.assertEqual(cross.shape, a.shape)
+        self.assertEqual(cross.dtype, a.dtype)
+        self.assertEqual(cross.split, a.split)
+        self.assertEqual(cross.comm, a.comm)
+        self.assertEqual(cross.device, a.device)
+        self.assertTrue(ht.equal(cross, ht.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])))
+
+        a = ht.eye(3, dtype=ht.int8, split=1)
+        b = ht.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=ht.int8, split=1)
+
+        cross = ht.cross(a, b, axis=0)
+        self.assertEqual(cross.shape, a.shape)
+        self.assertEqual(cross.dtype, a.dtype)
+        self.assertEqual(cross.split, a.split)
+        self.assertEqual(cross.comm, a.comm)
+        self.assertEqual(cross.device, a.device)
+        self.assertTrue(ht.equal(cross, ht.array([[0, 0, -1], [-1, 0, 0], [0, -1, 0]])))
+
+        with self.assertRaises(ValueError):
+            ht.cross(ht.eye(3), ht.eye(4))
+        with self.assertRaises(ValueError):
+            ht.cross(ht.eye(3, split=0), ht.eye(3, split=1))
+        if torch.cuda.is_available():
+            with self.assertRaises(ValueError):
+                ht.cross(ht.eye(3, device="gpu"), ht.eye(3, device="cpu"))
+        with self.assertRaises(TypeError):
+            ht.cross(ht.eye(3), ht.eye(3), axis="wasd")
+        with self.assertRaises(ValueError):
+            ht.cross(ht.eye(3, split=0), ht.eye(3, split=0), axis=0)
+
     def test_dot(self):
         # ONLY TESTING CORRECTNESS! ALL CALLS IN DOT ARE PREVIOUSLY TESTED
         # cases to test: