diff --git a/toqito/matrix_ops/inner_product.py b/toqito/matrix_ops/inner_product.py
index dde9210ac..2c4b8867b 100644
--- a/toqito/matrix_ops/inner_product.py
+++ b/toqito/matrix_ops/inner_product.py
@@ -9,15 +9,41 @@ def inner_product(v1: np.ndarray, v2: np.ndarray) -> float:
     The inner product is calculated as follows:
 
     .. math::
-        \left\langle \begin{pmatrix}a_1 \\ \vdots \\ a_n\end{pmatrix},\begin{pmatrix}b_1 \\ \vdots \\ b_n\end{pmatrix}\right\rangle = \begin{pmatrix} a_1,\cdots, a_n\end{pmatrix}\begin{pmatrix}b_1 \\ \vdots \\ b_n\end{pmatrix} = a_1 b_1 + \cdots + a_n b_n
+        \left\langle 
+        \begin{pmatrix} 
+            a_1 \\ \vdots \\ a_n
+        \end{pmatrix},
+        \begin{pmatrix} 
+            b_1 \\ \vdots \\ b_n
+        \end{pmatrix} \right\rangle = 
+        \begin{pmatrix} 
+            a_1, \cdots, a_n
+        \end{pmatrix}
+        \begin{pmatrix} 
+            b_1 \\ \vdots \\ b_n
+        \end{pmatrix} = a_1 b_1 + \cdots + a_n b_n
 
     Example
     ==========
 
-    The inner product of the vectors :math:`v1 = \begin{pmatrix}1 \\ 2 \\ 3 \end{pmatrix}` and :math:`v2 = \begin{pmatrix}4 \\ 5 \\ 6 \ \end{pmatrix}` looks as follows:
+    The inner product of the vectors :math:`v1 = \begin{pmatrix}1 \\ 2 \\ 3 \end{pmatrix}` and :math:`v2 =
+    \begin{pmatrix}4 \\ 5 \\ 6 \ \end{pmatrix}` looks as follows:
 
     .. math::
-        \left\langle \begin{pmatrix}1 \\ 2 \\ 3\end{pmatrix},\begin{pmatrix}4 \\ 5 \\ 6\end{pmatrix}\right\rangle = \begin{pmatrix} 1,2, 3\end{pmatrix}\begin{pmatrix}4 \\ 5 \\ 6\end{pmatrix} = 1\times 4 + 2\times 5 + 3\times 6 = 32
+        \left\langle 
+        \begin{pmatrix} 
+            1 \\ 2 \\ 3
+        \end{pmatrix},
+        \begin{pmatrix}
+            4 \\ 5 \\ 6
+        \end{pmatrix}\right\rangle =
+        \begin{pmatrix} 
+            1, 2, 3
+        \end{pmatrix}
+        \begin{pmatrix}
+            4 \\ 5 \\ 6
+        \end{pmatrix} = 1\times 4 + 2\times 5 + 3\times 6 =
+        32
 
     In :code:`toqito`, this looks like this:
 
@@ -33,16 +59,12 @@ def inner_product(v1: np.ndarray, v2: np.ndarray) -> float:
         https://en.wikipedia.org/wiki/Inner_product_space
 
     :raises ValueError: Vector dimensions are mismatched.
-    :param v1: v1 and v2, both vectors of dimenstions :math:`(n,1)` where :math:`n>1`.
-    :param v2: v1 and v2, both vectors of dimenstions :math:`(n,1)` where :math:`n>1`.
+    :param v1: v1 and v2, both vectors of dimensions :math:`(n,1)` where :math:`n>1`.
+    :param v2: v1 and v2, both vectors of dimensions :math:`(n,1)` where :math:`n>1`.
     :return: The computed inner product.
     """
-    # Check for dimensional validity
-    if not (v1.shape[0] == v2.shape[0] and v1.shape[0] > 1 and len(v1.shape) == 1):
-        raise ValueError("Dimension mismatch")
-
-    res = 0
-    for i in range(v1.shape[0]):
-        res += v1[i] * v2[i]
-
-    return res
+    if v1.ndim != 1 or v2.ndim != 1:
+        raise ValueError("Both v1 and v2 must be 1D vectors.")
+    if v1.shape[0] != v2.shape[0]:
+        raise ValueError("Dimension mismatch between v1 and v2.")
+    return np.vdot(v1, v2)
diff --git a/toqito/matrix_ops/outer_product.py b/toqito/matrix_ops/outer_product.py
index f73554f42..ba8cc86e9 100644
--- a/toqito/matrix_ops/outer_product.py
+++ b/toqito/matrix_ops/outer_product.py
@@ -13,11 +13,33 @@ def outer_product(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
 
     Example
     ==========
-
-    The outer product of the vectors :math:`v1 = \begin{pmatrix}1 \\ 2 \\ 3 \end{pmatrix}` and :math:`v2 = \begin{pmatrix}4 \\ 5 \\ 6 \ \end{pmatrix}` looks as follows:
+    The outer product of the vectors :math:`v1 = \begin{pmatrix}1 \\ 2 \\ 3 \end{pmatrix}` and :math:`v2 =
+    \begin{pmatrix}4 \\ 5 \\ 6 \ \end{pmatrix}` looks as follows:
 
     .. math::
-        \left|\begin{pmatrix}1\\2\\3\end{pmatrix}\right\rangle\left\langle\begin{pmatrix}4\\5\\6\end{pmatrix}\right|=\begin{pmatrix}1\\2\\3\end{pmatrix}\begin{pmatrix}4&5&6\end{pmatrix}=\begin{pmatrix}1\times4&1\times5&1\times6\\2\times4&2\times5&2\times6\\3\times4&3\times5&3\times6\end{pmatrix}=\begin{pmatrix}4&5&6\\8&10&12\\12&15&18\end{pmatrix}
+        \left|\begin{pmatrix}
+            1 \\ 2 \\ 3
+        \end{pmatrix}\right\rangle
+        \left\langle
+        \begin{pmatrix}
+            4 \\ 5 \\ 6
+        \end{pmatrix}\right|=
+        \begin{pmatrix}
+            1 \\ 2 \\ 3
+        \end{pmatrix}
+        \begin{pmatrix}
+            4 & 5 & 6
+        \end{pmatrix}=
+        \begin{pmatrix}
+            1 \times 4 & 1 \times 5 & 1 \times 6 \\ 
+            2 \times 4 & 2 \times 5 & 2 \times 6 \\ 
+            3 \times 4 & 3 \times 5 & 3 \times 6
+        \end{pmatrix}=
+        \begin{pmatrix}
+            4 & 5 & 6 \\ 
+            8 & 10 & 12 \\ 
+            12 & 15 & 18
+        \end{pmatrix}
 
     In :code:`toqito`, this looks like this:
 
@@ -35,16 +57,13 @@ def outer_product(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
         https://en.wikipedia.org/wiki/Outer_product
 
     :raises ValueError: Vector dimensions are mismatched.
-    :param v1: v1 and v2, both vectors of dimenstions :math:`(n,1)` where :math:`n>1`.
-    :param v2: v1 and v2, both vectors of dimenstions :math:`(n,1)` where :math:`n>1`.
+    :param v1: v1 and v2, both vectors of dimensions :math:`(n,1)` where :math:`n>1`.
+    :param v2: v1 and v2, both vectors of dimensions :math:`(n,1)` where :math:`n>1`.
     :return: The computed outer product.
     """
-    # Check for dimensional validity
-    if not (v1.shape[0] == v2.shape[0] and v1.shape[0] > 1 and len(v1.shape) == 1):
-        raise ValueError("Dimension mismatch")
-
-    res = np.ndarray((v1.shape[0], v1.shape[0]))
-    for i in range(v1.shape[0]):
-        for j in range(v1.shape[0]):
-            res[i, j] = v1[i] * v2[j]
-    return res
+    if v1.ndim != 1 or v2.ndim != 1:
+        raise ValueError("Both v1 and v2 must be 1D vectors.")
+    if v1.shape[0] != v2.shape[0]:
+        raise ValueError("Dimension mismatch between v1 and v2.")
+
+    return np.outer(v1, np.conj(v2))
\ No newline at end of file
diff --git a/toqito/matrix_ops/tests/test_inner_product.py b/toqito/matrix_ops/tests/test_inner_product.py
index 24e1ba5e9..aab770dc3 100644
--- a/toqito/matrix_ops/tests/test_inner_product.py
+++ b/toqito/matrix_ops/tests/test_inner_product.py
@@ -1,44 +1,28 @@
 """Test inner_product."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import inner_product
 
 
-def test_inner_product():
-    """Test with two vectors, no complications."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([4, 5, 6])
-    expected_res = 32
-    np.testing.assert_equal(inner_product(v1, v2), expected_res)
-
-
-def test_inner_product_negative_input():
-    """Test with two vectors, with negative input value."""
-
-    v1, v2 = np.array([-1, 2, 3]), np.array([4, 5, 6])
-    expected_res = 24
-    np.testing.assert_equal(inner_product(v1, v2), expected_res)
-
-
-def test_inner_product_negative_output():
-    """Test with two vectors, with negative expected output."""
-
-    v1, v2 = np.array([1, 2, -3]), np.array([4, 5, 6])
-    expected_res = -4
-    np.testing.assert_equal(inner_product(v1, v2), expected_res)
-
-
-def test_inner_product_different_dimensions():
-    """Test with two vectors of different dimensions."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([4, 5, 6, 7])
-    with np.testing.assert_raises(ValueError):
-        inner_product(v1, v2)
-
-
-def test_inner_product_different_dimensions_2():
-    """Test with a vector and a 2d array."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([[4, 5, 6], [7, 8, 9]])
-    with np.testing.assert_raises(ValueError):
+@pytest.mark.parametrize("v1, v2, expected_result", [
+    # Test with two vectors, no complications.
+    (np.array([1, 2, 3]), np.array([4, 5, 6]), 32),
+    # Test with two vectors, with negative input value.
+    (np.array([-1, 2, 3]), np.array([4, 5, 6]), 24),
+    # Test with two vectors, with negative expected output.
+    (np.array([1, 2, -3]), np.array([4, 5, 6]), -4),
+])
+def test_inner_product(v1, v2, expected_result):
+   assert inner_product(v1, v2) == expected_result
+
+
+@pytest.mark.parametrize("v1, v2", [
+    # Different dimensions of vectors.
+    (np.array([1, 2, 3]), np.array([4, 5, 6, 7])),
+    # Vector and 2D array.
+    (np.array([1, 2, 3]), np.array([[4, 5, 6], [7, 8, 9]])),
+])
+def test_inner_product_invalid_input(v1, v2):
+    with pytest.raises(ValueError):
         inner_product(v1, v2)
diff --git a/toqito/matrix_ops/tests/test_outer_product.py b/toqito/matrix_ops/tests/test_outer_product.py
index d48d87b3d..7d924fedb 100644
--- a/toqito/matrix_ops/tests/test_outer_product.py
+++ b/toqito/matrix_ops/tests/test_outer_product.py
@@ -1,36 +1,26 @@
 """Test outer_product."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import outer_product
 
 
-def test_outer_product():
-    """Test with two vectors, no complications."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([4, 5, 6])
-    expected_res = np.array([[4, 5, 6], [8, 10, 12], [12, 15, 18]])
-    np.testing.assert_equal(outer_product(v1, v2), expected_res)
-
-
-def test_outer_product_negative():
-    """Test with two vectors, with negative input/output values."""
-
-    v1, v2 = np.array([-1, 2, 3]), np.array([4, 5, 6])
-    expected_res = np.array([[-4, -5, -6], [8, 10, 12], [12, 15, 18]])
-    np.testing.assert_equal(outer_product(v1, v2), expected_res)
-
-
-def test_outer_product_different_dimensions():
-    """Test with two vectors of different dimensions."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([4, 5, 6, 7])
-    with np.testing.assert_raises(ValueError):
-        outer_product(v1, v2)
-
-
-def test_outer_product_different_dimensions_2():
-    """Test with a vector and a 2d array."""
-
-    v1, v2 = np.array([1, 2, 3]), np.array([[4, 5, 6], [7, 8, 9]])
-    with np.testing.assert_raises(ValueError):
+@pytest.mark.parametrize("v1, v2, expected_result", [
+    # Test with two vectors, no complications.
+    (np.array([1, 2, 3]), np.array([4, 5, 6]), np.array([[4, 5, 6], [8, 10, 12], [12, 15, 18]])),
+    # Test with two vectors, with negative input/output values.
+    (np.array([-1, 2, 3]), np.array([4, 5, 6]), np.array([[-4, -5, -6], [8, 10, 12], [12, 15, 18]])),
+])
+def test_outer_product(v1, v2, expected_result):
+    np.testing.assert_array_equal(outer_product(v1, v2), expected_result)
+
+
+@pytest.mark.parametrize("v1, v2", [
+    # Different dimensions of vectors.
+    (np.array([1, 2, 3]), np.array([4, 5, 6, 7])),
+    # Vector and 2D array.
+    (np.array([1, 2, 3]), np.array([[4, 5, 6], [7, 8, 9]])),
+])
+def test_outer_product_invalid_input(v1, v2):
+    with pytest.raises(ValueError):
         outer_product(v1, v2)
diff --git a/toqito/matrix_ops/tests/test_tensor.py b/toqito/matrix_ops/tests/test_tensor.py
index 0138b30d5..5ffc06fe6 100644
--- a/toqito/matrix_ops/tests/test_tensor.py
+++ b/toqito/matrix_ops/tests/test_tensor.py
@@ -1,5 +1,6 @@
 """Test tensor."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import tensor
 from toqito.states import basis
@@ -150,7 +151,28 @@ def test_tensor_list_3():
     np.testing.assert_equal(np.all(bool_mat), True)
 
 
+def test_tensor_with_three_or_more_matrices():
+    """Test tensor product with a numpy array containing three or more matrices."""
+    # Three matrices to be Kronecker multiplied
+    matrix1 = np.array([[1, 2]])
+    matrix2 = np.array([[3], [4]])
+    matrix3 = np.array([[5, 6]])
+    matrix4 = np.array([[7, 8]])
+
+    # The numpy array containing the matrices
+    matrices = np.array([matrix1, matrix2, matrix3, matrix4], dtype=object)
+
+    # Expected output: Kronecker product of matrix1, matrix2, and matrix3
+    expected_output = np.kron(np.kron(matrix1, np.kron(matrix2, matrix3)), matrix4)
+
+    # Call the tensor function
+    result = tensor(matrices)
+
+    # Assert that the result is as expected
+    np.testing.assert_array_equal(result, expected_output)
+
+
 def test_tensor_empty_args():
     r"""Test tensor with no arguments."""
-    with np.testing.assert_raises(ValueError):
-        tensor()
+    with pytest.raises(ValueError):
+        tensor()
\ No newline at end of file
diff --git a/toqito/matrix_ops/tests/test_unvec.py b/toqito/matrix_ops/tests/test_unvec.py
index 16169babb..fba402e0e 100644
--- a/toqito/matrix_ops/tests/test_unvec.py
+++ b/toqito/matrix_ops/tests/test_unvec.py
@@ -1,28 +1,15 @@
 """Test unvec."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import unvec
 
 
-def test_unvec():
-    """Test standard unvec operation on a vector."""
-    expected_res = np.array([[1, 2], [3, 4]])
-
-    test_input_vec = np.array([1, 3, 2, 4])
-
-    res = unvec(test_input_vec)
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
-
-
-def test_unvec_custom_dim():
-    """Test standard unvec operation on a vector with custom dimension."""
-    expected_res = np.array([[1], [3], [2], [4]])
-
-    test_input_vec = np.array([1, 3, 2, 4])
-
-    res = unvec(test_input_vec, [4, 1])
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
+@pytest.mark.parametrize("vector, shape, expected_result", [
+    # Test standard unvec operation on a vector.
+    (np.array([1, 3, 2, 4]), None, np.array([[1, 2], [3, 4]])),
+    # Test standard unvec operation on a vector with custom dimension.
+    (np.array([1, 3, 2, 4]), [4, 1], np.array([[1], [3], [2], [4]])),
+])
+def test_unvec(vector, shape, expected_result):
+    np.testing.assert_array_equal(unvec(vector, shape), expected_result)
diff --git a/toqito/matrix_ops/tests/test_vec.py b/toqito/matrix_ops/tests/test_vec.py
index 69cfd8254..f64da27ee 100644
--- a/toqito/matrix_ops/tests/test_vec.py
+++ b/toqito/matrix_ops/tests/test_vec.py
@@ -1,16 +1,13 @@
 """Test vec."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import vec
 
 
-def test_vec():
-    """Test standard vec operation on a matrix."""
-    expected_res = np.array([[1], [3], [2], [4]])
-
-    test_input_mat = np.array([[1, 2], [3, 4]])
-
-    res = vec(test_input_mat)
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
+@pytest.mark.parametrize("vector, expected_result", [
+    # Test standard vec operation on a vector.
+    (np.array(np.array([[1, 2], [3, 4]])), np.array([[1], [3], [2], [4]])),
+])
+def test_vec(vector, expected_result):
+    np.testing.assert_array_equal(vec(vector), expected_result)
diff --git a/toqito/matrix_ops/tests/test_vectors_from_gram_matrix.py b/toqito/matrix_ops/tests/test_vectors_from_gram_matrix.py
index 0ae0a4793..481791ed7 100644
--- a/toqito/matrix_ops/tests/test_vectors_from_gram_matrix.py
+++ b/toqito/matrix_ops/tests/test_vectors_from_gram_matrix.py
@@ -1,18 +1,26 @@
 """Test vectors_from_gram_matrix."""
 import numpy as np
+import pytest
 
 from toqito.matrix_ops import vectors_from_gram_matrix
 
 
-def test_vectors_from_gram_matrix():
+@pytest.mark.parametrize("gram, expected_result", [
+    # Gram matrix is identity matrix.
+    (
+        np.identity(4), 
+        [
+            np.array([1, 0, 0, 0]),
+            np.array([0, 1, 0, 0]),
+            np.array([0, 0, 1, 0]),
+            np.array([0, 0, 0, 1])
+         ]
+     ),
+])
+def test_vectors_from_gram_matrix(gram, expected_result):
     """Test able to extract vectors from Gram matrix."""
-    gram = np.identity(4)
     vectors = vectors_from_gram_matrix(gram)
-
-    assert np.allclose(vectors[0], np.array([1, 0, 0, 0]))
-    assert np.allclose(vectors[1], np.array([0, 1, 0, 0]))
-    assert np.allclose(vectors[2], np.array([0, 0, 1, 0]))
-    assert np.allclose(vectors[3], np.array([0, 0, 0, 1]))
+    np.testing.assert_allclose(vectors, expected_result)
 
 
 def test_vectors_from_gram_matrix_not_psd():
@@ -24,3 +32,14 @@ def test_vectors_from_gram_matrix_not_psd():
     assert np.allclose(vectors[0][0], 1)
     assert np.allclose(vectors[1][0], -1 / 2)
     assert np.allclose(vectors[2][0], -1 / 2)
+
+
+@pytest.mark.parametrize("gram", [
+    # Non-square matrix.
+    (np.array([[1, 2], [4, 5], [7, 8]])),
+    # Non-symmetric matrix.
+    (np.array([[1, 2], [3, 4]])),
+])
+def test_vectors_from_gram_matrix_invalid_input(gram):
+    with pytest.raises(np.linalg.LinAlgError):
+        vectors_from_gram_matrix(gram)
diff --git a/toqito/matrix_ops/tests/test_vectors_to_gram_matrix.py b/toqito/matrix_ops/tests/test_vectors_to_gram_matrix.py
index af585c69e..f62b87baf 100644
--- a/toqito/matrix_ops/tests/test_vectors_to_gram_matrix.py
+++ b/toqito/matrix_ops/tests/test_vectors_to_gram_matrix.py
@@ -1,18 +1,32 @@
 """Test vectors_to_gram_matrix."""
 import numpy as np
+import pytest
 
-from toqito.matrices import standard_basis
 from toqito.matrix_ops import vectors_to_gram_matrix
 
 
-def test_vectors_to_gram_matrix():
+e_0, e_1 = np.array([[1], [0]]), np.array([[0], [1]])
+trine = [
+    e_0,
+    1 / 2 * (-e_0 + np.sqrt(3) * e_1),
+    -1 / 2 * (e_0 + np.sqrt(3) * e_1),
+]
+
+
+@pytest.mark.parametrize("vectors, expected_result", [
+    # Trine states.
+    (trine, np.array([[1, -1 / 2, -1 / 2], [-1 / 2, 1, -1 / 2], [-1 / 2, -1 / 2, 1]])),
+])
+def test_vectors_to_gram_matrix(vectors, expected_result):
     """Test able to construct Gram matrix from vectors."""
-    e_0, e_1 = standard_basis(2)
-    trine = [
-        e_0,
-        1 / 2 * (-e_0 + np.sqrt(3) * e_1),
-        -1 / 2 * (e_0 + np.sqrt(3) * e_1),
-    ]
-    gram = vectors_to_gram_matrix(trine)
-    expected_gram = np.array([[1, -1 / 2, -1 / 2], [-1 / 2, 1, -1 / 2], [-1 / 2, -1 / 2, 1]])
-    assert np.allclose(gram, expected_gram)
+    vectors = vectors_to_gram_matrix(trine)
+    np.testing.assert_allclose(vectors, expected_result)
+
+
+@pytest.mark.parametrize("vectors", [
+    # Vectors of different sizes. 
+    ([np.array([1, 2, 3]), np.array([1, 2])]),
+])
+def test_vectors_to_gram_matrix_invalid_input(vectors):
+    with np.testing.assert_raises(ValueError):
+        vectors_to_gram_matrix(vectors)
\ No newline at end of file
diff --git a/toqito/matrix_ops/vectors_from_gram_matrix.py b/toqito/matrix_ops/vectors_from_gram_matrix.py
index 1fb1f9993..81ad7ac73 100644
--- a/toqito/matrix_ops/vectors_from_gram_matrix.py
+++ b/toqito/matrix_ops/vectors_from_gram_matrix.py
@@ -3,12 +3,47 @@
 
 
 def vectors_from_gram_matrix(gram: np.ndarray) -> list[np.ndarray]:
-    """Obtain the corresponding ensemble of states from the Gram matrix.
+    r"""Obtain the corresponding ensemble of states from the Gram matrix [WikGram]_.
 
-    :param gram: Input Gram matrix.
-    :return: list of ensemble states
+    The function attempts to compute the Cholesky decomposition of the given Gram matrix. If the matrix is positive
+    definite, the Cholesky decomposition is returned. If the matrix is not positive definite, the function falls back to
+    eigendecomposition.
+
+    Examples
+    ========
+
+    # Example of a positive definite matrix:
+
+    >>> gram_matrix = np.array([[2, -1], [-1, 2]])
+    >>> vectors = vectors_from_gram_matrix(gram_matrix)
+    >>> vectors
+    [array([1.41421356, 0.        ]), array([-0.70710678,  1.22474487])]
+
+    # Example of a matrix that is not positive definite:
+
+    >>> gram_matrix = np.array([[0, 1], [1, 0]])
+    >>> vectors = vectors_from_gram_matrix(gram_matrix)
+
+    Matrix is not positive semidefinite. Using eigendecomposition as alternative.
+    >>> vectors
+    [array([0.70710678+0.j, 0.        +0.j]), array([0.        +0.j, 0.70710678+0.j])]
+
+    References
+    ==========
+    .. [WikGram] Wikipedia: Gram matrix
+        https://en.wikipedia.org/wiki/Gram_matrix
+
+    :raises LinAlgError: If the Gram matrix is not symmetric.
+    :raises LinAlgError: If the Gram matrix is not square..
+    :param gram: A square, symmetric matrix representing the Gram matrix.
+    :return: A list of vectors (np.ndarray) corresponding to the ensemble of states.
     """
     dim = gram.shape[0]
+    if gram.shape[0] != gram.shape[1]:
+        raise np.linalg.LinAlgError("The Gram matrix must be square.")
+    if not np.allclose(gram, gram.T):
+        raise np.linalg.LinAlgError("The Gram matrix must be symmetric.")
+
     # If matrix is PD, can do Cholesky decomposition:
     try:
         decomp = np.linalg.cholesky(gram)
diff --git a/toqito/matrix_ops/vectors_to_gram_matrix.py b/toqito/matrix_ops/vectors_to_gram_matrix.py
index a72ac6b03..2b76e431e 100644
--- a/toqito/matrix_ops/vectors_to_gram_matrix.py
+++ b/toqito/matrix_ops/vectors_to_gram_matrix.py
@@ -3,10 +3,42 @@
 
 
 def vectors_to_gram_matrix(vectors: list[np.ndarray]) -> np.ndarray:
-    """Given a list of vectors, return the Gram matrix."""
-    n = len(vectors)
-    gram = np.zeros((n, n), dtype=complex)
-    for i in range(n):
-        for j in range(n):
-            gram[i, j] = (vectors[i].conj().T @ vectors[j].reshape(-1, 1)).item()
-    return gram
+    r"""Construct the Gram matrix from a list of vectors [WikGram]_.
+
+    The Gram matrix is a matrix of inner products, where the entry G[i, j] is the inner product of vectors[i] and
+    vectors[j]. This function computes the Gram matrix for a given list of vectors.
+
+    Examples
+    ========
+    # Example with real vectors:
+    >>> vectors = [np.array([1, 2]), np.array([3, 4])]
+    >>> gram_matrix = vectors_to_gram_matrix(vectors)
+    >>> gram_matrix
+    array([[ 5., 11.],
+           [11., 25.]])
+
+    # Example with complex vectors:
+    >>> vectors = [np.array([1+1j, 2+2j]), np.array([3+3j, 4+4j])]
+    >>> gram_matrix = vectors_to_gram_matrix(vectors)
+    >>> gram_matrix
+    array([[ 10.+0.j,  20.+0.j],
+           [ 20.+0.j,  40.+0.j]])
+
+    References
+    ==========
+    .. [WikGram] Wikipedia: Gram matrix
+        https://en.wikipedia.org/wiki/Gram_matrix
+
+    :raises ValueError: If the vectors are not all of the same length.
+    :param vectors: A list of vectors (1D numpy arrays). All vectors must be of the same length. 
+    :return: A list of vectors corresponding to the ensemble of states.
+    """
+    # Check that all vectors are of the same length
+    if not all(v.shape == vectors[0].shape for v in vectors):
+        raise ValueError("All vectors must be of the same length.")
+
+    # Stack vectors into a matrix
+    stacked_vectors = np.column_stack(vectors)
+
+    # Compute Gram matrix using vectorized operations
+    return np.dot(stacked_vectors.conj().T, stacked_vectors)
diff --git a/toqito/perms/antisymmetric_projection.py b/toqito/perms/antisymmetric_projection.py
index 253c28c94..58fa9432d 100644
--- a/toqito/perms/antisymmetric_projection.py
+++ b/toqito/perms/antisymmetric_projection.py
@@ -1,24 +1,19 @@
 """Antisymmetric projection operator."""
 from itertools import permutations
-
 import numpy as np
-
 from scipy import linalg, sparse
 
 from toqito.perms import perm_sign, permutation_operator
 
 
-def antisymmetric_projection(
-    dim: int, p_param: int = 2, partial: bool = False
-) -> sparse.lil_matrix:
+def antisymmetric_projection(dim: int, p_param: int = 2, partial: bool = False) -> sparse.lil_matrix:
     r"""
     Produce the projection onto the antisymmetric subspace [WikAsym]_.
 
     Produces the orthogonal projection onto the anti-symmetric subspace of :code:`p_param` copies of
-    :code:`dim`-dimensional space. If :code:`partial = True`, then the antisymmetric projection (PA)
-    isn't the orthogonal projection itself, but rather a matrix whose columns form an orthonormal
-    basis for the symmetric subspace (and hence the PA * PA' is the orthogonal projection onto the
-    symmetric subspace.)
+    :code:`dim`-dimensional space. If :code:`partial = True`, then the antisymmetric projection (PA) isn't the
+    orthogonal projection itself, but rather a matrix whose columns form an orthonormal basis for the symmetric subspace
+    (and hence the PA * PA' is the orthogonal projection onto the symmetric subspace.)
 
     Examples
     ==========
@@ -40,9 +35,8 @@ def antisymmetric_projection(
     [[1., 0.],
      [0., 1.]]
 
-    When the :math:`p` value is greater than the dimension of the antisymmetric projection, this
-    just gives the matrix consisting of all zero entries. For instance, when :math:`d = 2` and
-    :math:`p = 3` we have that
+    When the :math:`p` value is greater than the dimension of the antisymmetric projection, this just gives the matrix
+    consisting of all zero entries. For instance, when :math:`d = 2` and :math:`p = 3` we have that
 
     .. math::
         A_{2, 3} =
diff --git a/toqito/perms/perfect_matchings.py b/toqito/perms/perfect_matchings.py
index ffef0e9a0..3f1f436ab 100644
--- a/toqito/perms/perfect_matchings.py
+++ b/toqito/perms/perfect_matchings.py
@@ -1,6 +1,4 @@
 """Perfect matchings."""
-
-
 import numpy as np
 
 
@@ -8,11 +6,11 @@ def perfect_matchings(num: list[int] | int | np.ndarray) -> np.ndarray:
     r"""
     Give all perfect matchings of :code:`num` objects.
 
-    The input can be either an even natural number (the number of objects to be matched) or a
-    `numpy` array containing an even number of distinct objects to be matched.
+    The input can be either an even natural number (the number of objects to be matched) or a `numpy` array containing
+    an even number of distinct objects to be matched.
 
-    Returns all perfect matchings of a given list of objects. That is, it returns all ways of
-    grouping an even number of objects into pairs.
+    Returns all perfect matchings of a given list of objects. That is, it returns all ways of grouping an even number of
+    objects into pairs.
 
     This function is adapted from QETLAB.
 
diff --git a/toqito/perms/perm_sign.py b/toqito/perms/perm_sign.py
index ecf6426ec..09748f372 100644
--- a/toqito/perms/perm_sign.py
+++ b/toqito/perms/perm_sign.py
@@ -1,8 +1,5 @@
 """Calculate permutation sign."""
-
-
 import numpy as np
-
 from scipy import linalg
 
 
@@ -10,9 +7,8 @@ def perm_sign(perm: np.ndarray | list[int]) -> float:
     """
     Compute the "sign" of a permutation [WikParPerm]_.
 
-    The sign (either -1 or 1) of the permutation :code:`perm` is
-    :code:`-1**`inv`, where :code:`inv` is the number of inversions contained
-    in :code:`perm`.
+    The sign (either -1 or 1) of the permutation :code:`perm` is :code:`-1**`inv`, where :code:`inv` is the number of
+    inversions contained in :code:`perm`.
 
     Examples
     ==========
@@ -22,8 +18,8 @@ def perm_sign(perm: np.ndarray | list[int]) -> float:
     .. math::
         [1, 2, 3, 4]
 
-    the permutation sign is positive as the number of elements in the vector are even. This can be
-    performed in :code:`toqito` as follows.
+    the permutation sign is positive as the number of elements in the vector are even. This can be performed in
+    :code:`toqito` as follows.
 
     >>> from toqito.perms import perm_sign
     >>> perm_sign([1, 2, 3, 4])
@@ -34,8 +30,8 @@ def perm_sign(perm: np.ndarray | list[int]) -> float:
     .. math::
         [1, 2, 3, 4, 5]
 
-    the permutation sign is negative as the number of elements in the vector are odd. This can be
-    performed in :code:`toqito` as follows.
+    the permutation sign is negative as the number of elements in the vector are odd. This can be performed in
+    :code:`toqito` as follows.
 
     >>> from toqito.perms import perm_sign
     >>> perm_sign([1, 2, 4, 3, 5])
@@ -50,5 +46,4 @@ def perm_sign(perm: np.ndarray | list[int]) -> float:
     :return: The value 1 if the permutation is of even length and the value of
              -1 if the permutation is of odd length.
     """
-    eye = np.eye(len(perm))
-    return linalg.det(eye[:, np.array(perm) - 1])
+    return linalg.det(np.eye(len(perm))[:, np.array(perm) - 1])
diff --git a/toqito/perms/permutation_operator.py b/toqito/perms/permutation_operator.py
index 2b79cae0b..d2ef1522b 100644
--- a/toqito/perms/permutation_operator.py
+++ b/toqito/perms/permutation_operator.py
@@ -1,6 +1,4 @@
 """Permutation operator."""
-
-
 import numpy as np
 import scipy as sp
 
@@ -16,18 +14,17 @@ def permutation_operator(
     r"""
     Produce a unitary operator that permutes subsystems.
 
-    Generates a unitary operator that permutes the order of subsystems according to the permutation
-    vector :code:`perm`, where the :math:`i^{th}` subsystem has dimension :code:`dim[i]`.
+    Generates a unitary operator that permutes the order of subsystems according to the permutation vector :code:`perm`,
+    where the :math:`i^{th}` subsystem has dimension :code:`dim[i]`.
 
-    If :code:`inv_perm` = True, it implements the inverse permutation of :code:`perm`. The
-    permutation operator return is full is :code:`is_sparse` is :code:`False` and sparse if
-    :code:`is_sparse` is :code:`True`.
+    If :code:`inv_perm` = True, it implements the inverse permutation of :code:`perm`. The permutation operator return
+    is full is :code:`is_sparse` is :code:`False` and sparse if :code:`is_sparse` is :code:`True`.
 
     Examples
     ==========
 
-    The permutation operator obtained with dimension :math:`d = 2` is equivalent to the standard
-    swap operator on two qubits
+    The permutation operator obtained with dimension :math:`d = 2` is equivalent to the standard swap operator on two
+    qubits
 
     .. math::
         P_{2, [2, 1]} =
diff --git a/toqito/perms/permute_systems.py b/toqito/perms/permute_systems.py
index 4b23acfbc..46a9b31b2 100644
--- a/toqito/perms/permute_systems.py
+++ b/toqito/perms/permute_systems.py
@@ -1,11 +1,7 @@
 """Permute systems."""
-
-
 import functools
 import operator
-
 import numpy as np
-
 from scipy import sparse
 
 from toqito.matrix_ops import vec
@@ -21,28 +17,26 @@ def permute_systems(
     r"""
     Permute subsystems within a state or operator.
 
-    Permutes the order of the subsystems of the vector or matrix :code:`input_mat` according to the
-    permutation vector :code:`perm`, where the dimensions of the subsystems are given by the vector
-    :code:`dim`. If :code:`input_mat` is non-square and not a vector, different row and column
-    dimensions can be specified by putting the row dimensions in the first row of :code:`dim` and
-    the columns dimensions in the second row of :code:`dim`.
+    Permutes the order of the subsystems of the vector or matrix :code:`input_mat` according to the permutation vector
+    :code:`perm`, where the dimensions of the subsystems are given by the vector :code:`dim`. If :code:`input_mat` is
+    non-square and not a vector, different row and column dimensions can be specified by putting the row dimensions in
+    the first row of :code:`dim` and the columns dimensions in the second row of :code:`dim`.
 
-    If :code:`row_only = True`, then only the rows of :code:`input_mat` are permuted, but not the
-    columns -- this is equivalent to multiplying :code:`input_mat` on the left by the corresponding
-    permutation operator, but not on the right.
+    If :code:`row_only = True`, then only the rows of :code:`input_mat` are permuted, but not the columns -- this is
+    equivalent to multiplying :code:`input_mat` on the left by the corresponding permutation operator, but not on the
+    right.
 
-    If :code:`row_only = False`, then :code:`dim` only needs to contain the row dimensions of the
-    subsystems, even if :code:`input_mat` is not square. If :code:`inv_perm = True`, then the
-    inverse permutation of :code:`perm` is applied instead of :code:`perm` itself.
+    If :code:`row_only = False`, then :code:`dim` only needs to contain the row dimensions of the subsystems, even if
+    :code:`input_mat` is not square. If :code:`inv_perm = True`, then the inverse permutation of :code:`perm` is applied
+    instead of :code:`perm` itself.
 
     Examples
     ==========
 
-    For spaces :math:`\mathcal{A}` and :math:`\mathcal{B}` where
-    :math:`\text{dim}(\mathcal{A}) = \text{dim}(\mathcal{B}) = 2` we may consider an operator
-    :math:`X \in \mathcal{A} \otimes \mathcal{B}`. Applying the `permute_systems` function with
-    vector :math:`[2,1]` on :math:`X`, we may reorient the spaces such that
-    :math:`X \in \mathcal{B} \otimes \mathcal{A}`.
+    For spaces :math:`\mathcal{A}` and :math:`\mathcal{B}` where :math:`\text{dim}(\mathcal{A}) =
+    \text{dim}(\mathcal{B}) = 2` we may consider an operator :math:`X \in \mathcal{A} \otimes \mathcal{B}`. Applying the
+    `permute_systems` function with vector :math:`[2,1]` on :math:`X`, we may reorient the spaces such that :math:`X \in
+    \mathcal{B} \otimes \mathcal{A}`.
 
     For example, if we define :math:`X \in \mathcal{A} \otimes \mathcal{B}` as
 
@@ -54,8 +48,8 @@ def permute_systems(
             13 & 14 & 15 & 16
         \end{pmatrix},
 
-    then applying the `permute_systems` function on :math:`X` to obtain
-    :math:`X \in \mathcal{B} \otimes \mathcal{A}` yield the following matrix
+    then applying the `permute_systems` function on :math:`X` to obtain :math:`X \in \mathcal{B} \otimes \mathcal{A}`
+    yield the following matrix
 
     .. math::
         X_{[2,1]} = \begin{pmatrix}
@@ -76,11 +70,10 @@ def permute_systems(
      [ 5  7  6  8]
      [13 15 14 16]]
 
-    For spaces :math:`\mathcal{A}, \mathcal{B}`, and :math:`\mathcal{C}` where
-    :math:`\text{dim}(\mathcal{A}) = \text{dim}(\mathcal{B}) = \text{dim}(\mathcal{C}) = 2` we may
-    consider an operator :math:`X \in \mathcal{A} \otimes \mathcal{B} \otimes \mathcal{C}`. Applying
-    the :code:`permute_systems` function with vector :math:`[2,3,1]` on :math:`X`, we may reorient
-    the spaces such that :math:`X \in \mathcal{B} \otimes \mathcal{C} \otimes \mathcal{A}`.
+    For spaces :math:`\mathcal{A}, \mathcal{B}`, and :math:`\mathcal{C}` where :math:`\text{dim}(\mathcal{A}) =
+    \text{dim}(\mathcal{B}) = \text{dim}(\mathcal{C}) = 2` we may consider an operator :math:`X \in \mathcal{A} \otimes
+    \mathcal{B} \otimes \mathcal{C}`. Applying the :code:`permute_systems` function with vector :math:`[2,3,1]` on
+    :math:`X`, we may reorient the spaces such that :math:`X \in \mathcal{B} \otimes \mathcal{C} \otimes \mathcal{A}`.
 
     For example, if we define :math:`X \in \mathcal{A} \otimes \mathcal{B} \otimes \mathcal{C}` as
 
@@ -97,8 +90,8 @@ def permute_systems(
             57 & 58 & 59 & 60 & 61 & 62 & 63 & 64
         \end{pmatrix},
 
-    then applying the `permute_systems` function on :math:`X` to obtain
-    :math:`X \in \mathcal{B} \otimes \mathcal{C} \otimes \mathcal{C}` yield the following matrix
+    then applying the `permute_systems` function on :math:`X` to obtain :math:`X \in \mathcal{B} \otimes \mathcal{C}
+    \otimes \mathcal{C}` yield the following matrix
 
     .. math::
         X_{[2, 3, 1]} =
diff --git a/toqito/perms/swap.py b/toqito/perms/swap.py
index a504355ae..f6f50fc25 100644
--- a/toqito/perms/swap.py
+++ b/toqito/perms/swap.py
@@ -1,6 +1,4 @@
 """Swap."""
-
-
 import numpy as np
 
 from toqito.perms import permute_systems
@@ -15,17 +13,15 @@ def swap(
     r"""
     Swap two subsystems within a state or operator.
 
-    Swaps the two subsystems of the vector or matrix :code:`rho`, where the dimensions of the
-    (possibly more than 2) subsystems are given by :code:`dim` and the indices of the two subsystems
-    to be swapped are specified in the 1-by-2 vector :code:`sys`.
+    Swaps the two subsystems of the vector or matrix :code:`rho`, where the dimensions of the (possibly more than 2)
+    subsystems are given by :code:`dim` and the indices of the two subsystems to be swapped are specified in the 1-by-2
+    vector :code:`sys`.
 
-    If :code:`rho` is non-square and not a vector, different row and column dimensions can be
-    specified by putting the row dimensions in the first row of :code:`dim` and the column
-    dimensions in the second row of :code:`dim`.
+    If :code:`rho` is non-square and not a vector, different row and column dimensions can be specified by putting the
+    row dimensions in the first row of :code:`dim` and the column dimensions in the second row of :code:`dim`.
 
-    If :code:`row_only` is set to :code:`True`, then only the rows of :code:`rho` are swapped, but
-    not the columns -- this is equivalent to multiplying :code:`rho` on the left by the
-    corresponding swap operator, but not on the right.
+    If :code:`row_only` is set to :code:`True`, then only the rows of :code:`rho` are swapped, but not the columns --
+    this is equivalent to multiplying :code:`rho` on the left by the corresponding swap operator, but not on the right.
 
     Examples
     ==========
@@ -41,8 +37,7 @@ def swap(
             4 & 8 & 12 & 16
         \end{pmatrix}.
 
-    If we apply the :code:`swap` function provided by :code:`toqito` on :math:`X`, we should obtain
-    the following matrix
+    If we apply the :code:`swap` function provided by :code:`toqito` on :math:`X`, we should obtain the following matrix
 
     .. math::
         \text{Swap}(X) =
@@ -66,9 +61,9 @@ def swap(
      [ 2 10  6 14]
      [ 4 12  8 16]]
 
-    It is also possible to use the :code:`sys` and :code:`dim` arguments, it is possible to specify
-    the system and dimension on which to apply the swap operator. For instance for
-    :code:`sys = [1 ,2]` and :code:`dim = 2` we have that
+    It is also possible to use the :code:`sys` and :code:`dim` arguments, it is possible to specify the system and
+    dimension on which to apply the swap operator. For instance for :code:`sys = [1 ,2]` and :code:`dim = 2` we have
+    that
 
     .. math::
         \text{Swap}(X)_{2, [1, 2]} =
@@ -144,16 +139,9 @@ def swap(
 
     # Verify that the input sys makes sense.
     if any(sys) < 1 or any(sys) > num_sys:
-        val_error = """
-            InvalidSys: The subsystems in `sys` must be between 1 and 
-            `len(dim).` inclusive.
-        """
-        raise ValueError(val_error)
+        raise ValueError("InvalidSys: The subsystems in `sys` must be between 1 and `len(dim).` inclusive.")
     if len(sys) != 2:
-        val_error = """
-            InvalidSys: `sys` must be a vector with exactly two elements.
-        """
-        raise ValueError(val_error)
+        raise ValueError("InvalidSys: `sys` must be a vector with exactly two elements.")
 
     # Swap the indicated subsystems.
     perm = np.array(range(1, num_sys + 1))
diff --git a/toqito/perms/swap_operator.py b/toqito/perms/swap_operator.py
index c1ad47b97..786a5264f 100644
--- a/toqito/perms/swap_operator.py
+++ b/toqito/perms/swap_operator.py
@@ -1,6 +1,4 @@
 """Swap operator."""
-
-
 import numpy as np
 import scipy as sp
 
@@ -11,9 +9,8 @@ def swap_operator(dim: list[int] | int, is_sparse: bool = False) -> np.ndarray:
     r"""
     Produce a unitary operator that swaps two subsystems.
 
-    Provides the unitary operator that swaps two copies of :code:`dim`-dimensional space. If the two
-    subsystems are not of the same dimension, :code:`dim` should be a 1-by-2 vector containing the
-    dimension of the subsystems.
+    Provides the unitary operator that swaps two copies of :code:`dim`-dimensional space. If the two subsystems are not
+    of the same dimension, :code:`dim` should be a 1-by-2 vector containing the dimension of the subsystems.
 
     Examples
     ==========
diff --git a/toqito/perms/symmetric_projection.py b/toqito/perms/symmetric_projection.py
index 795bba722..d452f91e0 100644
--- a/toqito/perms/symmetric_projection.py
+++ b/toqito/perms/symmetric_projection.py
@@ -8,34 +8,29 @@
 from toqito.perms import permutation_operator
 
 
-def symmetric_projection(
-    dim: int, p_val: int = 2, partial: bool = False
-) -> [np.ndarray, scipy.sparse.lil_matrix]:
+def symmetric_projection( dim: int, p_val: int = 2, partial: bool = False ) -> [np.ndarray, scipy.sparse.lil_matrix]:
     r"""
     Produce the projection onto the symmetric subspace [CJKLZ14]_.
 
-    For a complex Euclidean space :math:`\mathcal{X}` and a positive integer :math:`n`, the
-    projection onto the symmetric subspace is given by
-
+    For a complex Euclidean space :math:`\mathcal{X}` and a positive integer :math:`n`, the projection onto the
+    symmetric subspace is given by
+    
     .. math::
         \frac{1}{n!} \sum_{\pi \in S_n} W_{\pi}
 
-    where :math:`W_{\pi}` is the swap operator and where :math:`S_n` is the symmetric group on
-    :math:`n` symbols.
+    where :math:`W_{\pi}` is the swap operator and where :math:`S_n` is the symmetric group on :math:`n` symbols.
 
-    Produces the orthogonal projection onto the symmetric subspace of :code:`p_val` copies of
-    `dim`-dimensional space. If `partial = True`, then the symmetric projection (PS) isn't the
-    orthogonal projection itself, but rather a matrix whose columns form an orthonormal basis for
-    the symmetric subspace (and hence the PS * PS' is the orthogonal projection onto the symmetric
-    subspace).
+    Produces the orthogonal projection onto the symmetric subspace of :code:`p_val` copies of `dim`-dimensional space.
+    If `partial = True`, then the symmetric projection (PS) isn't the orthogonal projection itself, but rather a matrix
+    whose columns form an orthonormal basis for the symmetric subspace (and hence the PS * PS' is the orthogonal
+    projection onto the symmetric subspace).
 
     This function was adapted from the QETLAB package.
 
     Examples
     ==========
 
-    The :math:`2`-dimensional symmetric projection with :math:`p=1` is given as
-    :math:`2`-by-:math:`2` identity matrix
+    The :math:`2`-dimensional symmetric projection with :math:`p=1` is given as :math:`2`-by-:math:`2` identity matrix
 
     .. math::
         \begin{pmatrix}
@@ -82,7 +77,7 @@ def symmetric_projection(
     :param partial: Default value of 0.
     :return: Projection onto the symmetric subspace.
     """
-    dimp = dim**p_val
+    dimp = dim ** p_val
 
     if p_val == 1:
         return np.eye(dim)
diff --git a/toqito/perms/tests/test_antisymmetric_projection.py b/toqito/perms/tests/test_antisymmetric_projection.py
index 3f6d3d050..bb308ad65 100644
--- a/toqito/perms/tests/test_antisymmetric_projection.py
+++ b/toqito/perms/tests/test_antisymmetric_projection.py
@@ -1,37 +1,32 @@
 """Test antisymmetric_projection."""
 import numpy as np
+import pytest
 
 from toqito.perms import antisymmetric_projection
 
 
-def test_antisymmetric_projection_d_2_p_1():
-    """Dimension is 2 and p is equal to 1."""
-    res = antisymmetric_projection(2, 1).todense()
-    expected_res = np.array([[1, 0], [0, 1]])
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
-
-
-def test_antisymmetric_projection_p_larger_than_d():
-    """The `p` value is greater than the dimension `d`."""
-    res = antisymmetric_projection(2, 3).todense()
-    expected_res = np.zeros((8, 8))
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
-
-
-def test_antisymmetric_projection_2():
-    """The dimension is 2."""
-    res = antisymmetric_projection(2).todense()
-    expected_res = np.array([[0, 0, 0, 0], [0, 0.5, -0.5, 0], [0, -0.5, 0.5, 0], [0, 0, 0, 0]])
-
-    bool_mat = np.isclose(res, expected_res)
-    np.testing.assert_equal(np.all(bool_mat), True)
-
-
-def test_antisymmetric_projection_3_3_true():
-    """The `dim` is 3, the `p` is 3, and `partial` is True."""
-    res = antisymmetric_projection(3, 3, True).todense()
-    np.testing.assert_equal(np.isclose(res[5].item(), -0.40824829), True)
+# Create a zero vector of length 27
+anti_proj_3_3_partial = np.zeros((27, 1))
+
+# Set specific indices to -0.40824829 and 0.40824829
+anti_proj_3_3_partial[5] = -0.40824829
+anti_proj_3_3_partial[7] = 0.40824829
+anti_proj_3_3_partial[11] = 0.40824829
+anti_proj_3_3_partial[15] = -0.40824829
+anti_proj_3_3_partial[19] = -0.40824829
+anti_proj_3_3_partial[21] = 0.40824829
+
+
+@pytest.mark.parametrize("dim, p_param, partial, expected_result", [
+    # Dimension is 2 and p is equal to 1.
+    (2, 1, False, np.array([[1, 0], [0, 1]])),
+    # The `p` value is greater than the dimension `d`.
+    (2, 3, False, np.zeros((8, 8))),
+    # The dimension is 2.
+    (2, 2, False, np.array([[0, 0, 0, 0], [0, 0.5, -0.5, 0], [0, -0.5, 0.5, 0], [0, 0, 0, 0]])),
+    # The `dim` is 3, the `p` is 3, and `partial` is True.
+    (3, 3, True, anti_proj_3_3_partial)
+])
+def test_antisymmetric_projection(dim, p_param, partial, expected_result):
+    proj = antisymmetric_projection(dim=dim, p_param=p_param, partial=partial).todense()
+    np.testing.assert_allclose(proj, expected_result)
diff --git a/toqito/rand/random_density_matrix.py b/toqito/rand/random_density_matrix.py
index be6b11737..e6626d772 100644
--- a/toqito/rand/random_density_matrix.py
+++ b/toqito/rand/random_density_matrix.py
@@ -1,6 +1,4 @@
 """Generate random density matrix."""
-
-
 import numpy as np
 
 from toqito.rand import random_unitary
@@ -15,15 +13,13 @@ def random_density_matrix(
     r"""
     Generate a random density matrix.
 
-    Generates a random :code:`dim`-by-:code:`dim` density matrix distributed according to the
-    Hilbert-Schmidt measure. The matrix is of rank <= :code:`k_param` distributed according to the
-    distribution :code:`distance_metric` If :code:`is_real = True`, then all of its entries will be
-    real. The variable :code:`distance_metric` must be one of:
+    Generates a random :code:`dim`-by-:code:`dim` density matrix distributed according to the Hilbert-Schmidt measure.
+    The matrix is of rank <= :code:`k_param` distributed according to the distribution :code:`distance_metric` If
+    :code:`is_real = True`, then all of its entries will be real. The variable :code:`distance_metric` must be one of:
 
         - :code:`haar` (default):
-            Generate a larger pure state according to the Haar measure and
-            trace out the extra dimensions. Sometimes called the
-            Hilbert-Schmidt measure when :code:`k_param = dim`.
+            Generate a larger pure state according to the Haar measure and trace out the extra dimensions. Sometimes
+            called the Hilbert-Schmidt measure when :code:`k_param = dim`.
 
         - :code:`bures`:
             The Bures measure.
@@ -31,8 +27,8 @@ def random_density_matrix(
     Examples
     ==========
 
-    Using :code:`toqito`, we may generate a random complex-valued :math:`n`- dimensional density
-    matrix. For :math:`d=2`, this can be accomplished as follows.
+    Using :code:`toqito`, we may generate a random complex-valued :math:`n`- dimensional density matrix. For
+    :math:`d=2`, this can be accomplished as follows.
 
     >>> from toqito.rand import random_density_matrix
     >>> complex_dm = random_density_matrix(2)
@@ -40,8 +36,8 @@ def random_density_matrix(
     [[0.34903796+0.j       0.4324904 +0.103298j]
      [0.4324904 -0.103298j 0.65096204+0.j      ]]
 
-    We can verify that this is in fact a valid density matrix using the :code:`is_denisty` function
-    from :code:`toqito` as follows
+    We can verify that this is in fact a valid density matrix using the :code:`is_denisty` function from :code:`toqito`
+    as follows
 
     >>> from toqito.matrix_props import is_density
     >>> is_density(complex_dm)
@@ -61,8 +57,8 @@ def random_density_matrix(
     >>> is_density(real_dm)
     True
 
-    By default, the random density operators are constructed using the Haar measure. We can select
-    to generate the random density matrix according to the Bures metric instead as follows.
+    By default, the random density operators are constructed using the Haar measure. We can select to generate the
+    random density matrix according to the Bures metric instead as follows.
 
     >>> from toqito.rand import random_density_matrix
     >>> bures_mat = random_density_matrix(2, distance_metric="bures")
diff --git a/toqito/rand/random_ginibre.py b/toqito/rand/random_ginibre.py
index 254ada793..4aa149f48 100644
--- a/toqito/rand/random_ginibre.py
+++ b/toqito/rand/random_ginibre.py
@@ -2,17 +2,14 @@
 import numpy as np
 
 
-def random_ginibre(
-    dim_n: int,
-    dim_m: int,
-) -> np.ndarray:
+def random_ginibre(dim_n: int, dim_m: int) -> np.ndarray:
     r"""
     Generate a Ginibre random matrix [WIKCIRC]_.
 
     Generates a random :code:`dim_n`-by-:code:`dim_m` Ginibre matrix.
 
-    A *Ginibre random matrix* is a matrix with independent and identically distributed complex
-    standard Gaussian entries.
+    A *Ginibre random matrix* is a matrix with independent and identically distributed complex standard Gaussian
+    entries.
 
     Ginibre random matrices are used in the construction of Wishart-random POVMs [HMN20]_.
 
diff --git a/toqito/rand/random_povm.py b/toqito/rand/random_povm.py
index c4ba06cd1..289f7e0fd 100644
--- a/toqito/rand/random_povm.py
+++ b/toqito/rand/random_povm.py
@@ -9,10 +9,9 @@ def random_povm(dim: int, num_inputs: int, num_outputs: int) -> np.ndarray:
     Examples
     ==========
 
-    We can generate a set of `dim`-by-`dim` POVMs consisting of a specific dimension along with
-    a given number of measurement inputs and measurement outputs. As an example, we can
-    construct a random set of :math:`2`-by-:math:`2` POVMs of dimension with :math:`2` inputs
-    and :math:`2` outputs.
+    We can generate a set of `dim`-by-`dim` POVMs consisting of a specific dimension along with a given number of
+    measurement inputs and measurement outputs. As an example, we can construct a random set of :math:`2`-by-:math:`2`
+    POVMs of dimension with :math:`2` inputs and :math:`2` outputs.
 
     >>> from toqito.rand import random_povm
     >>> import numpy as np
@@ -29,8 +28,8 @@ def random_povm(dim: int, num_inputs: int, num_outputs: int) -> np.ndarray:
      [[ 0.452533  +0.j,  0.547467  +0.j],
       [ 0.34692158+0.j,  0.65307842+0.j]]]]
 
-    We can verify that this constitutes a valid set of POVM elements as checking that these
-    operators all sum to the identity operator.
+    We can verify that this constitutes a valid set of POVM elements as checking that these operators all sum to the
+    identity operator.
 
     >>> np.round(povms[:, :, 0, 0] + povms[:, :, 0, 1])
     [[1.+0.j, 0.+0.j],
diff --git a/toqito/rand/random_state_vector.py b/toqito/rand/random_state_vector.py
index 7e05f4d19..f7bbca79c 100644
--- a/toqito/rand/random_state_vector.py
+++ b/toqito/rand/random_state_vector.py
@@ -1,15 +1,11 @@
 """Generate random state vector."""
-
-
 import numpy as np
 
 from toqito.perms import swap
 from toqito.states import max_entangled
 
 
-def random_state_vector(
-    dim: list[int] | int, is_real: bool = False, k_param: int = 0
-) -> np.ndarray:
+def random_state_vector(dim: list[int] | int, is_real: bool = False, k_param: int = 0) -> np.ndarray:
     r"""Generate a random pure state vector.
 
     Examples
@@ -24,8 +20,8 @@ def random_state_vector(
     [[0.50993973+0.15292408j],
      [0.27787332+0.79960122j]]
 
-    We can verify that this is in fact a valid state vector by computing the corresponding density
-    matrix of the vector and checking if the density matrix is pure.
+    We can verify that this is in fact a valid state vector by computing the corresponding density matrix of the vector
+    and checking if the density matrix is pure.
 
     >>> from toqito.state_props import is_pure
     >>> dm = vec.conj().T * vec
diff --git a/toqito/rand/random_unitary.py b/toqito/rand/random_unitary.py
index 692951eb0..722381ee1 100644
--- a/toqito/rand/random_unitary.py
+++ b/toqito/rand/random_unitary.py
@@ -1,6 +1,4 @@
 """Generate random unitary."""
-
-
 import numpy as np
 
 
@@ -75,6 +73,9 @@ def random_unitary(dim: list[int] | int, is_real: bool = False) -> np.ndarray:
     """
     if isinstance(dim, int):
         dim = [dim, dim]
+    
+    if dim[0] != dim[1]:
+        raise ValueError("Unitary matrix must be square.")
 
     # Construct the Ginibre ensemble.
     gin = np.random.rand(dim[0], dim[1])
diff --git a/toqito/rand/tests/test_random_density_matrix.py b/toqito/rand/tests/test_random_density_matrix.py
index ac09afc94..f0862eddc 100644
--- a/toqito/rand/tests/test_random_density_matrix.py
+++ b/toqito/rand/tests/test_random_density_matrix.py
@@ -1,35 +1,49 @@
 """Test random_density_matrix."""
 import numpy as np
+import pytest
 
-from toqito.matrix_props import is_density
+from toqito.matrix_props import is_density, is_positive_semidefinite
 from toqito.rand import random_density_matrix
 
 
-def test_random_density_not_real():
-    """Generate random non-real density matrix."""
-    mat = random_density_matrix(2)
-    np.testing.assert_equal(is_density(mat), True)
-
-
-def test_random_density_real():
-    """Generate random real density matrix."""
-    mat = random_density_matrix(2, True)
-    np.testing.assert_equal(is_density(mat), True)
-
-
-def test_random_density_not_real_bures():
-    """Random non-real density matrix according to Bures metric."""
-    mat = random_density_matrix(2, distance_metric="bures")
-    np.testing.assert_equal(is_density(mat), True)
-
-
-def test_random_density_not_real_k_param():
-    """Generate random non-real density matrix wih k_param."""
-    mat = random_density_matrix(2, distance_metric="bures")
-    np.testing.assert_equal(is_density(mat), True)
-
-
-def test_random_density_not_real_all_params():
-    """Generate random non-real density matrix all params."""
-    mat = random_density_matrix(2, True, 2, "haar")
-    np.testing.assert_equal(is_density(mat), True)
+@pytest.mark.parametrize("dim, is_real, k_param, distance_metric", [
+    # Generate random non-real density matrix.
+    (2, False, None, "haar"),
+    # Generate random real density matrix.
+    (2, True, None, "haar"),
+    # Random non-real density matrix according to Bures metric.
+    (2, False, None, "bures"),
+    # Generate random non-real density matrix all params.
+    (2, True, 2, "haar"),
+])
+def test_random_density(dim, is_real, k_param, distance_metric):
+    if k_param == dim:
+        mat = random_density_matrix(
+            dim=dim, is_real=is_real, k_param=k_param, distance_metric=distance_metric
+        )
+        np.testing.assert_equal(is_density(mat), True)
+
+
+@pytest.mark.parametrize("dim, is_real, distance_metric", [
+    (2, False, "haar"),
+    (2, True, "haar"),
+    (3, False, "bures"),
+    (3, True, "bures")
+])
+def test_random_density_matrix(dim, is_real, distance_metric):
+    dm = random_density_matrix(dim, is_real, distance_metric=distance_metric)
+
+    # Check if the matrix is a valid density matrix
+    assert is_density(dm), "Matrix should be a valid density matrix"
+    
+    # Check if the matrix is positive semidefinite
+    assert is_positive_semidefinite(dm), "Matrix should be positive semidefinite"
+
+    # Check if the trace is 1
+    assert np.isclose(np.trace(dm), 1), "Trace of the matrix should be 1"
+
+    # Check if the matrix is real or complex as expected
+    if is_real:
+        assert np.all(np.isreal(dm)), "Matrix should be real"
+    else:
+        assert np.any(np.iscomplex(dm)), "Matrix should be complex"
\ No newline at end of file
diff --git a/toqito/rand/tests/test_random_ginibre.py b/toqito/rand/tests/test_random_ginibre.py
index 556b1f507..b1947d678 100644
--- a/toqito/rand/tests/test_random_ginibre.py
+++ b/toqito/rand/tests/test_random_ginibre.py
@@ -1,11 +1,27 @@
 """Test random_ginibre."""
 import numpy as np
+import pytest
 
 from toqito.rand import random_ginibre
 
 
-def test_random_ginibre_dims():
+@pytest.mark.parametrize("dim_n", range(0, 8))
+@pytest.mark.parametrize("dim_m", range(0, 8))
+def test_random_ginibre_dims(dim_n, dim_m):
     """Generate random Ginibre matrix and check proper dimensions."""
-    gin_mat = random_ginibre(2, 2)
-    np.testing.assert_equal(gin_mat.shape[0], 2)
-    np.testing.assert_equal(gin_mat.shape[1], 2)
+    gin_mat = random_ginibre(dim_n, dim_m)
+    np.testing.assert_equal(gin_mat.shape, (dim_n, dim_m))
+
+
+@pytest.mark.parametrize("dim_n, dim_m", [
+    # Negative dim_n.
+    (-1, 3),
+    # Negative dim_m.
+    (3, -2),
+    # Negative dim_n and dim_m.
+    (-4, -2),
+])
+def test_random_ginibre_negative_dims(dim_n, dim_m):
+    """Negative dimensions are not allowed."""
+    with pytest.raises(ValueError, match="negative dimensions are not allowed"):
+        random_ginibre(dim_n, dim_m)
diff --git a/toqito/rand/tests/test_random_povm.py b/toqito/rand/tests/test_random_povm.py
index 3842ffa08..3a28ad6f1 100644
--- a/toqito/rand/tests/test_random_povm.py
+++ b/toqito/rand/tests/test_random_povm.py
@@ -1,32 +1,24 @@
 """Test random_povm."""
-import unittest
-
 import numpy as np
+import pytest
 
 from toqito.rand import random_povm
 
-
-class TestRandomPOVM(unittest.TestCase):
-    """Unit test for NonlocalGame."""
-
-    def test_random_povm_unitary_not_real(self):
-        """Generate random POVMs and check that they sum to the identity."""
-        dim, num_inputs, num_outputs = 2, 2, 2
-        povms = random_povm(dim, num_inputs, num_outputs)
-
-        self.assertEqual(povms.shape, (dim, dim, num_inputs, num_outputs))
-
-        np.testing.assert_allclose(
-            povms[:, :, 0, 0] + povms[:, :, 0, 1], np.identity(dim), atol=1e-7
-        )
-
-    def test_random_povm_uneven_dimensions(self):
-        """Generate random POVMs of uneven dimensions"""
-        dim, num_inputs, num_outputs = 2, 3, 4
-        povms = random_povm(dim, num_inputs, num_outputs)
-
-        self.assertEqual(povms.shape, (dim, dim, num_inputs, num_outputs))
-
-        for i in range(num_inputs):
-            povm_sum = np.sum(povms[:, :, i, :], axis=-1)
-            np.testing.assert_allclose(povm_sum, np.identity(dim), atol=1e-7)
+@pytest.mark.parametrize("dim", range(2, 8))
+@pytest.mark.parametrize("num_inputs", range(2, 8))
+@pytest.mark.parametrize("num_outputs", range(2, 8))
+def test_random_povm_dimensions(dim, num_inputs, num_outputs):
+    """Verify that the output has the correct shape as specified by the input parameters."""
+    povms = random_povm(dim=dim, num_inputs=num_inputs, num_outputs=num_outputs)
+    np.testing.assert_equal(povms.shape, (dim, dim, num_inputs, num_outputs))
+
+
+@pytest.mark.parametrize("dim", range(2, 8))
+@pytest.mark.parametrize("num_inputs", range(2, 8))
+@pytest.mark.parametrize("num_outputs", range(2, 8))
+def test_random_povm_validity(dim, num_inputs, num_outputs):
+    """Each set of POVMs for a given input sums up to the identity matrix."""
+    povms = random_povm(dim=dim, num_inputs=num_inputs, num_outputs=num_outputs)
+    for i in range(num_inputs):
+        sum_povms = sum(povms[:, :, i, j] for j in range(num_outputs))
+        assert np.allclose(sum_povms, np.identity(dim))
diff --git a/toqito/rand/tests/test_random_state_vector.py b/toqito/rand/tests/test_random_state_vector.py
index 512bbdf3a..021d4e9f0 100644
--- a/toqito/rand/tests/test_random_state_vector.py
+++ b/toqito/rand/tests/test_random_state_vector.py
@@ -1,40 +1,24 @@
 """Test random_state_vector."""
 import numpy as np
+import pytest
 
 from toqito.rand import random_state_vector
-from toqito.state_props import is_pure, purity
-
-
-def test_random_state_vector_complex_state_purity():
-    """Check that complex state vector from random state vector is pure."""
-    vec = random_state_vector(2)
-    mat = vec * vec.conj().T
-    np.testing.assert_equal(is_pure(mat), True)
-
-
-def test_random_state_vector_complex_state_purity_k_param():
-    """Check that complex state vector with k_param > 0."""
-    vec = random_state_vector(2, False, 1).reshape(-1, 1)
-    mat = vec * vec.conj().T
-    np.testing.assert_equal(np.isclose(purity(mat), 1), True)
-
-
-def test_random_state_vector_complex_state_purity_k_param_dim_list():
-    """Check that complex state vector with k_param > 0 and dim list."""
-    vec = random_state_vector([2, 2], False, 1).reshape(-1, 1)
-    mat = vec * vec.conj().T
-    np.testing.assert_equal(np.isclose(purity(mat), 1), True)
-
-
-def test_random_state_vector_real_state_purity_with_k_param():
-    """Check that real state vector with k_param > 0."""
-    vec = random_state_vector(2, True, 1).reshape(-1, 1)
-    mat = vec * vec.conj().T
-    np.testing.assert_equal(np.isclose(purity(mat), 1), True)
-
-
-def test_random_state_vector_real_state_purity():
-    """Check that real state vector from random state vector is pure."""
-    vec = random_state_vector(2, True).reshape(-1, 1)
-    mat = vec * vec.conj().T
-    np.testing.assert_equal(np.isclose(purity(mat), 1), True)
+from toqito.state_props import is_pure
+
+@pytest.mark.parametrize("dim, is_real, k_param", [
+    # Check that complex state vector from random state vector is pure.
+    (2, False, 0),
+    # Check that complex state vector with k_param > 0.
+    (2, False, 1),
+    # Check that complex state vector with k_param > 0 and dim list.
+    ([2, 2], False, 1),
+    # Check that real state vector with k_param > 0.
+    (2, True, 1),
+    # Check that real state vector from random state vector is pure.
+    (2, True, 0),
+])
+def test_random_state_vector(dim, is_real, k_param):
+    # We expect the density matrix of any random state vector to be pure.
+    vec = random_state_vector(dim=dim, is_real=is_real, k_param=k_param).reshape(-1, 1)
+    mat = vec @ vec.conj().T
+    assert is_pure(mat)
diff --git a/toqito/rand/tests/test_random_unitary.py b/toqito/rand/tests/test_random_unitary.py
index f3ab7242c..87ea53cdf 100644
--- a/toqito/rand/tests/test_random_unitary.py
+++ b/toqito/rand/tests/test_random_unitary.py
@@ -1,23 +1,36 @@
 """Test random_unitary."""
 import numpy as np
+import pytest
 
 from toqito.matrix_props import is_unitary
 from toqito.rand import random_unitary
 
 
-def test_random_unitary_not_real():
-    """Generate random non-real unitary matrix."""
-    mat = random_unitary(2)
-    np.testing.assert_equal(is_unitary(mat), True)
+@pytest.mark.parametrize("dim", range(2, 8))
+@pytest.mark.parametrize("is_real", [True, False])
+def test_random_unitary_int_dim(dim, is_real):
+    mat = random_unitary(
+        dim=dim, is_real=is_real
+    )
+    assert is_unitary(mat)
 
 
-def test_random_unitary_real():
-    """Generate random real unitary matrix."""
-    mat = random_unitary(2, True)
-    np.testing.assert_equal(is_unitary(mat), True)
+@pytest.mark.parametrize("dim_n", range(2, 8))
+@pytest.mark.parametrize("dim_m", range(2, 8))
+@pytest.mark.parametrize("is_real", [True, False])
+def test_random_unitary_list_dims(dim_n, dim_m, is_real):
+    if dim_n == dim_m:
+        mat = random_unitary(
+            dim=[dim_n, dim_m], is_real=is_real
+        )
+        assert is_unitary(mat)
 
 
-def test_random_unitary_vec_dim():
-    """Generate random non-real unitary matrix."""
-    mat = random_unitary([4, 4], True)
-    np.testing.assert_equal(is_unitary(mat), True)
+@pytest.mark.parametrize("dim_n", range(2, 8))
+@pytest.mark.parametrize("dim_m", range(2, 8))
+@pytest.mark.parametrize("is_real", [True, False])
+def test_non_square_dims(dim_n, dim_m, is_real):
+    """Unitary matrices must be square are not allowed."""
+    if dim_n != dim_m:
+        with pytest.raises(ValueError, match="Unitary matrix must be square."):
+            random_unitary(dim=[dim_n, dim_m], is_real=is_real)
diff --git a/toqito/state_metrics/bures_angle.py b/toqito/state_metrics/bures_angle.py
index f77366ede..6979f6778 100644
--- a/toqito/state_metrics/bures_angle.py
+++ b/toqito/state_metrics/bures_angle.py
@@ -8,17 +8,14 @@ def bures_angle(rho_1: np.ndarray, rho_2: np.ndarray, decimals: int = 10) -> flo
     r"""
     Compute the Bures angle of two density matrices [WikBures]_.
 
-    Calculate the Bures angle between two density matrices :code:`rho_1` and :code:`rho_2`
-    defined by:
+    Calculate the Bures angle between two density matrices :code:`rho_1` and :code:`rho_2` defined by:
 
     .. math::
         \arccos{\sqrt{F (\rho_1, \rho_2)}}
 
-    where :math:`F(\cdot)` denotes the fidelity between :math:`\rho_1` and
-    :math:`\rho_2`. The return is a value between :math:`0` and :math:`\pi / 2`,
-    with :math:`0` corresponding to matrices :code:`rho_1 = rho_2` and
-    :math:`\pi / 2` corresponding to the case :code:`rho_1` and :code:`rho_2`
-    with orthogonal support.
+    where :math:`F(\cdot)` denotes the fidelity between :math:`\rho_1` and :math:`\rho_2`. The return is a value between
+    :math:`0` and :math:`\pi / 2`, with :math:`0` corresponding to matrices :code:`rho_1 = rho_2` and :math:`\pi / 2`
+    corresponding to the case :code:`rho_1` and :code:`rho_2` with orthogonal support.
 
     Examples
     ==========
@@ -38,8 +35,8 @@ def bures_angle(rho_1: np.ndarray, rho_2: np.ndarray, decimals: int = 10) -> flo
                          1 & 0 & 0 & 1
                        \end{pmatrix} \in \text{D}(\mathcal{X}).
 
-    In the event where we calculate the Bures angle between states that are identical, we
-    should obtain the value of :math:`0`. This can be observed in :code:`toqito` as follows.
+    In the event where we calculate the Bures angle between states that are identical, we should obtain the value of
+    :math:`0`. This can be observed in :code:`toqito` as follows.
 
     >>> from toqito.state_metrics import bures_angle
     >>> import numpy as np
@@ -64,7 +61,7 @@ def bures_angle(rho_1: np.ndarray, rho_2: np.ndarray, decimals: int = 10) -> flo
     :param decimals: Number of decimal places to round to (default 10).
     :return: The Bures angle between :code:`rho_1` and :code:`rho_2`.
     """
-    # Perform error checking
+    # Perform error checking.
     if not np.all(rho_1.shape == rho_2.shape):
         raise ValueError("InvalidDim: `rho_1` and `rho_2` must be matrices of the same size.")
     # Round fidelity to only 10 decimals to avoid error when :code:`rho_1 = rho_2`.
diff --git a/toqito/state_metrics/bures_distance.py b/toqito/state_metrics/bures_distance.py
index c7286e3d9..3ae138d87 100644
--- a/toqito/state_metrics/bures_distance.py
+++ b/toqito/state_metrics/bures_distance.py
@@ -8,16 +8,14 @@ def bures_distance(rho_1: np.ndarray, rho_2: np.ndarray, decimals: int = 10) ->
     r"""
     Compute the Bures distance of two density matrices [WikBures]_.
 
-    Calculate the Bures distance between two density matrices :code:`rho_1` and :code:`rho_2`
-    defined by:
+    Calculate the Bures distance between two density matrices :code:`rho_1` and :code:`rho_2` defined by:
 
     .. math::
         \sqrt{2 (1 - F(\rho_1, \rho_2))},
 
-    where :math:`F(\cdot)` denotes the fidelity between :math:`\rho_1` and :math:`\rho_2`. The
-    return is a value between :math:`0` and :math:`\sqrt{2}`,with :math:`0` corresponding to
-    matrices: :code:`rho_1 = rho_2` and :math:`\sqrt{2}` corresponding to the case: :code:`rho_1`
-    and :code:`rho_2` with orthogonal support.
+    where :math:`F(\cdot)` denotes the fidelity between :math:`\rho_1` and :math:`\rho_2`. The return is a value between
+    :math:`0` and :math:`\sqrt{2}`,with :math:`0` corresponding to matrices: :code:`rho_1 = rho_2` and :math:`\sqrt{2}`
+    corresponding to the case: :code:`rho_1` and :code:`rho_2` with orthogonal support.
 
     Examples
     ==========
@@ -37,8 +35,8 @@ def bures_distance(rho_1: np.ndarray, rho_2: np.ndarray, decimals: int = 10) ->
                          1 & 0 & 0 & 1
                        \end{pmatrix} \in \text{D}(\mathcal{X}).
 
-    In the event where we calculate the Bures distance between states that are identical, we
-    should obtain the value of :math:`0`. This can be observed in :code:`toqito` as follows.
+    In the event where we calculate the Bures distance between states that are identical, we should obtain the value of
+    :math:`0`. This can be observed in :code:`toqito` as follows.
 
     >>> from toqito.state_metrics import bures_distance
     >>> import numpy as np
diff --git a/toqito/state_metrics/fidelity.py b/toqito/state_metrics/fidelity.py
index f951a9835..860cf3bea 100644
--- a/toqito/state_metrics/fidelity.py
+++ b/toqito/state_metrics/fidelity.py
@@ -10,15 +10,14 @@ def fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     r"""
     Compute the fidelity of two density matrices [WikFid]_.
 
-    Calculate the fidelity between the two density matrices :code:`rho` and :code:`sigma`, defined
-    by:
+    Calculate the fidelity between the two density matrices :code:`rho` and :code:`sigma`, defined by:
 
     .. math::
         ||\sqrt(\rho) \sqrt(\sigma)||_1,
 
-    where :math:`|| \cdot ||_1` denotes the trace norm. The return is a value between :math:`0` and
-    :math:`1`, with :math:`0` corresponding to matrices :code:`rho` and :code:`sigma` with
-    orthogonal support, and :math:`1` corresponding to the case :code:`rho = sigma`.
+    where :math:`|| \cdot ||_1` denotes the trace norm. The return is a value between :math:`0` and :math:`1`, with
+    :math:`0` corresponding to matrices :code:`rho` and :code:`sigma` with orthogonal support, and :math:`1`
+    corresponding to the case :code:`rho = sigma`.
 
     Examples
     ==========
@@ -38,8 +37,8 @@ def fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
                          1 & 0 & 0 & 1
                        \end{pmatrix} \in \text{D}(\mathcal{X}).
 
-    In the event where we calculate the fidelity between states that are identical, we should obtain
-    the value of :math:`1`. This can be observed in :code:`toqito` as follows.
+    In the event where we calculate the fidelity between states that are identical, we should obtain the value of
+    :math:`1`. This can be observed in :code:`toqito` as follows.
 
     >>> from toqito.state_metrics import fidelity
     >>> import numpy as np
@@ -67,12 +66,9 @@ def fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     if not np.all(rho.shape == sigma.shape):
         raise ValueError("InvalidDim: `rho` and `sigma` must be matrices of the same size.")
 
-    # If `rho` or `sigma` is a cvxpy variable then compute fidelity via
-    # semidefinite programming, so that this function can be used in the
-    # objective function or constraints of other cvxpy optimization problems.
-    if isinstance(rho, cvxpy.atoms.affine.vstack.Vstack) or isinstance(
-        sigma, cvxpy.atoms.affine.vstack.Vstack
-    ):
+    # If `rho` or `sigma` is a cvxpy variable then compute fidelity via semidefinite programming, so that this function
+    # can be used in the objective function or constraints of other cvxpy optimization problems.
+    if isinstance(rho, cvxpy.atoms.affine.vstack.Vstack) or isinstance(sigma, cvxpy.atoms.affine.vstack.Vstack):
         z_var = cvxpy.Variable(rho.shape, complex=True)
         objective = cvxpy.Maximize(cvxpy.real(cvxpy.trace(z_var + z_var.H)))
         constraints = [cvxpy.bmat([[rho, z_var], [z_var.H, sigma]]) >> 0]
@@ -83,8 +79,7 @@ def fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     if not is_density(rho) or not is_density(sigma):
         raise ValueError("Fidelity is only defined for density operators.")
 
-    # If `rho` or `sigma` are *not* cvxpy variables, compute fidelity normally,
-    # since this is much faster.
+    # If `rho` or `sigma` are *not* cvxpy variables, compute fidelity normally, since this is much faster.
     sq_rho = scipy.linalg.sqrtm(rho)
     sq_fid = scipy.linalg.sqrtm(sq_rho @ sigma @ sq_rho)
     return np.real(np.trace(sq_fid))
diff --git a/toqito/state_metrics/fidelity_of_separability.py b/toqito/state_metrics/fidelity_of_separability.py
index e7206b082..70059ffaf 100644
--- a/toqito/state_metrics/fidelity_of_separability.py
+++ b/toqito/state_metrics/fidelity_of_separability.py
@@ -1,9 +1,7 @@
 """Add function for fidelity of separability as defined in [Phil23]_.
 
-The constrainsts for this function are positive partial transpose (PPT)
-& k-extendible states.
+The constraints for this function are positive partial transpose (PPT) & k-extendible states.
 """
-
 import numpy as np
 import picos
 
@@ -22,24 +20,20 @@ def fidelity_of_separability(
     r"""
     Define the first benchmark introduced in Appendix H of [Phil23]_.
 
-    If you would like to instead use the benchmark introduced in Appendix I,
-    go to :obj:`toqito.channel_metrics.fidelity_of_separability`.
+    If you would like to instead use the benchmark introduced in Appendix I, go to
+    :obj:`toqito.channel_metrics.fidelity_of_separability`.
 
-    In [Phil23]_ a variational quantum algorithm (VQA) is introduced to test
-    the separability of a general bipartite state. The algorithm utilizes
-    quantum steering between two separated systems such that the separability
-    of the state is quantified.
+    In [Phil23]_ a variational quantum algorithm (VQA) is introduced to test the separability of a general bipartite
+    state. The algorithm utilizes quantum steering between two separated systems such that the separability of the state
+    is quantified.
 
-    Due to the limitations of currently available quantum computers, two
-    optimization semidefinite programs (SDP) benchmarks were introduced to
-    maximize the fidelity of separability subject to some state constraints
-    (Positive Partial Transpose (PPT), symmetric extensions (k-extendibility
-    ) [Hay12]_ ) This function approximites the fidelity of separability by
-    maximizing over PPT states & k-extendible states i.e. an optimization
-    problem over states [TBWat18]_ .
+    Due to the limitations of currently available quantum computers, two optimization semidefinite programs (SDP)
+    benchmarks were introduced to maximize the fidelity of separability subject to some state constraints (Positive
+    Partial Transpose (PPT), symmetric extensions (k-extendibility ) [Hay12]_ ) This function approximates the fidelity
+    of separability by maximizing over PPT states & k-extendible states i.e. an optimization problem over states
+    [TBWat18]_ .
 
-    The following expression (Equation (H2) from [Phil23]_ ) defines the
-    constraints for approxiamting
+    The following expression (Equation (H2) from [Phil23]_ ) defines the constraints for approximating
     :math:`\sqrt{\widetilde{F}_s^1}(\rho_{AB}) {:}=`
 
     .. math::
@@ -64,11 +58,10 @@ def fidelity_of_separability(
     approximated but this function returns
     :math:`\widetilde{F}_s^1(\rho_{AB})`.
 
-    :math:`\operatorname{Re}[\operatorname{Tr}[X_{AB}]]` is the maximization
-    problem subject to PPT & k-extendibile state constraints.
+    :math:`\operatorname{Re}[\operatorname{Tr}[X_{AB}]]` is the maximization problem subject to PPT & k-extendibile
+    state constraints.
 
-    Here, :math:`\mathcal{L}(\mathcal{H}_{AB})` is the space of linear
-    operators over space :math:`\mathcal{H}_{AB}`.
+    Here, :math:`\mathcal{L}(\mathcal{H}_{AB})` is the space of linear operators over space :math:`\mathcal{H}_{AB}`.
 
     :math:`\sigma_{AB^{k}}` is a k-extension of :math:`\rho_{AB}`.
 
@@ -80,11 +73,10 @@ def fidelity_of_separability(
 
     Examples
     ==========
-    Let's consider a density matrix of a state that we know is pure and
-    separable; :math:`|00 \rangle = |0 \rangle \otimes |0 \rangle`.
+    Let's consider a density matrix of a state that we know is pure and separable; :math:`|00 \rangle = |0 \rangle
+    \otimes |0 \rangle`.
 
-    The expected approximation of fidelity of separability is the maximum
-    value possible i.e. very close to 1.
+    The expected approximation of fidelity of separability is the maximum value possible i.e. very close to 1.
 
     .. math::
         \rho_{AB} = |00 \rangle \langle 00|
@@ -102,7 +94,6 @@ def fidelity_of_separability(
 
     >>> 0.9999999998278968
 
-
     References
     ==========
     .. [Hay12] Hayden, Patrick et.al.
@@ -151,18 +142,17 @@ def fidelity_of_separability(
     if not is_separable(input_state_rho):
         raise ValueError("Provided input state is entangled.")
 
-    # Infer the dimension of Alice and Bob's system.
-    # subsystem-dimensions in rho_AB
+    # Infer the dimension of Alice and Bob's system. subsystem-dimensions in rho_AB
     dim_a, dim_b = input_state_rho_dims
 
-    # Extend the number of dimensions based on the level `k`.
-    # new dims for AB with k-extendibility in subsystem B
+    # Extend the number of dimensions based on the level `k`. new dims for AB with k-extendibility in subsystem B
     dim_direct_sum_ab_k = [dim_a] + [dim_b] * (k)
     # new dims for a linear op acting on the space of sigma_ab_k
     dim_op_sigma_ab_k = dim_a * dim_b**k
 
     # A list of the symmetrically extended subsystems based on the level `k`.
     sub_sys_ext = list(range(2, 2 + k - 1))
+
     # unitary permutation operator in B1,B2,...,Bk
     permutation_op = symmetric_projection(dim_b, k)
 
@@ -187,8 +177,7 @@ def fidelity_of_separability(
 
     # k-extendible constraint:
     problem.add_constraint(
-        (picos.I(dim_a) @ permutation_op) * sigma_ab_k * (picos.I(dim_a) @ permutation_op)
-        == sigma_ab_k
+        (picos.I(dim_a) @ permutation_op) * sigma_ab_k * (picos.I(dim_a) @ permutation_op) == sigma_ab_k
     )
 
     # PPT constraint:
diff --git a/toqito/state_metrics/helstrom_holevo.py b/toqito/state_metrics/helstrom_holevo.py
index 00e72f999..aebc58aaf 100644
--- a/toqito/state_metrics/helstrom_holevo.py
+++ b/toqito/state_metrics/helstrom_holevo.py
@@ -8,8 +8,8 @@ def helstrom_holevo(rho: np.ndarray, sigma: np.ndarray) -> float:
     r"""
     Compute the Helstrom-Holevo distance between density matrices [WikHeHo]_.
 
-    In general, the best success probability to discriminate two mixed states represented by
-    :math:`\rho` and :math:`\sigma` is given by [WikHeHo]_.
+    In general, the best success probability to discriminate two mixed states represented by :math:`\rho` and
+    :math:`\sigma` is given by [WikHeHo]_.
 
     .. math::
          \frac{1}{2}+\frac{1}{2} \left(\frac{1}{2} \left|\rho - \sigma \right|_1\right).
diff --git a/toqito/state_metrics/hilbert_schmidt.py b/toqito/state_metrics/hilbert_schmidt.py
index 4a13211c1..8653bd9c6 100644
--- a/toqito/state_metrics/hilbert_schmidt.py
+++ b/toqito/state_metrics/hilbert_schmidt.py
@@ -8,8 +8,7 @@ def hilbert_schmidt(rho: np.ndarray, sigma: np.ndarray) -> float:
     r"""
     Compute the Hilbert-Schmidt distance between two states [WikHS]_.
 
-    The Hilbert-Schmidt distance between density operators :math:`\rho` and :math:`\sigma` is
-    defined as
+    The Hilbert-Schmidt distance between density operators :math:`\rho` and :math:`\sigma` is defined as
 
     .. math::
         D_{\text{HS}}(\rho, \sigma) = \text{Tr}((\rho - \sigma)^2) = \left\lVert \rho - \sigma
@@ -18,8 +17,8 @@ def hilbert_schmidt(rho: np.ndarray, sigma: np.ndarray) -> float:
     Examples
     ==========
 
-    One may consider taking the Hilbert-Schmidt distance between two Bell states. In :code:`toqito`,
-    one may accomplish this as
+    One may consider taking the Hilbert-Schmidt distance between two Bell states. In :code:`toqito`, one may accomplish
+    this as
 
     >>> from toqito.states import bell
     >>> from toqito.state_metrics import hilbert_schmidt
diff --git a/toqito/state_metrics/hilbert_schmidt_inner_product.py b/toqito/state_metrics/hilbert_schmidt_inner_product.py
index 9868720cf..05530666a 100644
--- a/toqito/state_metrics/hilbert_schmidt_inner_product.py
+++ b/toqito/state_metrics/hilbert_schmidt_inner_product.py
@@ -6,23 +6,20 @@ def hilbert_schmidt_inner_product(a_mat: np.ndarray, b_mat: np.ndarray) -> compl
     r"""
     Compute the Hilbert-Schmidt inner product between two matrices [WikHSO]_.
 
-    The Hilbert-Schmidt inner product between :code:`a_mat` and :code:`b_mat` is
-    defined as
+    The Hilbert-Schmidt inner product between :code:`a_mat` and :code:`b_mat` is defined as
 
     .. math::
 
         HS = (A|B) = Tr[A^\dagger B]
 
-    where :math:`|B\rangle = \text{vec}(B)` and :math:`\langle A|` is the dual
-    vector to :math:`|A \rangle`.
+    where :math:`|B\rangle = \text{vec}(B)` and :math:`\langle A|` is the dual vector to :math:`|A \rangle`.
 
     Note: This function has been adapted from [Rigetti21]_.
 
     Examples
     ==========
 
-    One may consider taking the Hilbert-Schmidt distance between two Hadamard
-    matrices.
+    One may consider taking the Hilbert-Schmidt distance between two Hadamard matrices.
 
     >>> from toqito.matrices import hadamard
     >>> from toqito.matrices import pauli
diff --git a/toqito/state_metrics/matsumoto_fidelity.py b/toqito/state_metrics/matsumoto_fidelity.py
index d24c63b4a..a6ec2218b 100644
--- a/toqito/state_metrics/matsumoto_fidelity.py
+++ b/toqito/state_metrics/matsumoto_fidelity.py
@@ -10,8 +10,7 @@ def matsumoto_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     r"""
     Compute the Matsumoto fidelity of two density matrices [Mat10]_.
 
-    Calculate the Matsumoto fidelity between the two density matrices :code:`rho`
-    and :code:`sigma`, defined by:
+    Calculate the Matsumoto fidelity between the two density matrices :code:`rho` and :code:`sigma`, defined by:
 
     .. math::
         \mathrm{tr}(\rho\#\sigma),
@@ -26,9 +25,9 @@ def matsumoto_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     .. math::
         \rho\#\sigma = \lim_{\epsilon\to0}(\rho+\epsilon\mathbb{I})\#(+\epsilon\mathbb{I}).
 
-    The return is a value between :math:`0` and :math:`1`, with :math:`0` corresponding to matrices
-    :code:`rho` and :code:`sigma` with orthogonal support, and :math:`1` corresponding to the case
-    :code:`rho = sigma`. The Matsumoto fidelity is a lower bound for the fidelity.
+    The return is a value between :math:`0` and :math:`1`, with :math:`0` corresponding to matrices :code:`rho` and
+    :code:`sigma` with orthogonal support, and :math:`1` corresponding to the case :code:`rho = sigma`. The Matsumoto
+    fidelity is a lower bound for the fidelity.
 
     Examples
     ==========
@@ -48,8 +47,8 @@ def matsumoto_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
                          1 & 0 & 0 & 1
                        \end{pmatrix} \in \text{D}(\mathcal{X}).
 
-    In the event where we calculate the Matsumoto fidelity between states that are identical,
-    we should obtain the value of :math:`1`. This can be observed in :code:`toqito` as follows.
+    In the event where we calculate the Matsumoto fidelity between states that are identical, we should obtain the value
+    of :math:`1`. This can be observed in :code:`toqito` as follows.
 
     >>> from toqito.state_metrics import matsumoto_fidelity
     >>> import numpy as np
diff --git a/toqito/state_metrics/sub_fidelity.py b/toqito/state_metrics/sub_fidelity.py
index 320bb491b..a1a087796 100644
--- a/toqito/state_metrics/sub_fidelity.py
+++ b/toqito/state_metrics/sub_fidelity.py
@@ -14,8 +14,8 @@ def sub_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
         E(\rho, \sigma) = \text{Tr}(\rho \sigma) +
         \sqrt{2 \left[ \text{Tr}(\rho \sigma)^2 - \text{Tr}(\rho \sigma \rho \sigma) \right]},
 
-    where :math:`\sigma` and :math:`\rho` are density matrices. The sub-fidelity serves as an lower
-    bound for the fidelity.
+    where :math:`\sigma` and :math:`\rho` are density matrices. The sub-fidelity serves as an lower bound for the
+    fidelity.
 
     Examples
     ==========
@@ -29,8 +29,8 @@ def sub_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
         \sigma = \frac{1}{8}|0 \rangle \langle 0| +
                  \frac{7}{8}|1 \rangle \langle 1|.
 
-    Calculating the fidelity between the states :math:`\rho` and :math:`\sigma` as
-    :math:`F(\rho, \sigma) \approx 0.774`. This can be observed in :code:`toqito` as
+    Calculating the fidelity between the states :math:`\rho` and :math:`\sigma` as :math:`F(\rho, \sigma) \approx
+    0.774`. This can be observed in :code:`toqito` as
 
     >>> from toqito.states import basis
     >>> from toqito.state_metrics import fidelity
@@ -40,9 +40,8 @@ def sub_fidelity(rho: np.ndarray, sigma: np.ndarray) -> float:
     >>> fidelity(rho, sigma)
     0.77389339119464
 
-    As the sub-fidelity is a lower bound on the fidelity, that is
-    :math:`E(\rho, \sigma) \leq F(\rho, \sigma)`, we can use :code:`toqito` to observe that
-    :math:`E(\rho, \sigma) \approx 0.599\leq F(\rho, \sigma \approx 0.774`.
+    As the sub-fidelity is a lower bound on the fidelity, that is :math:`E(\rho, \sigma) \leq F(\rho, \sigma)`, we can
+    use :code:`toqito` to observe that :math:`E(\rho, \sigma) \approx 0.599\leq F(\rho, \sigma \approx 0.774`.
 
     >>> from toqito.states import basis
     >>> from toqito.state_metrics import sub_fidelity
diff --git a/toqito/state_metrics/trace_distance.py b/toqito/state_metrics/trace_distance.py
index 1291a0630..e9a0ccd29 100644
--- a/toqito/state_metrics/trace_distance.py
+++ b/toqito/state_metrics/trace_distance.py
@@ -34,8 +34,8 @@ def trace_distance(rho: np.ndarray, sigma: np.ndarray) -> float:
                          1 & 0 & 0 & 1
                        \end{pmatrix} \in \text{D}(\mathcal{X}).
 
-    The trace distance between :math:`\rho` and another state :math:`\sigma` is equal to :math:`0`
-    if any only if :math:`\rho = \sigma`. We can check this using the :code:`toqito` package.
+    The trace distance between :math:`\rho` and another state :math:`\sigma` is equal to :math:`0` if any only if
+    :math:`\rho = \sigma`. We can check this using the :code:`toqito` package.
 
     >>> from toqito.states import bell
     >>> from toqito.state_metrics import trace_norm
diff --git a/toqito/states/basis.py b/toqito/states/basis.py
index 43ec6b27d..63df9f5a0 100644
--- a/toqito/states/basis.py
+++ b/toqito/states/basis.py
@@ -50,5 +50,4 @@ def basis(dim: int, pos: int) -> np.ndarray:
 
     ret = np.array(list(int(x) for x in list(f"{0:0{dim}}")))
     ret[pos] = 1
-    ret = ret.conj().T.reshape(-1, 1)
-    return ret
+    return ret.conj().T.reshape(-1, 1)
diff --git a/toqito/states/bell.py b/toqito/states/bell.py
index 5fa5a1b9c..2952fc6fa 100644
--- a/toqito/states/bell.py
+++ b/toqito/states/bell.py
@@ -50,12 +50,13 @@ def bell(idx: int) -> np.ndarray:
     :return: Bell state with index :code:`idx`.
     """
     e_0, e_1 = basis(2, 0), basis(2, 1)
-    if idx == 0:
-        return 1 / np.sqrt(2) * (np.kron(e_0, e_0) + np.kron(e_1, e_1))
-    if idx == 1:
-        return 1 / np.sqrt(2) * (np.kron(e_0, e_0) - np.kron(e_1, e_1))
-    if idx == 2:
-        return 1 / np.sqrt(2) * (np.kron(e_0, e_1) + np.kron(e_1, e_0))
-    if idx == 3:
-        return 1 / np.sqrt(2) * (np.kron(e_0, e_1) - np.kron(e_1, e_0))
+    match idx:
+        case 0:
+            return 1 / np.sqrt(2) * (np.kron(e_0, e_0) + np.kron(e_1, e_1))
+        case 1:
+            return 1 / np.sqrt(2) * (np.kron(e_0, e_0) - np.kron(e_1, e_1))
+        case 2:
+            return 1 / np.sqrt(2) * (np.kron(e_0, e_1) + np.kron(e_1, e_0))
+        case 3:
+            return 1 / np.sqrt(2) * (np.kron(e_0, e_1) - np.kron(e_1, e_0))
     raise ValueError("Invalid integer value for Bell state.")
diff --git a/toqito/states/brauer.py b/toqito/states/brauer.py
index 49299958c..06754bebc 100644
--- a/toqito/states/brauer.py
+++ b/toqito/states/brauer.py
@@ -10,11 +10,10 @@ def brauer(dim: int, p_val: int) -> np.ndarray:
     r"""
     Produce all Brauer states [WikBrauer]_.
 
-    Produce a matrix whose columns are all of the (unnormalized) "Brauer" states: states that are
-    the :code:`p_val`-fold tensor product of the standard maximally-entangled pure state on
-    :code:`dim` local dimensions. There are many such states, since there are many different ways to
-    group the :code:`2 * p_val` parties into :code:`p_val` pairs (with each pair corresponding to
-    one maximally-entangled state).
+    Produce a matrix whose columns are all of the (unnormalized) "Brauer" states: states that are the :code:`p_val`-fold
+    tensor product of the standard maximally-entangled pure state on :code:`dim` local dimensions. There are many such
+    states, since there are many different ways to group the :code:`2 * p_val` parties into :code:`p_val` pairs (with
+    each pair corresponding to one maximally-entangled state).
 
     The exact number of such states is:
 
diff --git a/toqito/states/breuer.py b/toqito/states/breuer.py
index 7871846da..1d1ae642f 100644
--- a/toqito/states/breuer.py
+++ b/toqito/states/breuer.py
@@ -9,18 +9,16 @@ def breuer(dim: int, lam: float) -> np.ndarray:
     r"""
     Produce a Breuer state [HPBreuer]_.
 
-    Gives a Breuer bound entangled state for two qudits of local dimension :code:`dim`, with the
-    :code:`lam` parameter describing the weight of the singlet component as described in
-    [HPBreuer]_.
+    Gives a Breuer bound entangled state for two qudits of local dimension :code:`dim`, with the :code:`lam` parameter
+    describing the weight of the singlet component as described in [HPBreuer]_.
 
     This function was adapted from the QETLAB package.
 
     Examples
     ==========
 
-    We can generate a Breuer state of dimension :math:`4` with weight :math:`0.1`. For any weight
-    above :math:`0`, the state will be bound entangled, that is, it will satisfy the PPT criterion,
-    but it will be entangled.
+    We can generate a Breuer state of dimension :math:`4` with weight :math:`0.1`. For any weight above :math:`0`, the
+    state will be bound entangled, that is, it will satisfy the PPT criterion, but it will be entangled.
 
     >>> from toqito.states import breuer
     >>> breuer(2, 0.1)
@@ -46,6 +44,4 @@ def breuer(dim: int, lam: float) -> np.ndarray:
     max_entangled(dim)
     psi = np.dot(np.kron(np.identity(dim), v_mat), max_entangled(dim))
 
-    return lam * (psi * psi.conj().T) + (1 - lam) * 2 * symmetric_projection(dim) / (
-        dim * (dim + 1)
-    )
+    return lam * (psi * psi.conj().T) + (1 - lam) * 2 * symmetric_projection(dim) / (dim * (dim + 1))
diff --git a/toqito/states/chessboard.py b/toqito/states/chessboard.py
index 3f2d84b9b..f47c49f5c 100644
--- a/toqito/states/chessboard.py
+++ b/toqito/states/chessboard.py
@@ -6,10 +6,9 @@ def chessboard(mat_params: list[float], s_param: float = None, t_param: float =
     r"""
     Produce a chessboard state [BP00]_.
 
-    Generates the chessboard state defined in [BP00]_. Note that, for certain choices of
-    :code:`s_param` and :code:`t_param`, this state will not have positive partial transpose, and
-    thus may not be bound entangled.
-
+    Generates the chessboard state defined in [BP00]_. Note that, for certain choices of :code:`s_param` and
+    :code:`t_param`, this state will not have positive partial transpose, and thus may not be bound entangled.
+    
     Examples
     ==========
 
@@ -36,10 +35,6 @@ def chessboard(mat_params: list[float], s_param: float = None, t_param: float =
      [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
        0.        ,  0.        ,  0.        ,  0.        ]]
 
-    See Also
-    ========
-    chessboard
-
     References
     ==========
     .. [BP00] Three qubits can be entangled in two inequivalent ways.
@@ -58,26 +53,8 @@ def chessboard(mat_params: list[float], s_param: float = None, t_param: float =
         t_param = mat_params[0] * mat_params[3] / mat_params[4]
 
     v_1 = np.array([[mat_params[4], 0, s_param, 0, mat_params[5], 0, 0, 0, 0]])
-
     v_2 = np.array([[0, mat_params[0], 0, mat_params[1], 0, mat_params[2], 0, 0, 0]])
-
     v_3 = np.array([[np.conj(mat_params[5]), 0, 0, 0, -np.conj(mat_params[4]), 0, t_param, 0, 0]])
-
-    v_4 = np.array(
-        [
-            [
-                0,
-                np.conj(mat_params[1]),
-                0,
-                -np.conj(mat_params[0]),
-                0,
-                0,
-                0,
-                mat_params[3],
-                0,
-            ]
-        ]
-    )
-
+    v_4 = np.array([[0, np.conj(mat_params[1]), 0, -np.conj(mat_params[0]), 0, 0, 0, mat_params[3], 0]])
     rho = v_1.conj().T * v_1 + v_2.conj().T * v_2 + v_3.conj().T * v_3 + v_4.conj().T * v_4
     return rho / np.trace(rho)
diff --git a/toqito/states/domino.py b/toqito/states/domino.py
index 7aaa63049..be47d0ab3 100644
--- a/toqito/states/domino.py
+++ b/toqito/states/domino.py
@@ -88,22 +88,23 @@ def domino(idx: int) -> np.ndarray:
     :return: Domino state of index :code:`idx`.
     """
     e_0, e_1, e_2 = basis(3, 0), basis(3, 1), basis(3, 2)
-    if idx == 0:
-        return np.kron(e_1, e_1)
-    if idx == 1:
-        return np.kron(e_0, 1 / np.sqrt(2) * (e_0 + e_1))
-    if idx == 2:
-        return np.kron(e_0, 1 / np.sqrt(2) * (e_0 - e_1))
-    if idx == 3:
-        return np.kron(e_2, 1 / np.sqrt(2) * (e_1 + e_2))
-    if idx == 4:
-        return np.kron(e_2, 1 / np.sqrt(2) * (e_1 - e_2))
-    if idx == 5:
-        return np.kron(1 / np.sqrt(2) * (e_1 + e_2), e_0)
-    if idx == 6:
-        return np.kron(1 / np.sqrt(2) * (e_1 - e_2), e_0)
-    if idx == 7:
-        return np.kron(1 / np.sqrt(2) * (e_0 + e_1), e_2)
-    if idx == 8:
-        return np.kron(1 / np.sqrt(2) * (e_0 - e_1), e_2)
+    match idx:
+        case 0:
+            return np.kron(e_1, e_1)
+        case 1:
+            return np.kron(e_0, 1 / np.sqrt(2) * (e_0 + e_1))
+        case 2:
+            return np.kron(e_0, 1 / np.sqrt(2) * (e_0 - e_1))
+        case 3:
+            return np.kron(e_2, 1 / np.sqrt(2) * (e_1 + e_2))
+        case 4:
+            return np.kron(e_2, 1 / np.sqrt(2) * (e_1 - e_2))
+        case 5:
+            return np.kron(1 / np.sqrt(2) * (e_1 + e_2), e_0)
+        case 6:
+            return np.kron(1 / np.sqrt(2) * (e_1 - e_2), e_0)
+        case 7:
+            return np.kron(1 / np.sqrt(2) * (e_0 + e_1), e_2)
+        case 8:
+            return np.kron(1 / np.sqrt(2) * (e_0 - e_1), e_2)
     raise ValueError("Invalid integer value for Domino state.")
diff --git a/toqito/states/gen_bell.py b/toqito/states/gen_bell.py
index 138343c0f..f4bbb4c4b 100644
--- a/toqito/states/gen_bell.py
+++ b/toqito/states/gen_bell.py
@@ -98,4 +98,4 @@ def gen_bell(k_1: int, k_2: int, dim: int) -> np.ndarray:
     :param dim: The dimension of the generalized Bell state.
     """
     gen_pauli_w = gen_pauli(k_1, k_2, dim)
-    return 1 / dim * vec(gen_pauli_w) * vec(gen_pauli_w).conj().T
+    return 1 / dim * vec(gen_pauli_w) @ vec(gen_pauli_w).conj().T
diff --git a/toqito/states/ghz.py b/toqito/states/ghz.py
index 605b8f075..f9d98820f 100644
--- a/toqito/states/ghz.py
+++ b/toqito/states/ghz.py
@@ -8,9 +8,8 @@ def ghz(dim: int, num_qubits: int, coeff: list[int] = None) -> sparse:
     r"""
     Generate a (generalized) GHZ state [GHZ07]_.
 
-    Returns a :code:`num_qubits`-partite GHZ state acting on :code:`dim` local dimensions, described
-    in [GHZ07]_. For example, :code:`ghz(2, 3)` returns the standard 3-qubit GHZ state on qubits.
-    The output of this function is sparse.
+    Returns a :code:`num_qubits`-partite GHZ state acting on :code:`dim` local dimensions, described in [GHZ07]_. For
+    example, :code:`ghz(2, 3)` returns the standard 3-qubit GHZ state on qubits.  The output of this function is sparse.
 
     For a system of :code:`num_qubits` qubits (i.e., :code:`dim = 2`), the GHZ state can be written
     as
@@ -40,8 +39,8 @@ def ghz(dim: int, num_qubits: int, coeff: list[int] = None) -> sparse:
      [0.        ],
      [0.70710678]]
 
-    As this function covers the generalized GHZ state, we can consider higher dimensions. For
-    instance here is the GHZ state in :math:`\mathbb{C}^{4^{\otimes 7}}` as
+    As this function covers the generalized GHZ state, we can consider higher dimensions. For instance here is the GHZ
+    state in :math:`\mathbb{C}^{4^{\otimes 7}}` as
 
     .. math::
         \frac{1}{\sqrt{30}} \left(|0000000 \rangle + 2|1111111 \rangle +
diff --git a/toqito/states/horodecki.py b/toqito/states/horodecki.py
index 797ccccf1..dde03ec36 100644
--- a/toqito/states/horodecki.py
+++ b/toqito/states/horodecki.py
@@ -6,14 +6,12 @@ def horodecki(a_param: float, dim: list[int] = None) -> np.ndarray:
     r"""
     Produce a Horodecki state [HOR]_, [CHR]_.
 
-    Returns the Horodecki state in either :math:`(3 \otimes 3)`-dimensional space or
-    :math:`(2 \otimes 4)`-dimensional space, depending on the dimensions in the 1-by-2 vector
-    :code:`dim`.
+    Returns the Horodecki state in either :math:`(3 \otimes 3)`-dimensional space or :math:`(2 \otimes 4)`-dimensional
+    space, depending on the dimensions in the 1-by-2 vector :code:`dim`.
 
-    The Horodecki state was introduced in [1] which serves as an example in
-    :math:`\mathbb{C}^3 \otimes \mathbb{C}` or :math:`\mathbb{C}^2 \otimes \mathbb{C}^4` of an
-    entangled state that is positive under partial transpose (PPT). The state is PPT for all
-    :math:`a \in [0, 1]` and separable only for :code:`a_param = 0` or :code:`a_param = 1`.
+    The Horodecki state was introduced in [1] which serves as an example in :math:`\mathbb{C}^3 \otimes \mathbb{C}` or
+    :math:`\mathbb{C}^2 \otimes \mathbb{C}^4` of an entangled state that is positive under partial transpose (PPT). The
+    state is PPT for all :math:`a \in [0, 1]` and separable only for :code:`a_param = 0` or :code:`a_param = 1`.
 
     These states have the following definitions:
 
@@ -53,15 +51,14 @@ def horodecki(a_param: float, dim: list[int] = None) -> np.ndarray:
         \end{equation}
 
     .. note::
-        Refer to [CHR]_ (specifically equations (1) and (2)) for more information on this state
-        and its properties. The 3x3 Horodecki state is defined explicitly in Section 4.1 of [HOR]_
-        and the 2x4 Horodecki state is defined explicitly in Section 4.2 of [HOR]_.
+        Refer to [CHR]_ (specifically equations (1) and (2)) for more information on this state and its properties. The
+        3x3 Horodecki state is defined explicitly in Section 4.1 of [HOR]_ and the 2x4 Horodecki state is defined
+        explicitly in Section 4.2 of [HOR]_.
 
     Examples
     ==========
 
-    The following code generates a Horodecki state in
-    :math:`\mathbb{C}^3 \otimes \mathbb{C}^3`
+    The following code generates a Horodecki state in :math:`\mathbb{C}^3 \otimes \mathbb{C}^3`
 
     >>> from toqito.states import horodecki
     >>> horodecki(0.5, [3, 3])
diff --git a/toqito/states/isotropic.py b/toqito/states/isotropic.py
index daf252699..41987bfcd 100644
--- a/toqito/states/isotropic.py
+++ b/toqito/states/isotropic.py
@@ -8,8 +8,8 @@ def isotropic(dim: int, alpha: float) -> np.ndarray:
     r"""
     Produce a isotropic state [HH99]_.
 
-    Returns the isotropic state with parameter :code:`alpha` acting on
-    (:code:`dim`-by-:code:`dim`)-dimensional space. The isotropic state has the following form
+    Returns the isotropic state with parameter :code:`alpha` acting on (:code:`dim`-by-:code:`dim`)-dimensional space.
+    The isotropic state has the following form
 
     .. math::
         \begin{equation}
@@ -18,14 +18,14 @@ def isotropic(dim: int, alpha: float) -> np.ndarray:
             \mathbb{C}^d \otimes \mathbb{C}^2
         \end{equation}
 
-    where :math:`|\psi_+ \rangle = \frac{1}{\sqrt{d}} \sum_j |j \rangle \otimes |j \rangle` is the
-    maximally entangled state.
+    where :math:`|\psi_+ \rangle = \frac{1}{\sqrt{d}} \sum_j |j \rangle \otimes |j \rangle` is the maximally entangled
+    state.
 
     Examples
     ==========
 
-    To generate the isotropic state with parameter :math:`\alpha=1/2`, we can make the following
-    call to :code:`toqito` as
+    To generate the isotropic state with parameter :math:`\alpha=1/2`, we can make the following call to :code:`toqito`
+    as
 
     >>> from toqito.states import isotropic
     >>> isotropic(3, 1 / 2)
@@ -59,4 +59,4 @@ def isotropic(dim: int, alpha: float) -> np.ndarray:
     :return: Isotropic state of dimension :code:`dim`.
     """
     psi = max_entangled(dim, False, False)
-    return (1 - alpha) * np.identity(dim**2) / dim**2 + alpha * psi * psi.conj().T / dim
+    return (1 - alpha) * np.identity(dim**2) / dim**2 + alpha * psi @ psi.conj().T / dim
diff --git a/toqito/states/max_entangled.py b/toqito/states/max_entangled.py
index cb17d22bd..faef5aec8 100644
--- a/toqito/states/max_entangled.py
+++ b/toqito/states/max_entangled.py
@@ -3,16 +3,14 @@
 import scipy as sp
 
 
-def max_entangled(
-    dim: int, is_sparse: bool = False, is_normalized: bool = True
-) -> [np.ndarray, sp.sparse.dia_matrix]:
+def max_entangled(dim: int, is_sparse: bool = False, is_normalized: bool = True) -> [np.ndarray, sp.sparse.dia_matrix]:
     r"""
     Produce a maximally entangled bipartite pure state [WikEnt]_.
 
-    Produces a maximally entangled pure state as above that is sparse if :code:`is_sparse = True`
-    and is full if :code:`is_sparse = False`. The pure state is normalized to have Euclidean norm 1
-    if :code:`is_normalized = True`, and it is unnormalized (i.e. each entry in the vector is 0 or 1
-    and the Euclidean norm of the vector is :code:`sqrt(dim)` if :code:`is_normalized = False`.
+    Produces a maximally entangled pure state as above that is sparse if :code:`is_sparse = True` and is full if
+    :code:`is_sparse = False`. The pure state is normalized to have Euclidean norm 1 if :code:`is_normalized = True`,
+    and it is unnormalized (i.e. each entry in the vector is 0 or 1 and the Euclidean norm of the vector is
+    :code:`sqrt(dim)` if :code:`is_normalized = False`.
 
     Examples
     ==========
diff --git a/toqito/states/max_mixed.py b/toqito/states/max_mixed.py
index e4e77c082..456ad2c53 100644
--- a/toqito/states/max_mixed.py
+++ b/toqito/states/max_mixed.py
@@ -8,8 +8,7 @@ def max_mixed(dim: int, is_sparse: bool = False) -> [np.ndarray, sparse.dia_matr
     r"""
     Produce the maximally mixed state [AAR6]_.
 
-    Produces the maximally mixed state on of :code:`dim` dimensions. The maximally mixed state is
-    defined as
+    Produces the maximally mixed state on of :code:`dim` dimensions. The maximally mixed state is defined as
 
     .. math::
         \omega = \frac{1}{d} \begin{pmatrix}
@@ -24,11 +23,11 @@ def max_mixed(dim: int, is_sparse: bool = False) -> [np.ndarray, sparse.dia_matr
     .. math::
         \omega = \frac{\mathbb{I}}{\text{dim}(\mathcal{X})}
 
-    for some complex Euclidean space :math:`\mathcal{X}`. The maximally mixed state is sometimes
-    also referred to as the tracial state.
+    for some complex Euclidean space :math:`\mathcal{X}`. The maximally mixed state is sometimes also referred to as the
+    tracial state.
 
-    The maximally mixed state is returned as a sparse matrix if :code:`is_sparse = True` and is full
-    if :code:`is_sparse = False`.
+    The maximally mixed state is returned as a sparse matrix if :code:`is_sparse = True` and is full if :code:`is_sparse
+    = False`.
 
     Examples
     ==========
diff --git a/toqito/states/singlet.py b/toqito/states/singlet.py
index 53ee74a72..ebed14132 100644
--- a/toqito/states/singlet.py
+++ b/toqito/states/singlet.py
@@ -36,8 +36,8 @@ def singlet(dim: int) -> np.ndarray:
          [ 0. , -0.5,  0.5,  0. ],
          [ 0. ,  0. ,  0. ,  0. ]]
 
-    It is possible for us to consider higher dimensional singlet states. For instance, we
-    can consider the :math:`3`-dimensional Singlet state as follows:
+    It is possible for us to consider higher dimensional singlet states. For instance, we can consider the
+    :math:`3`-dimensional Singlet state as follows:
 
     >>> from toqito.states import singlet
     >>> dim = 3
diff --git a/toqito/states/tests/test_werner.py b/toqito/states/tests/test_werner.py
index 67929fd63..522c59cd9 100644
--- a/toqito/states/tests/test_werner.py
+++ b/toqito/states/tests/test_werner.py
@@ -1,7 +1,9 @@
 """Test werner."""
 import numpy as np
+import pytest
 
 from toqito.states import werner
+from toqito.matrix_props import is_density
 
 
 def test_werner_qutrit():
@@ -17,7 +19,25 @@ def test_werner_multipartite():
     np.testing.assert_equal(np.isclose(res[0][0], 0.1127, atol=1e-02), True)
 
 
-def test_werner_invalid_alpha():
-    """Test for invalid `alpha` parameter."""
-    with np.testing.assert_raises(ValueError):
-        werner(3, [1, 2])
+def test_werner_multipartite_valid():
+    """Test multipartite Werner states with valid alpha lengths."""
+    # Valid alpha length for p=3 (2!-1 = 1)
+    alpha = [0.5]
+    dim = 2
+    state = werner(dim, alpha)
+    np.testing.assert_equal(is_density(state), True)
+
+
+@pytest.mark.parametrize("dim, alpha", [
+    # Invalid alpha length (not matching p!-1 for any integer p > 1)
+    (2, [0.5, 0.6, 0.7]),
+    # Test with an integer (which is not a valid type for alpha)
+    (2, 5),
+    # Test with a string (which is not a valid type for alpha)
+    (2, "invalid"),
+    # Test with a dictionary (which is not a valid type for alpha)
+    (2, {"key": "value"}),
+])
+def test_werner_state_invalid(dim, alpha):
+    with pytest.raises(ValueError):
+        werner(dim, alpha)
diff --git a/toqito/states/tile.py b/toqito/states/tile.py
index d0af4b092..da59eb47b 100644
--- a/toqito/states/tile.py
+++ b/toqito/states/tile.py
@@ -8,8 +8,7 @@ def tile(idx: int) -> np.ndarray:
     r"""
     Produce a Tile state [UPBTile99]_.
 
-    The Tile states constitute five states on 3-by-3 dimensional space that form a UPB
-    (unextendible product basis).
+    The Tile states constitute five states on 3-by-3 dimensional space that form a UPB (unextendible product basis).
 
     Returns one of the following five tile states depending on the value of :code:`idx`:
 
@@ -62,14 +61,15 @@ def tile(idx: int) -> np.ndarray:
     :return: Tile state.
     """
     e_0, e_1, e_2 = basis(3, 0), basis(3, 1), basis(3, 2)
-    if idx == 0:
-        return 1 / np.sqrt(2) * np.kron(e_0, (e_0 - e_1))
-    if idx == 1:
-        return 1 / np.sqrt(2) * np.kron((e_0 - e_1), e_2)
-    if idx == 2:
-        return 1 / np.sqrt(2) * np.kron(e_2, (e_1 - e_2))
-    if idx == 3:
-        return 1 / np.sqrt(2) * np.kron((e_1 - e_2), e_0)
-    if idx == 4:
-        return 1 / 3 * np.kron((e_0 + e_1 + e_2), (e_0 + e_1 + e_2))
+    match idx:
+        case 0:
+            return 1 / np.sqrt(2) * np.kron(e_0, (e_0 - e_1))
+        case 1:
+            return 1 / np.sqrt(2) * np.kron((e_0 - e_1), e_2)
+        case 2:
+            return 1 / np.sqrt(2) * np.kron(e_2, (e_1 - e_2))
+        case 3:
+            return 1 / np.sqrt(2) * np.kron((e_1 - e_2), e_0)
+        case 4:
+            return 1 / 3 * np.kron((e_0 + e_1 + e_2), (e_0 + e_1 + e_2))
     raise ValueError("Invalid integer value for Tile state.")
diff --git a/toqito/states/w_state.py b/toqito/states/w_state.py
index f518b7dff..4b5ffd653 100644
--- a/toqito/states/w_state.py
+++ b/toqito/states/w_state.py
@@ -37,8 +37,7 @@ def w_state(num_qubits: int, coeff: list[int] = None) -> np.ndarray:
      [0.    ],
      [0.    ]]
 
-    We may also generate a generalized :math:`W`-state. For instance, here is a
-    :math:`4`-dimensional :math:`W`-state
+    We may also generate a generalized :math:`W`-state. For instance, here is a :math:`4`-dimensional :math:`W`-state
 
     .. math::
         \frac{1}{\sqrt{30}} \left( |1000 \rangle + 2|0100 \rangle + 3|0010
diff --git a/toqito/states/werner.py b/toqito/states/werner.py
index 1b194d1ec..c46921bbe 100644
--- a/toqito/states/werner.py
+++ b/toqito/states/werner.py
@@ -1,8 +1,5 @@
 """Werner state."""
-
-
 import itertools
-
 import numpy as np
 
 from toqito.perms import permutation_operator, swap_operator
@@ -21,20 +18,16 @@ def werner(dim: int, alpha: float | list[float]) -> np.ndarray:
             \mathbb{I} - \alpha S \right) \in \mathbb{C}^d \otimes \mathbb{C}^d.
         \end{equation}
 
-    Yields a Werner state with parameter :code:`alpha` acting on :code:`(dim * dim)`- dimensional
-    space. More specifically, :math:`\rho` is the density operator defined by
-    :math:`(\mathbb{I} - `alpha` S)` (normalized to have trace 1), where :math:`\mathbb{I}` is the
-    density operator and :math:`S` is the operator that swaps two copies of :code:`dim`-dimensional
-    space (see swap and swap_operator for example).
+    Yields a Werner state with parameter :code:`alpha` acting on :code:`(dim * dim)`- dimensional space. More
+    specifically, :math:`\rho` is the density operator defined by :math:`(\mathbb{I} - `alpha` S)` (normalized to have
+    trace 1), where :math:`\mathbb{I}` is the density operator and :math:`S` is the operator that swaps two copies of
+    :code:`dim`-dimensional space (see swap and swap_operator for example).
 
-    If :code:`alpha` is a vector with :math:`p!-1` entries, for some integer :math:`p > 1`, then a
-    multipartite Werner state is returned. This multipartite Werner state is the normalization of
-    I - `alpha(1)*P(2)` - ... - `alpha(p!-1)*P(p!)`, where P(i) is the operator that permutes p
-    subsystems according to the i-th permutation when they are written in lexicographical order
-    (for example, the lexicographical ordering when p = 3 is:
-    `[1, 2, 3], [1, 3, 2], [2, 1,3], [2, 3, 1], [3, 1, 2], [3, 2, 1],`
-
-    so P(4) in this case equals permutation_operator(dim, [2, 3, 1]).
+    If :code:`alpha` is a vector with :math:`p!-1` entries, for some integer :math:`p > 1`, then a multipartite Werner
+    state is returned. This multipartite Werner state is the normalization of I - `alpha(1)*P(2)` - ... -
+    `alpha(p!-1)*P(p!)`, where P(i) is the operator that permutes p subsystems according to the i-th permutation when
+    they are written in lexicographical order (for example, the lexicographical ordering when p = 3 is: `[1, 2, 3], [1,
+    3, 2], [2, 1,3], [2, 3, 1], [3, 1, 2], [3, 2, 1],` so P(4) in this case equals permutation_operator(dim, [2, 3, 1]).
 
     Examples
     ==========
@@ -94,14 +87,9 @@ def werner(dim: int, alpha: float | list[float]) -> np.ndarray:
     :param alpha: Parameter to specify Werner state.
     :return: A Werner state of dimension :code:`dim`.
     """
-    # The total number of permutation operators.
-    if isinstance(alpha, float):
-        n_fac = 2
-    else:
-        n_fac = len(alpha) + 1
-
     # Multipartite Werner state.
-    if n_fac > 2:
+    if isinstance(alpha, list):
+        n_fac = len(alpha) + 1
         # Compute the number of parties from `len(alpha)`.
         n_var = n_fac
         # We won't actually go all the way to `n_fac`.
@@ -111,20 +99,25 @@ def werner(dim: int, alpha: float | list[float]) -> np.ndarray:
                 break
             if n_var < i:
                 raise ValueError(
-                    "InvalidAlpha: The `alpha` vector must contain"
-                    " p!-1 entries for some integer p > 1."
+                    "InvalidAlpha: The `alpha` vector must contain p!-1 entries for some integer p > 1."
                 )
 
-        # Done error checking and computing the number of parties -- now
-        # compute the Werner state.
-        perms = list(itertools.permutations(np.arange(n_var)))
+        # Done error checking and computing the number of parties
+        # -- now compute the Werner state.
+        perms = list(itertools.permutations(range(n_var)))
         sorted_perms = np.argsort(perms, axis=1) + 1
 
+        rho = np.identity(dim**n_var)
         for i in range(2, n_fac):
-            rho = np.identity(dim**n_var) - alpha[i - 1] * permutation_operator(
-                dim, sorted_perms[i, :], False, True
+            rho -= alpha[i - 1] * permutation_operator(
+                dim, sorted_perms[i - 1, :], False, True
             )
         rho = rho / np.trace(rho)
         return rho
-    # Bipartite Werner state.
-    return (np.identity(dim**2) - alpha * swap_operator(dim, True)) / (dim * (dim - alpha))
+    
+    # Bipartite Werner state (executed only if alpha is a float).
+    if isinstance(alpha, float):
+        n_fac = 2
+        return (np.identity(dim**2) - alpha * swap_operator(dim, True)) / (dim * (dim - alpha))
+
+    raise ValueError("Alpha must be either a float or a list of floats.")