Different dtypes testing

heat/core/tests/test_suites/basic_test.py

@@ -1,4 +1,5 @@
 from unittest import TestCase
+
 from heat.core import dndarray, MPICommunication, MPI, types, factories
 import heat as ht
 
@@ -89,6 +90,12 @@ def assert_array_equal(self, heat_array, expected_array):
                 heat_array.shape, expected_array.shape
             ),
         )
+
+        if not heat_array.is_balanced():
+            # Array is distributed correctly
+            print("Heat array is unbalanced, balancing now")
+            heat_array.balance_()
+
         split = heat_array.split
         offset, local_shape, slices = heat_array.comm.chunk(heat_array.gshape, split)
         self.assertEqual(
@@ -99,10 +106,13 @@ def assert_array_equal(self, heat_array, expected_array):
         )
         local_numpy = heat_array._DNDarray__array.numpy()
 
-        # Array is distributed correctly
         equal_res = np.array(compare_func(local_numpy, expected_array[slices]))
         self.comm.Allreduce(MPI.IN_PLACE, equal_res, MPI.LAND)
-        self.assertTrue(equal_res, "Local tensors do not match the corresponding numpy slices.")
+        self.assertTrue(
+            equal_res,
+            "Local tensors do not match the corresponding numpy slices. "
+            "dtype was {}, split was {}".format(heat_array.dtype, heat_array.split),
+        )
         self.assertEqual(
             local_numpy.dtype,
             expected_array.dtype,
@@ -118,36 +128,111 @@ def assert_array_equal(self, heat_array, expected_array):
 
     def assert_func_equal(
         self,
-        tensor,
+        shape,
         heat_func,
         numpy_func,
         distributed_result=True,
         heat_args=None,
         numpy_args=None,
+        data_types=(np.int32, np.int64, np.float32, np.float64),
+        low=-10000,
+        high=10000,
     ):
         """
-        Checks if a heat function works equivalent to a given numpy function. The initial array is iteratively split
-        along each axis before the heat function is applied. After that the resulting arrays are compared using the
-        assert_array_equal function.
+        This function will create random tensors of the given shape with different data types.
+        All of these tensors will be tested with `ht.assert_func_equal_for_tensor`.
 
         Parameters
         ----------
-        tensor: tuple, list, numpy.ndarray or torch.Tensor
-            The tensor to check the function against.
-            If a tuple or list is provided instead a random array of the specified shape is generated
+        shape: tuple or list
+            The shape of which a random tensors will be created and tested against
         heat_func: function
             The function that is to be tested
         numpy_func: function
             The numpy implementation of an equivalent function to test against
+        heat_args: dictionary, optional
+            The keyword arguments that will be passed to the heat function. Array and split function don't need to be
+            specified. Default is {}.
+        numpy_args: dictionary, optional
+            The keyword arguments that will be passed to the numpy function. Array doesn't need to be specified.
+            Default is {}.
         distributed_result: bool, optional
             Specify whether the result of the heat function is distributed across all nodes or all nodes have the full
             result. Default is True.
+        data_types: list of numpy dtypes, optional
+            Tensors with all of these dtypes will be created and tested. Each type must to be a numpy dtype.
+            Default is [numpy.int32, numpy.int64, numpy.float32, numpy.float64]
+        low: int, optional
+            In case one of the data_types has integer types, this is the lower bound for the random values.
+            Default is -10000
+        high: int, optional
+            In case one of the data_types has integer types, this is the upper bound for the random values.
+            Default is 10000
+
+        Raises
+        ------
+        AssertionError if the functions to not perform equally.
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> import heat as ht
+        >>> self.assert_func_equal((2, 2), ht.exp, np.exp)
+
+        >>> self.assert_func_equal((2, 2), ht.exp, np.log)
+        AssertionError: [...]
+        >>> self.assert_func_equal((1, 3, 5), ht.any, np.any, distributed_result=False)
+
+        >>> heat_args = {'sorted': True, 'axis': 0}
+        >>> numpy_args = {'axis': 0}
+        >>> self.assert_func_equal([5, 5, 5, 5], ht.unique, np.unique, heat_arg=heat_args, numpy_args=numpy_args)
+        """
+        if not isinstance(shape, tuple) and not isinstance(shape, list):
+            raise ValueError(
+                "The shape must be either a list or a tuple but was {}".format(type(shape))
+            )
+
+        for dtype in data_types:
+            tensor = self._create_random_array(shape, dtype=dtype, low=low, high=high)
+            self.assert_func_equal_for_tensor(
+                tensor=tensor,
+                heat_func=heat_func,
+                numpy_func=numpy_func,
+                heat_args=heat_args,
+                numpy_args=numpy_args,
+                distributed_result=distributed_result,
+            )
+
+    def assert_func_equal_for_tensor(
+        self,
+        tensor,
+        heat_func,
+        numpy_func,
+        heat_args=None,
+        numpy_args=None,
+        distributed_result=True,
+    ):
+        """
+        This function will create random tensors of the given shape with different data types.
+        All of these tensors will be tested with `ht.assert_func_equal_for_tensor`.
+
+        Parameters
+        ----------
+        tensor: torch.Tensor or numpy.ndarray
+            The tensor on which the heat function will be executed.
+        heat_func: function
+            The function that is to be tested
+        numpy_func: function
+            The numpy implementation of an equivalent function to test against
         heat_args: dictionary, optional
             The keyword arguments that will be passed to the heat function. Array and split function don't need to be
             specified. Default is {}.
         numpy_args: dictionary, optional
             The keyword arguments that will be passed to the numpy function. Array doesn't need to be specified.
             Default is {}.
+        distributed_result: bool, optional
+            Specify whether the result of the heat function is distributed across all nodes or all nodes have the full
+            result. Default is True.
 
         Raises
         ------
@@ -158,20 +243,19 @@ def assert_func_equal(
         >>> import numpy as np
         >>> import heat as ht
         >>> a = np.arange(10)
-        >>> self.assert_func_equal(a, ht.exp, np.exp)
+        >>> self.assert_func_equal_for_tensor(a, ht.exp, np.exp)
 
-        >>> self.assert_func_equal(a, ht.exp, np.log)
+        >>> self.assert_func_equal_for_tensor(a, ht.exp, np.log)
         AssertionError: [...]
-        >>> self.assert_func_equal((1, 3, 5), ht.any, np.any, distributed_result=False)
+        >>> self.assert_func_equal_for_tensor(a, ht.any, np.any, distributed_result=False)
 
+        >>> a = torch.ones([5, 5, 5, 5])
         >>> heat_args = {'sorted': True, 'axis': 0}
         >>> numpy_args = {'axis': 0}
-        >>> self.assert_func_equal([5, 5, 5, 5], ht.unique, np.unique, heat_arg=heat_args, numpy_args=numpy_args)
+        >>> self.assert_func_equal_for_tensor(a, ht.unique, np.unique, heat_arg=heat_args, numpy_args=numpy_args)
         """
         self.assertTrue(callable(heat_func))
         self.assertTrue(callable(numpy_func))
-        if isinstance(tensor, tuple) or isinstance(tensor, list):
-            tensor = self._create_random_array(tensor)
 
         if heat_args is None:
             heat_args = {}
@@ -209,9 +293,50 @@ def assert_func_equal(
             else:
                 self.assertTrue(np.array_equal(ht_res._DNDarray__array.numpy(), np_res))
 
-    def _create_random_array(self, shape):
+    def _create_random_array(self, shape, dtype=np.float64, low=-10000, high=10000):
+        """
+        Creates a random array based on the input parameters.
+        The used seed will be printed to stdout for debugging purposes.
+
+        Parameters
+        ----------
+        shape: list or tuple
+            The shape of the random array to be created.
+        dtype: np.dtype, optional
+            The datatype of the resulting array.
+            If dtype is subclass of np.floating then numpy.random.randn is used to create random values.
+            If dtype is subclass of np.integer then numpy.random.randint is called with the low and high argument to
+            create the random values.
+            Default is numpy.float64
+        low: int, optional
+            In case dtype is an integer type, this is the lower bound for the random values.
+            Default is -10000
+        low: int, optional
+            In case dtype is an integer type, this is the upper bound for the random values.
+            Default is 10000
+
+        Returns
+        -------
+        res: numpy.ndarray
+            An array of random values with the specified shape and dtype.
+
+        Raises
+        ------
+        ValueError if the dtype is not a subtype of numpy.integer or numpy.floating
+        """
         seed = np.random.randint(1000000, size=(1,))
+        print("using seed {} for random values".format(seed))
         self.comm.Bcast(seed, root=0)
         np.random.seed(seed=seed.item())
-        array = np.random.randn(*shape)
+        if issubclass(dtype, np.floating):
+            array = np.random.randn(*shape)
+        elif issubclass(dtype, np.integer):
+            array = np.random.randint(low=low, high=high, size=shape)
+        else:
+            raise ValueError(
+                "Unsupported dtype. Expected a subclass of `np.floating` or `np.integer` but got {}".format(
+                    dtype
+                )
+            )
+        array = array.astype(dtype)
         return array

heat/core/tests/test_suites/test_basic_test.py

@@ -24,29 +24,45 @@ def test_assert_array_equal(self):
         self.assert_array_equal(heat_array, expected_array)
 
     def test_assert_func_equal(self):
+        # Testing with random values
+        shape = (5, 3, 2, 9)
+
+        self.assert_func_equal(shape, heat_func=ht.exp, numpy_func=np.exp, low=-100, high=100)
+
+        self.assert_func_equal(shape, heat_func=ht.exp2, numpy_func=np.exp2, low=-100, high=100)
+
+        # np.random.randn eventually creates values < 0 which will result in math.nan.
+        # Because math.nan != math.nan this would always produce an exception.
+        self.assert_func_equal(
+            shape, heat_func=ht.log, numpy_func=np.log, data_types=[np.int32, np.int64], low=1
+        )
+
+        with self.assertRaises(AssertionError):
+            self.assert_func_equal(shape, heat_func=ht.exp, numpy_func=np.exp2, low=-100, high=100)
+
+        with self.assertRaises(ValueError):
+            self.assert_func_equal(np.ones(shape), heat_func=np.exp, numpy_func=np.exp)
+
+    def test_assert_func_equal_for_tensor(self):
         array = np.ones((self.get_size(), 20), dtype=np.int8)
         ht_func = ht.any
         np_func = np.any
-        self.assert_func_equal(array, ht_func, np_func, distributed_result=False)
+        self.assert_func_equal_for_tensor(array, ht_func, np_func, distributed_result=False)
 
         array = np.array([[1, 2, 4, 1, 3], [1, 4, 7, 5, 1]], dtype=np.int8)
         ht_func = ht.expand_dims
         np_func = np.expand_dims
         ht_args = {"axis": 1}
         np_args = {"axis": 1}
-        self.assert_func_equal(array, ht_func, np_func, heat_args=ht_args, numpy_args=np_args)
-
-        # Testing with random values
-        shape = (5, 3, 2, 9)
-        ht_func = ht.exp
-        np_func = np.exp
-        self.assert_func_equal(shape, heat_func=ht_func, numpy_func=np_func)
+        self.assert_func_equal_for_tensor(
+            array, ht_func, np_func, heat_args=ht_args, numpy_args=np_args
+        )
 
         array = torch.randn(15, 15)
         ht_func = ht.exp
         np_func = np.exp
-        self.assert_func_equal(array, heat_func=ht_func, numpy_func=np_func)
+        self.assert_func_equal_for_tensor(array, heat_func=ht_func, numpy_func=np_func)
 
         array = ht.ones((15, 15))
         with self.assertRaises(TypeError):
-            self.assert_func_equal(array, heat_func=ht_func, numpy_func=np_func)
+            self.assert_func_equal_for_tensor(array, heat_func=ht_func, numpy_func=np_func)