fix: only assert same shape and chunk shape in spatial dimension

apps/hiopy/_distribute_work.py

@@ -105,11 +105,11 @@ def distribute_work(grouped_data_vars, size):
     variables = grouped_data_vars.pop(group)
 
     assert all(
-        variables[0].shape == v.shape for v in variables
-    ), "All variables in a group need to have the same shape"
+        variables[0].shape[-1] == v.shape[-1] for v in variables
+    ), "All variables in a group need to have the same spatial shape"
     assert all(
-        variables[0].chunks == v.chunks for v in variables
-    ), "All variables in a group need to have the same chunk shape"
+        variables[0].chunks[-1] == v.chunks[-1] for v in variables
+    ), "All variables in a group need to have the same spatial chunk shape"
 
     if gsize > bytes_per_rank:
         # Compute the number of ranks required for this group