fixing errors

Author: aman-coder03

toqito/matrix_ops/partial_trace.py

@@ -143,16 +143,19 @@ def partial_trace(
         if n % dim != 0:
             raise ValueError("Invalid: If `dim` is a scalar, it must evenly divide matrix dimension.")
         dim = np.array([dim, n // dim])
-    elif isinstance(dim, list):
+    elif isinstance(dim, (list, tuple, np.ndarray)):
+        dim = np.array(dim)
+        if dim.ndim != 1:
+            raise ValueError("Invalid: `dim` must be a 1D array-like of ints.")
         if len(dim) == 1:
             d = dim[0]
             if n % d != 0:
-                raise ValueError("Invalid: If `dim` is a scalar, it must evenly divide matrix dimension.")
+                raise ValueError(
+                    "Invalid: If `dim` is a scalar, it must evenly divide matrix dimension."
+                )
             dim = np.array([d, n // d])
-        else:
-            dim = np.array(dim)
     else:
-        raise ValueError("Invalid: `dim` must be int or list of ints.")
+        raise ValueError("Invalid: `dim` must be int or array-like of ints.")
 
     num_sys = len(dim)
     prod_dim = np.prod(dim)

toqito/matrix_ops/tests/test_partial_trace.py

@@ -684,15 +684,15 @@ def test_dim_list_branches(dim_value):
         # dim=None and non-perfect-square
         (np.ones((6, 6)), [0], None, "Cannot infer subsystem dimensions directly"),
 
-        # Invalid dim type
-        (np.eye(4), [0], (2, 2), None),
+        # dim not 1D (2D array)
+        (np.eye(4), [0], np.array([[2, 2]]), "1D"),
+
+        # completely invalid dim type
+        (np.eye(4), [0], "invalid_dim", "int or array-like"),
     ],
 )
 def test_partial_trace_invalid_inputs(input_mat, sys_value, dim_value, error_msg):
     """Test various invalid parameter combinations."""
-    if error_msg:
-        with pytest.raises(ValueError, match=error_msg):
-            partial_trace(input_mat, sys_value, dim_value)
-    else:
-        with pytest.raises(ValueError):
-            partial_trace(input_mat, sys_value, dim_value)
+    with pytest.raises(ValueError, match=error_msg):
+        partial_trace(input_mat, sys_value, dim_value)
+