diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py index 5f17f3ba5c..8cf1dff2cc 100644 --- a/heat/core/dndarray.py +++ b/heat/core/dndarray.py @@ -56,6 +56,8 @@ def __init__( comm: Communication, balanced: bool, ): + dtype = types.canonical_heat_type(dtype) + self.__array = array self.__gshape = gshape self.__dtype = dtype @@ -69,8 +71,17 @@ def __init__( self.__partitions_dict__ = None self.__lshape_map = None - # check for inconsistencies between torch and heat devices + # check for inconsistencies between local and global arrays assert str(array.device) == device.torch_device + assert self.ndim == array.ndim, ( + f"Local dimension {array.ndim} exceeds global dimension {self.ndim}!" + ) + assert all([gshape[i] >= array.shape[i] for i in range(self.ndim)]), ( + f"Local shape {array.shape} is larger than global shape {gshape}" + ) + assert dtype == types.canonical_heat_type(array.dtype), ( + f"Local datatype {array.dtype} is incompatible with global datatype {dtype}" + ) @property def balanced(self) -> bool: