From bf1493f23219fc8a6530bdb7f82f9a6a5351da2d Mon Sep 17 00:00:00 2001
From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com>
Date: Fri, 20 Feb 2026 09:48:50 +0100
Subject: [PATCH 1/2] Add basic sanity checks to DNDarray

---
 heat/core/dndarray.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py
index 16ce355700..41cc13ce69 100644
--- a/heat/core/dndarray.py
+++ b/heat/core/dndarray.py
@@ -71,6 +71,8 @@ def __init__(
         comm: Communication,
         balanced: bool,
     ):
+        dtype = types.canonical_heat_type(dtype)
+
         self.__array = array
         self.__gshape = gshape
         self.__dtype = dtype
@@ -84,8 +86,18 @@ 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}!"
+        )
+        if self.ndim == array.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:

From 8c396b9bf7e59130ae141a0b345555f7dd652c56 Mon Sep 17 00:00:00 2001
From: Thomas Baumann <39156931+brownbaerchen@users.noreply.github.com>
Date: Mon, 23 Feb 2026 11:36:36 +0100
Subject: [PATCH 2/2] Require same dimension in global array as local data

---
 heat/core/dndarray.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/heat/core/dndarray.py b/heat/core/dndarray.py
index 41cc13ce69..eec9b29c13 100644
--- a/heat/core/dndarray.py
+++ b/heat/core/dndarray.py
@@ -88,13 +88,12 @@ def __init__(
 
         # check for inconsistencies between local and global arrays
         assert str(array.device) == device.torch_device
-        assert self.ndim >= array.ndim, (
+        assert self.ndim == array.ndim, (
             f"Local dimension {array.ndim} exceeds global dimension {self.ndim}!"
         )
-        if self.ndim == array.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 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}"
         )