@@ -2122,7 +2122,12 @@ class _VectorizedFinalizedBuilderMixin(_FinalizedBuilderMixin):

             except Exception as exc:

                 _raise_user_error(exc, val)

-        ham = system._normalize_matrix_blocks(ham, len(to_site_array))

+        expected_shape = (

+            len(to_site_array),

+            to_family.norbs,

+            from_family.norbs if not is_onsite else to_family.norbs,

+        )

+        ham = system._normalize_matrix_blocks(ham, expected_shape)

         return ham

@@ -2446,8 +2451,17 @@ def _make_onsite_terms(builder, sites, site_arrays, term_offset):

     # Normalize any constant values and check that the shapes are consistent.

     onsite_term_values = [

         val if callable(val)

-        else system._normalize_matrix_blocks(val, len(val))

-        for val in onsite_term_values

+        else

+            system._normalize_matrix_blocks(

+                val,

+                (

+                    len(val),

+                    site_arrays[sa].family.norbs,

+                    site_arrays[sa].family.norbs,

+                ),

+            )

+        for (_, sa), val in

+        zip(onsite_to_term_nr.keys(), onsite_term_values)

     ]

     # Sort the sites in each term, and also sort the values

     # in the same way if they are a constant (as opposed to a callable).

@@ -2560,8 +2574,17 @@ def _make_hopping_terms(builder, graph, sites, site_arrays, cell_size, term_offs

         # Normalize any constant values and check that the shapes are consistent.

         hopping_term_values = [

             val if callable(val)

-            else system._normalize_matrix_blocks(val, len(val))

-            for val in hopping_term_values

+            else

+                system._normalize_matrix_blocks(

+                    val,

+                    (

+                        len(val),

+                        site_arrays[to_sa].family.norbs,

+                        site_arrays[from_sa].family.norbs,

+                    ),

+                )

+            for (_, to_sa, from_sa), val in

+            zip(hopping_to_term_nr.keys(), hopping_term_values)

         ]

         # Sort the hoppings in each term, and also sort the values

         # in the same way if they are a constant (as opposed to a callable).

@@ -327,12 +327,14 @@ class _FunctionalOnsite:

     def __call__(self, site_range, site_offsets, *args):

         sites = self.sites

-        offset = self.site_ranges[site_range][0]

+        offset, norbs, _ = self.site_ranges[site_range]

         try:

             ret = [self.onsite(sites[offset + i], *args) for i in site_offsets]

         except Exception as exc:

             _raise_user_error(exc, self.onsite, vectorized=False)

-        return _normalize_matrix_blocks(ret, len(site_offsets))

+

+        expected_shape = (len(site_offsets), norbs, norbs)

+        return _normalize_matrix_blocks(ret, expected_shape)

 class _VectorizedFunctionalOnsite:

@@ -349,7 +351,9 @@ class _VectorizedFunctionalOnsite:

             ret = self.onsite(sites, *args)

         except Exception as exc:

             _raise_user_error(exc, self.onsite, vectorized=True)

-        return _normalize_matrix_blocks(ret, len(site_offsets))

+

+        expected_shape = (len(sites), sites.family.norbs, sites.family.norbs)

+        return _normalize_matrix_blocks(ret, expected_shape)

 class _FunctionalOnsiteNoTransform:

@@ -359,12 +363,15 @@ class _FunctionalOnsiteNoTransform:

         self.site_ranges = site_ranges

     def __call__(self, site_range, site_offsets, *args):

-        site_ids = self.site_ranges[site_range][0] + site_offsets

+        offset, norbs, _ = self.site_ranges[site_range]

+        site_ids = offset + site_offsets

         try:

             ret = [self.onsite(id, *args) for id in site_ids]

         except Exception as exc:

             _raise_user_error(exc, self.onsite, vectorized=False)

-        return _normalize_matrix_blocks(ret, len(site_offsets))

+

+        expected_shape = (len(site_offsets), norbs, norbs)

+        return _normalize_matrix_blocks(ret, expected_shape)

 class _DictOnsite:

@@ -376,7 +383,9 @@ class _DictOnsite:

     def __call__(self, site_range, site_offsets, *args):

         fam = self.range_family_map[site_range]

         ret = [self.onsite[fam]] * len(site_offsets)

-        return _normalize_matrix_blocks(ret, len(site_offsets))

+

+        expected_shape = (len(site_offsets), fam.norbs, fam.norbs)

+        return _normalize_matrix_blocks(ret, expected_shape)

 def _normalize_onsite(syst, onsite, check_hermiticity):

@@ -993,11 +1002,13 @@ cdef class _LocalOperator:

                         syst.hamiltonian(where[i, 0], where[i, 1], *args, params=params)

                         for i in which

                     ]

-                data = _normalize_matrix_blocks(data, len(which))

-                # Checks for data consistency

+

                 (to_sr, from_sr) = term_id

-                to_norbs = syst.site_ranges[to_sr][1]

-                from_norbs = syst.site_ranges[from_sr][1]

+                to_norbs, from_norbs = syst.site_ranges[to_sr][1], syst.site_ranges[from_sr][1]

+                expected_shape = (len(which), to_norbs, from_norbs)

+                data = _normalize_matrix_blocks(data, expected_shape)

+                # Checks for data consistency

+

                 _check_hams(data, to_norbs, from_norbs, is_onsite and check_hermiticity)

                 return data

@@ -742,13 +742,13 @@ def is_vectorized(syst):

     return isinstance(syst, (FiniteVectorizedSystem, InfiniteVectorizedSystem))

-def _normalize_matrix_blocks(blocks, expected_length):

+def _normalize_matrix_blocks(blocks, expected_shape):

     """Normalize a sequence of matrices into a single 3D numpy array

     Parameters

     ----------

     blocks : sequence of complex array-like

-    expected_length : int

+    expected_shape : (int, int, int)

     """

     try:

         blocks = np.asarray(blocks, dtype=complex)

@@ -757,21 +757,19 @@ def _normalize_matrix_blocks(blocks, expected_length):

             "Matrix elements declared with incompatible shapes."

         ) from None

     if len(blocks.shape) == 0:  # scalar → broadcast to vector of 1x1 matrices

-        blocks = np.tile(blocks, (expected_length, 1, 1))

+        blocks = np.tile(blocks, (expected_shape[0], 1, 1))

     elif len(blocks.shape) == 1:  # vector → interpret as vector of 1x1 matrices

         blocks = blocks.reshape(-1, 1, 1)

     elif len(blocks.shape) == 2:  # matrix → broadcast to vector of matrices

-        blocks = np.tile(blocks, (expected_length, 1, 1))

+        blocks = np.tile(blocks, (expected_shape[0], 1, 1))

-    if len(blocks.shape) != 3:

+    if blocks.shape != expected_shape:

         msg = (

-            "Vectorized value functions must return an array of"

-            "scalars or an array of matrices."

+            "Expected values of shape {}, but received values of shape {}"

+            .format(expected_shape, blocks.shape)

         )

         raise ValueError(msg)

-    if blocks.shape[0] != expected_length:

-        raise ValueError("Value functions must return a single value per "

-                         "onsite/hopping.")

+

     return blocks