From cd87cb29d15e4ce82aa2f0fb26ceee84db71caa9 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Sun, 9 Oct 2022 22:49:33 -0700
Subject: [PATCH 01/12] Add Qudit to Qubit transforms

---
 docs/quantum_chess/concepts.ipynb         |  2 +-
 requirements.txt                          |  1 +
 unitary/alpha/qudit_to_qubit_transform.py | 46 +++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 unitary/alpha/qudit_to_qubit_transform.py

diff --git a/docs/quantum_chess/concepts.ipynb b/docs/quantum_chess/concepts.ipynb
index 9e52cdb3..fda17746 100644
--- a/docs/quantum_chess/concepts.ipynb
+++ b/docs/quantum_chess/concepts.ipynb
@@ -1625,7 +1625,7 @@
     "    print(f\"Qubit {q} maps to {dynamic_mapping[q]}\")\n",
     "dynamically_mapped_circuit = cirq.optimize_for_target_gateset(\n",
     "    exclusion_with_sqrt_iswaps, gateset=cirq.SqrtIswapTargetGateset()\n",
-    ").transform_qubits(qubit_map=lambda q: dynamic_mapping[q])\n",    
+    ").transform_qubits(qubit_map=lambda q: dynamic_mapping[q])\n",
     "print(dynamically_mapped_circuit)"
    ]
   },
diff --git a/requirements.txt b/requirements.txt
index a76f17a4..1210459c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,6 +2,7 @@ cirq-core>=0.15.0
 cirq-google>=0.15.0
 # When changing Cirq requirements be sure to update dev_tools/write-ci-requirements.py
 
+numpy
 seaborn
 sphinx
 ipython
diff --git a/unitary/alpha/qudit_to_qubit_transform.py b/unitary/alpha/qudit_to_qubit_transform.py
new file mode 100644
index 00000000..7b950290
--- /dev/null
+++ b/unitary/alpha/qudit_to_qubit_transform.py
@@ -0,0 +1,46 @@
+# Copyright 2022 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+
+
+def qudit_to_qubit_state(dim_single_qudit: int, num_qudits: int, qudit_state_vector: np.ndarray, pad_value: np.complex_ = 0):
+    """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
+    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit)/np.log(2))
+    # Reshape the state vector to a `num_qudits` rank tensor.
+    state_tensor = qudit_state_vector.reshape((dim_single_qudit,)*num_qudits)
+    # Number of extra elements needed in each dimension if represented using qubits.
+    padding_amount = int((2**num_qubits_per_qudit)-dim_single_qudit)
+    # Expand the number of elements in each dimension by the padding_amount. Fill
+    # the new elements with the pad_value.
+    padded_state_tensor = np.pad(state_tensor,
+                                 (0, padding_amount),
+                                 constant_values=pad_value)
+    # Return a flattened state vector view of the final tensor.
+    return np.ravel(padded_state_tensor)
+
+
+def qudit_to_qubit_unitary(dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray):
+    """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
+    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit)/np.log(2))
+    num_qubits = num_qubits_per_qudit * num_qudits
+    dim_qubit_space = int(2**num_qubits)
+
+    # Treat the unitary as a num_qubits^2 system's state vector and represent it using qubits (pad with 0s).
+    padded_unitary = qudit_to_qubit_state(dim_single_qudit, num_qudits**2, np.ravel(qudit_unitary))
+    # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
+    pad_qubits_vector = qudit_to_qubit_state(dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1)
+    # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the pad_qubits_vector.
+    # This addition ensures that the invalid states with the "padding" bits map to identity, preserving unitarity.
+    return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(pad_qubits_vector)

From 2467f086f531671400e112998c0d2501d357b020 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Sun, 9 Oct 2022 22:54:42 -0700
Subject: [PATCH 02/12] Formating

---
 unitary/alpha/qudit_to_qubit_transform.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/unitary/alpha/qudit_to_qubit_transform.py b/unitary/alpha/qudit_to_qubit_transform.py
index 7b950290..17d757e4 100644
--- a/unitary/alpha/qudit_to_qubit_transform.py
+++ b/unitary/alpha/qudit_to_qubit_transform.py
@@ -15,7 +15,8 @@
 import numpy as np
 
 
-def qudit_to_qubit_state(dim_single_qudit: int, num_qudits: int, qudit_state_vector: np.ndarray, pad_value: np.complex_ = 0):
+def qudit_to_qubit_state(
+        dim_single_qudit: int, num_qudits: int, qudit_state_vector: np.ndarray, pad_value: np.complex_ = 0):
     """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
     num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit)/np.log(2))
     # Reshape the state vector to a `num_qudits` rank tensor.
@@ -40,7 +41,9 @@ def qudit_to_qubit_unitary(dim_single_qudit: int, num_qudits: int, qudit_unitary
     # Treat the unitary as a num_qubits^2 system's state vector and represent it using qubits (pad with 0s).
     padded_unitary = qudit_to_qubit_state(dim_single_qudit, num_qudits**2, np.ravel(qudit_unitary))
     # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
-    pad_qubits_vector = qudit_to_qubit_state(dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1)
-    # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the pad_qubits_vector.
-    # This addition ensures that the invalid states with the "padding" bits map to identity, preserving unitarity.
+    pad_qubits_vector = qudit_to_qubit_state(
+        dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1)
+    # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the
+    # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map to identity,
+    # preserving unitarity.
     return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(pad_qubits_vector)

From 66a3d62d62a4d4284cfe71f57c8d4feaf4afa858 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Thu, 13 Oct 2022 23:47:28 -0700
Subject: [PATCH 03/12] Add some memoization

---
 unitary/alpha/qudit_to_qubit_transform.py | 73 +++++++++++++++++++----
 1 file changed, 60 insertions(+), 13 deletions(-)

diff --git a/unitary/alpha/qudit_to_qubit_transform.py b/unitary/alpha/qudit_to_qubit_transform.py
index 17d757e4..d418eb31 100644
--- a/unitary/alpha/qudit_to_qubit_transform.py
+++ b/unitary/alpha/qudit_to_qubit_transform.py
@@ -15,35 +15,82 @@
 import numpy as np
 
 
-def qudit_to_qubit_state(
-        dim_single_qudit: int, num_qudits: int, qudit_state_vector: np.ndarray, pad_value: np.complex_ = 0):
+def _nearest_power_of_two_ceiling(x: int) -> int:
+    exp = 0
+    while x:
+        x = x >> 1
+        exp += 1
+    return 2**exp
+
+
+def qudit_to_qubit_state_impl(
+    dim_single_qudit: int,
+    num_qudits: int,
+    qudit_state_vector: np.ndarray,
+    pad_value: np.complex_ = 0,
+):
     """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
-    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit)/np.log(2))
     # Reshape the state vector to a `num_qudits` rank tensor.
-    state_tensor = qudit_state_vector.reshape((dim_single_qudit,)*num_qudits)
+    state_tensor = qudit_state_vector.reshape((dim_single_qudit,) * num_qudits)
     # Number of extra elements needed in each dimension if represented using qubits.
-    padding_amount = int((2**num_qubits_per_qudit)-dim_single_qudit)
+    padding_amount = _nearest_power_of_two_ceiling(dim_single_qudit) - dim_single_qudit
     # Expand the number of elements in each dimension by the padding_amount. Fill
     # the new elements with the pad_value.
-    padded_state_tensor = np.pad(state_tensor,
-                                 (0, padding_amount),
-                                 constant_values=pad_value)
+    padded_state_tensor = np.pad(
+        state_tensor, (0, padding_amount), constant_values=pad_value
+    )
     # Return a flattened state vector view of the final tensor.
     return np.ravel(padded_state_tensor)
 
 
-def qudit_to_qubit_unitary(dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray):
+def qudit_to_qubit_from_map(
+    qudit_state_vector: np.ndarray, qubit_to_qudit_index_map: np.ndarray
+):
+    return [
+        qudit_state_vector[qubit_to_qudit_index_map[j] - 1]
+        if qubit_to_qudit_index_map[j]
+        else 0
+        for j in range(len(qubit_to_qudit_index_map))
+    ]
+
+
+def qudit_to_qubit_state(
+    dim_single_qudit: int,
+    num_qudits: int,
+    qudit_state_vectors: list[np.ndarray],
+    pad_value: np.complex_ = 0,
+) -> list[np.ndarray]:
+    qubit_to_qudit_index_map = qudit_to_qubit_state_impl(
+        dim_single_qudit,
+        num_qudits,
+        np.array([i + 1 for i in range(dim_single_qudit**num_qudits)]),
+        pad_value,
+    )
+    return [
+        qudit_to_qubit_from_map(state, qubit_to_qudit_index_map)
+        for state in qudit_state_vectors
+    ]
+
+
+def qudit_to_qubit_unitary(
+    dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray
+):
     """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
-    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit)/np.log(2))
+    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit) / np.log(2))
     num_qubits = num_qubits_per_qudit * num_qudits
     dim_qubit_space = int(2**num_qubits)
 
     # Treat the unitary as a num_qubits^2 system's state vector and represent it using qubits (pad with 0s).
-    padded_unitary = qudit_to_qubit_state(dim_single_qudit, num_qudits**2, np.ravel(qudit_unitary))
+    padded_unitary = qudit_to_qubit_state(
+        dim_single_qudit, num_qudits**2, np.ravel(qudit_unitary)
+    )
     # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
     pad_qubits_vector = qudit_to_qubit_state(
-        dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1)
+        dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1
+    )
     # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the
     # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map to identity,
     # preserving unitarity.
-    return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(pad_qubits_vector)
+    return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(
+        pad_qubits_vector
+    )

From 03c192f507b5b5f1ecd1780ae796b4e9baec23a3 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Sat, 15 Oct 2022 16:52:46 -0700
Subject: [PATCH 04/12] Unitary implementation

---
 unitary/alpha/qudit_to_qubit_transform.py | 24 +++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/unitary/alpha/qudit_to_qubit_transform.py b/unitary/alpha/qudit_to_qubit_transform.py
index d418eb31..f25ecfd6 100644
--- a/unitary/alpha/qudit_to_qubit_transform.py
+++ b/unitary/alpha/qudit_to_qubit_transform.py
@@ -58,13 +58,11 @@ def qudit_to_qubit_state(
     dim_single_qudit: int,
     num_qudits: int,
     qudit_state_vectors: list[np.ndarray],
-    pad_value: np.complex_ = 0,
 ) -> list[np.ndarray]:
     qubit_to_qudit_index_map = qudit_to_qubit_state_impl(
         dim_single_qudit,
         num_qudits,
         np.array([i + 1 for i in range(dim_single_qudit**num_qudits)]),
-        pad_value,
     )
     return [
         qudit_to_qubit_from_map(state, qubit_to_qudit_index_map)
@@ -72,6 +70,28 @@ def qudit_to_qubit_state(
     ]
 
 
+def qudit_to_qubit_unitary(
+    dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray
+) -> np.ndarray:
+    qubit_to_qudit_index_map = qudit_to_qubit_state_impl(
+        dim_single_qudit,
+        num_qudits,
+        np.array([i + 1 for i in range(dim_single_qudit**num_qudits)]),
+    )
+    result = np.identity(dim_single_qudit**num_qudits)
+    for i, j in (
+        range(len(qubit_to_qudit_index_map)),
+        range(len(qubit_to_qudit_index_map)),
+    ):
+        result[i][j] = (
+            qudit_unitary[qubit_to_qudit_index_map[i] - 1][
+                qubit_to_qudit_index_map[j] - 1
+            ]
+            if qubit_to_qudit_index_map[i] and qubit_to_qudit_index_map[j]
+            else 0
+        )
+
+
 def qudit_to_qubit_unitary(
     dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray
 ):

From d600231b971fc174ffae2358ba29a485531d7006 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Mon, 24 Oct 2022 00:14:36 -0700
Subject: [PATCH 05/12] Inverse functions, tests and cleanup

---
 unitary/alpha/qudit_state_transform.py      | 129 ++++++++++++++++++++
 unitary/alpha/qudit_state_transform_test.py |  53 ++++++++
 unitary/alpha/qudit_to_qubit_transform.py   | 116 ------------------
 3 files changed, 182 insertions(+), 116 deletions(-)
 create mode 100644 unitary/alpha/qudit_state_transform.py
 create mode 100644 unitary/alpha/qudit_state_transform_test.py
 delete mode 100644 unitary/alpha/qudit_to_qubit_transform.py

diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
new file mode 100644
index 00000000..89ac50cc
--- /dev/null
+++ b/unitary/alpha/qudit_state_transform.py
@@ -0,0 +1,129 @@
+# Copyright 2022 Google
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import itertools
+
+import numpy as np
+from typing import List
+
+
+def _nearest_power_of_two_ceiling(qudit_dim: int) -> int:
+    if qudit_dim == 0:
+        return 0
+    # Max index of the single qudit state.
+    x = qudit_dim - 1
+    # Number of (qu)bits needed to represent the max index.
+    bits = 0
+    while x:
+        x = x >> 1
+        bits += 1
+    # Total number of states in the qubit representation i.e. Dimension of the single qudit
+    # post-conversion.
+    return 2**bits
+
+
+def qudit_to_qubit_state(
+    qudit_dimension: int,
+    num_qudits: int,
+    qudit_state_vector: np.ndarray,
+    pad_value: np.complex_ = 0,
+):
+    """Converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
+    # Reshape the state vector to a `num_qudits` rank tensor.
+    state_tensor = qudit_state_vector.reshape((qudit_dimension,) * num_qudits)
+    # Number of extra elements needed in each dimension if represented using qubits.
+    padding_amount = _nearest_power_of_two_ceiling(qudit_dimension) - qudit_dimension
+    # Expand the number of elements in each dimension by the padding_amount. Fill
+    # the new elements with the pad_value.
+    padded_state_tensor = np.pad(
+        state_tensor, pad_width=(0, padding_amount), constant_values=pad_value
+    )
+    # Return a flattened state vector view of the final tensor.
+    return np.ravel(padded_state_tensor)
+
+
+def qubit_to_qudit_state(
+    qudit_dimension: int,
+    num_qudits: int,
+    qubit_state_vector: np.ndarray,
+):
+    """Converts a m-qubit-per-qudit column vector to a qudit-space quantum state vector."""
+    mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
+    # Reshape the state vector to a `num_qudits` rank tensor.
+    state_tensor = qubit_state_vector.reshape((mbit_dimension,) * num_qudits)
+    # Shrink the number of elements in each dimension up to the qudit_dimension, ignoring the rest.
+    trimmed_state_tensor = state_tensor[(slice(qudit_dimension),) * num_qudits]
+    # Return a flattened state vector view of the final tensor.
+    return np.ravel(trimmed_state_tensor)
+
+
+def qudit_to_qubit_unitary(
+    qudit_dimension: int,
+    num_qudits: int,
+    qudit_unitary: np.ndarray,
+    memoize: bool = False,
+) -> np.ndarray:
+    """Converts a qudit-space quantum unitary to m-qubit-per-qudit unitary."""
+    dim_qubit_space = _nearest_power_of_two_ceiling(qudit_dimension) ** num_qudits
+
+    if memoize:
+        d_to_b_index_map = qubit_to_qudit_state(
+            qudit_dimension,
+            num_qudits,
+            np.array([i + 1 for i in range(dim_qubit_space)]),
+        )
+        result = np.identity(dim_qubit_space, dtype=qudit_unitary.dtype)
+        iter_range = range(qudit_dimension**num_qudits)
+        for i, j in itertools.product(iter_range, iter_range):
+            result[d_to_b_index_map[i] - 1][d_to_b_index_map[j] - 1] = qudit_unitary[i][
+                j
+            ]
+        return result
+
+    # Treat the unitary as a num_qudits^2 system's state vector and represent it using qubits (pad
+    # with 0s).
+    padded_unitary = qudit_to_qubit_state(
+        qudit_dimension, num_qudits * 2, np.ravel(qudit_unitary)
+    )
+    # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
+    pad_qubits_vector = qudit_to_qubit_state(
+        qudit_dimension, num_qudits, np.zeros(qudit_dimension**num_qudits), 1
+    )
+    # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the
+    # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map
+    # to identity, preserving unitarity.
+    return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(
+        pad_qubits_vector
+    )
+
+
+def qubit_to_qudit_unitary(
+    qudit_dimension: int,
+    num_qudits: int,
+    qubit_unitary: np.ndarray,
+):
+    """Converts a m-qubit-per-qudit unitary to a qudit-space quantum unitary."""
+    mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
+    # Treat unitary as a `num_qudits*2` qudit system state vector.
+    effective_num_qudits = num_qudits * 2
+    # Reshape the state vector to a `num_qudits*2` rank tensor.
+    unitary_tensor = qubit_unitary.reshape((mbit_dimension,) * effective_num_qudits)
+    # Shrink the number of elements in each dimension up to the qudit_dimension, ignoring the rest.
+    trimmed_unitary_tensor = unitary_tensor[
+        (slice(qudit_dimension),) * effective_num_qudits
+    ]
+    # Return a flat unitary view of the final tensor.
+    return trimmed_unitary_tensor.reshape(
+        qudit_dimension**num_qudits, qudit_dimension**num_qudits
+    )
diff --git a/unitary/alpha/qudit_state_transform_test.py b/unitary/alpha/qudit_state_transform_test.py
new file mode 100644
index 00000000..bab0c4e1
--- /dev/null
+++ b/unitary/alpha/qudit_state_transform_test.py
@@ -0,0 +1,53 @@
+import pytest
+import numpy as np
+from unitary.alpha import qudit_state_transform
+
+
+@pytest.mark.parametrize("qudit_dim", range(10))
+@pytest.mark.parametrize("num_qudits", range(4))
+def test_qudit_state_and_unitary_transform_equivalence(qudit_dim, num_qudits):
+    qudit_state_space = qudit_dim**num_qudits
+    qudit_unitary_shape = (qudit_state_space, qudit_state_space)
+    # Run each configuration with 3 random states and unitaries.
+    for i in range(3):
+        # Random complex state vector in the qudit space.
+        random_state = np.random.rand(qudit_state_space) + 1j * np.random.rand(
+            qudit_state_space
+        )
+        # Random complex unitary in the qudit space.
+        random_unitary = np.random.rand(*qudit_unitary_shape) + 1j * np.random.rand(
+            *qudit_unitary_shape
+        )
+        # Apply the unitary on the state vector.
+        expected_product = np.matmul(random_unitary, random_state)
+        # Qubit space representation of the qudit state vector.
+        transformed_state = qudit_state_transform.qudit_to_qubit_state(
+            qudit_dim, num_qudits, random_state
+        )
+        # Qubit space representation of the qudit unitary.
+        transformed_unitary = qudit_state_transform.qudit_to_qubit_unitary(
+            qudit_dim, num_qudits, random_unitary, memoize=True
+        )
+        # Apply the transformed unitary on the transformed state vector.
+        transformed_product = np.matmul(transformed_unitary, transformed_state)
+        # Convert the transformed product back to the qudit space.
+        product_in_qudit_space = qudit_state_transform.qubit_to_qudit_state(
+            qudit_dim, num_qudits, transformed_product
+        )
+        # Assert that the transform back from qubit space is the inverse of the transform to qubit
+        # space.
+        np.testing.assert_allclose(
+            qudit_state_transform.qubit_to_qudit_state(
+                qudit_dim, num_qudits, transformed_state
+            ),
+            random_state,
+        )
+        np.testing.assert_allclose(
+            qudit_state_transform.qubit_to_qudit_unitary(
+                qudit_dim, num_qudits, transformed_unitary
+            ),
+            random_unitary,
+        )
+        # Assert that the operations in the qubit space are equivalent to the operations in the
+        # qudit space.
+        np.testing.assert_allclose(product_in_qudit_space, expected_product)
diff --git a/unitary/alpha/qudit_to_qubit_transform.py b/unitary/alpha/qudit_to_qubit_transform.py
deleted file mode 100644
index f25ecfd6..00000000
--- a/unitary/alpha/qudit_to_qubit_transform.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# Copyright 2022 Google
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import numpy as np
-
-
-def _nearest_power_of_two_ceiling(x: int) -> int:
-    exp = 0
-    while x:
-        x = x >> 1
-        exp += 1
-    return 2**exp
-
-
-def qudit_to_qubit_state_impl(
-    dim_single_qudit: int,
-    num_qudits: int,
-    qudit_state_vector: np.ndarray,
-    pad_value: np.complex_ = 0,
-):
-    """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
-    # Reshape the state vector to a `num_qudits` rank tensor.
-    state_tensor = qudit_state_vector.reshape((dim_single_qudit,) * num_qudits)
-    # Number of extra elements needed in each dimension if represented using qubits.
-    padding_amount = _nearest_power_of_two_ceiling(dim_single_qudit) - dim_single_qudit
-    # Expand the number of elements in each dimension by the padding_amount. Fill
-    # the new elements with the pad_value.
-    padded_state_tensor = np.pad(
-        state_tensor, (0, padding_amount), constant_values=pad_value
-    )
-    # Return a flattened state vector view of the final tensor.
-    return np.ravel(padded_state_tensor)
-
-
-def qudit_to_qubit_from_map(
-    qudit_state_vector: np.ndarray, qubit_to_qudit_index_map: np.ndarray
-):
-    return [
-        qudit_state_vector[qubit_to_qudit_index_map[j] - 1]
-        if qubit_to_qudit_index_map[j]
-        else 0
-        for j in range(len(qubit_to_qudit_index_map))
-    ]
-
-
-def qudit_to_qubit_state(
-    dim_single_qudit: int,
-    num_qudits: int,
-    qudit_state_vectors: list[np.ndarray],
-) -> list[np.ndarray]:
-    qubit_to_qudit_index_map = qudit_to_qubit_state_impl(
-        dim_single_qudit,
-        num_qudits,
-        np.array([i + 1 for i in range(dim_single_qudit**num_qudits)]),
-    )
-    return [
-        qudit_to_qubit_from_map(state, qubit_to_qudit_index_map)
-        for state in qudit_state_vectors
-    ]
-
-
-def qudit_to_qubit_unitary(
-    dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray
-) -> np.ndarray:
-    qubit_to_qudit_index_map = qudit_to_qubit_state_impl(
-        dim_single_qudit,
-        num_qudits,
-        np.array([i + 1 for i in range(dim_single_qudit**num_qudits)]),
-    )
-    result = np.identity(dim_single_qudit**num_qudits)
-    for i, j in (
-        range(len(qubit_to_qudit_index_map)),
-        range(len(qubit_to_qudit_index_map)),
-    ):
-        result[i][j] = (
-            qudit_unitary[qubit_to_qudit_index_map[i] - 1][
-                qubit_to_qudit_index_map[j] - 1
-            ]
-            if qubit_to_qudit_index_map[i] and qubit_to_qudit_index_map[j]
-            else 0
-        )
-
-
-def qudit_to_qubit_unitary(
-    dim_single_qudit: int, num_qudits: int, qudit_unitary: np.ndarray
-):
-    """This function converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
-    num_qubits_per_qudit = np.ceil(np.log(dim_single_qudit) / np.log(2))
-    num_qubits = num_qubits_per_qudit * num_qudits
-    dim_qubit_space = int(2**num_qubits)
-
-    # Treat the unitary as a num_qubits^2 system's state vector and represent it using qubits (pad with 0s).
-    padded_unitary = qudit_to_qubit_state(
-        dim_single_qudit, num_qudits**2, np.ravel(qudit_unitary)
-    )
-    # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
-    pad_qubits_vector = qudit_to_qubit_state(
-        dim_single_qudit, num_qudits, np.zeros(dim_single_qudit**num_qudits), 1
-    )
-    # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the
-    # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map to identity,
-    # preserving unitarity.
-    return padded_unitary.reshape(dim_qubit_space, dim_qubit_space) + np.diag(
-        pad_qubits_vector
-    )

From 818b13ac92f1f69cd8502babc9305272955967b6 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Mon, 24 Oct 2022 00:25:11 -0700
Subject: [PATCH 06/12] Remove unintentional change to ipynb

---
 docs/quantum_chess/concepts.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/quantum_chess/concepts.ipynb b/docs/quantum_chess/concepts.ipynb
index fda17746..9e52cdb3 100644
--- a/docs/quantum_chess/concepts.ipynb
+++ b/docs/quantum_chess/concepts.ipynb
@@ -1625,7 +1625,7 @@
     "    print(f\"Qubit {q} maps to {dynamic_mapping[q]}\")\n",
     "dynamically_mapped_circuit = cirq.optimize_for_target_gateset(\n",
     "    exclusion_with_sqrt_iswaps, gateset=cirq.SqrtIswapTargetGateset()\n",
-    ").transform_qubits(qubit_map=lambda q: dynamic_mapping[q])\n",
+    ").transform_qubits(qubit_map=lambda q: dynamic_mapping[q])\n",    
     "print(dynamically_mapped_circuit)"
    ]
   },

From 594ad9347e6c4c0bae0a274faef96e5bcd2cfde1 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Mon, 24 Oct 2022 00:46:10 -0700
Subject: [PATCH 07/12] Remove unneeded (index+1)s

---
 unitary/alpha/qudit_state_transform.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
index 89ac50cc..7171818c 100644
--- a/unitary/alpha/qudit_state_transform.py
+++ b/unitary/alpha/qudit_state_transform.py
@@ -81,14 +81,12 @@ def qudit_to_qubit_unitary(
         d_to_b_index_map = qubit_to_qudit_state(
             qudit_dimension,
             num_qudits,
-            np.array([i + 1 for i in range(dim_qubit_space)]),
+            np.arange(dim_qubit_space),
         )
         result = np.identity(dim_qubit_space, dtype=qudit_unitary.dtype)
         iter_range = range(qudit_dimension**num_qudits)
         for i, j in itertools.product(iter_range, iter_range):
-            result[d_to_b_index_map[i] - 1][d_to_b_index_map[j] - 1] = qudit_unitary[i][
-                j
-            ]
+            result[d_to_b_index_map[i]][d_to_b_index_map[j]] = qudit_unitary[i][j]
         return result
 
     # Treat the unitary as a num_qudits^2 system's state vector and represent it using qubits (pad

From 47bcc50dda779eedbeb79a01eb177b8709a7d5a4 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Tue, 25 Oct 2022 00:20:15 -0700
Subject: [PATCH 08/12] More clean up, more tests, add to __init__.py

---
 unitary/alpha/__init__.py                   |   7 ++
 unitary/alpha/qudit_state_transform.py      |  98 +++++++++++++---
 unitary/alpha/qudit_state_transform_test.py | 119 +++++++++++++++++++-
 3 files changed, 207 insertions(+), 17 deletions(-)

diff --git a/unitary/alpha/__init__.py b/unitary/alpha/__init__.py
index 24df674a..93bdf1fb 100644
--- a/unitary/alpha/__init__.py
+++ b/unitary/alpha/__init__.py
@@ -13,6 +13,13 @@
 # limitations under the License.
 #
 
+from unitary.alpha.qudit_state_transform import (
+    qubit_to_qudit_state,
+    qubit_to_qudit_unitary,
+    qudit_to_qubit_state,
+    qudit_to_qubit_unitary,
+)
+
 from unitary.alpha.quantum_world import (
     QuantumWorld,
 )
diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
index 7171818c..1690547a 100644
--- a/unitary/alpha/qudit_state_transform.py
+++ b/unitary/alpha/qudit_state_transform.py
@@ -15,21 +15,15 @@
 import itertools
 
 import numpy as np
-from typing import List
 
 
 def _nearest_power_of_two_ceiling(qudit_dim: int) -> int:
+    """Returns the smallest power of two greater than or equal to qudit_dim."""
     if qudit_dim == 0:
         return 0
-    # Max index of the single qudit state.
-    x = qudit_dim - 1
-    # Number of (qu)bits needed to represent the max index.
     bits = 0
-    while x:
-        x = x >> 1
+    while 2**bits < qudit_dim:
         bits += 1
-    # Total number of states in the qubit representation i.e. Dimension of the single qudit
-    # post-conversion.
     return 2**bits
 
 
@@ -38,8 +32,24 @@ def qudit_to_qubit_state(
     num_qudits: int,
     qudit_state_vector: np.ndarray,
     pad_value: np.complex_ = 0,
-):
-    """Converts a qudit-space quantum state vector to m-qubit-per-qudit column vector."""
+) -> np.ndarray:
+    """Converts a qudit-space quantum state vector to m-qubit-per-qudit column vector.
+
+    Each qudit is replaced by a set of qubits. Since the set of qubits can represent a larger state
+    space than the qudit, the state vector needs to be padded with 0s for those extra elements.
+
+    Args:
+        qudit_dimension: The dimension of a single qudit i.e. the number of states it can
+          represent.
+        num_qudits: The number of qudits in the given state vector.
+        qudit_state_vector: A flat, 1-D numpy array representing the state vector in qudit form.
+        pad_value: The value to set for the elements in the extra state space created in the
+          conversion. This field is mostly private, mostly for use by other methods. Do not set
+          this unless you really need to.
+
+    Returns:
+        A flat numpy array representing the input state vector using qubits.
+    """
     # Reshape the state vector to a `num_qudits` rank tensor.
     state_tensor = qudit_state_vector.reshape((qudit_dimension,) * num_qudits)
     # Number of extra elements needed in each dimension if represented using qubits.
@@ -57,8 +67,23 @@ def qubit_to_qudit_state(
     qudit_dimension: int,
     num_qudits: int,
     qubit_state_vector: np.ndarray,
-):
-    """Converts a m-qubit-per-qudit column vector to a qudit-space quantum state vector."""
+) -> np.ndarray:
+    """Converts a m-qubit-per-qudit column vector to a qudit-space quantum state vector.
+
+    Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger
+    state space than the qudit, the state vector needs to be sliced up to the qudit length in each
+    dimension.
+
+    Args:
+        qudit_dimension: The dimension of a single qudit i.e. the number of states it can
+          represent.
+        num_qudits: The number of qudits in the given/output state vector.
+        qubit_state_vector: A flat, 1-D numpy array representing the state vector in an
+          m-qubit-per-qudit form.
+
+    Returns:
+        A flat numpy array representing the input state vector using qudits.
+    """
     mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
     # Reshape the state vector to a `num_qudits` rank tensor.
     state_tensor = qubit_state_vector.reshape((mbit_dimension,) * num_qudits)
@@ -74,18 +99,46 @@ def qudit_to_qubit_unitary(
     qudit_unitary: np.ndarray,
     memoize: bool = False,
 ) -> np.ndarray:
-    """Converts a qudit-space quantum unitary to m-qubit-per-qudit unitary."""
+    """Converts a qudit-space quantum unitary to m-qubit-per-qudit unitary.
+
+    Each qudit is replaced by a set of qubits. Since the set of qubits can represent a larger state
+    space than the qudit, the unitary needs to be padded with 0s for those extra elements. A
+    unitary is treated similar to a 2*num_qudits system's state vector and padded using the state
+    vector protocol. The resulting unitary is updated to have the extra dimensions map to
+    themselves (identity) to preserve unitarity.
+
+    Args:
+        qudit_dimension: The dimension of a single qudit i.e. the number of states it can
+            represent.
+        num_qudits: The number of qudits in the given unitary.
+        qudit_unitary: A 2-D numpy array representing the unitary in qudit form.
+        memoize: Currently, this method has two independent implementations. If memoize is True, an
+            alternate implementation than above is used. A special state vector is passed to the
+            state vector protocol to get a mapping from qudit state indices to qubit state indices.
+            This mapping is then iteratively applied to the input unitary's elements.
+
+    Returns:
+        A numpy array representing the input unitary using qubits.
+    """
     dim_qubit_space = _nearest_power_of_two_ceiling(qudit_dimension) ** num_qudits
 
     if memoize:
+        # Perform the transform of the below array from qubit to qudit space so that the indices
+        # represent the position in qudit space and the values represent the position in the qubit
+        # space.
         d_to_b_index_map = qubit_to_qudit_state(
             qudit_dimension,
             num_qudits,
+            # An array of ints from 0 to dim_qubit_space. Each element represents the original
+            # index.
             np.arange(dim_qubit_space),
         )
+        # Initialize the result to the identity unitary in the qubit space.
         result = np.identity(dim_qubit_space, dtype=qudit_unitary.dtype)
+        # Iterate over each element in the qudit space dimension x qudit space dimension.
         iter_range = range(qudit_dimension**num_qudits)
         for i, j in itertools.product(iter_range, iter_range):
+            # Use the index map to populate the appropriate element in the qubit representation.
             result[d_to_b_index_map[i]][d_to_b_index_map[j]] = qudit_unitary[i][j]
         return result
 
@@ -94,7 +147,8 @@ def qudit_to_qubit_unitary(
     padded_unitary = qudit_to_qubit_state(
         qudit_dimension, num_qudits * 2, np.ravel(qudit_unitary)
     )
-    # A qubit-based state vector with the extra padding bits having 1s and rest having 0s.
+    # A qubit-based state vector with the extra padding bits having 1s and rest having 0s. This
+    # vector marks only the bits that are padded.
     pad_qubits_vector = qudit_to_qubit_state(
         qudit_dimension, num_qudits, np.zeros(qudit_dimension**num_qudits), 1
     )
@@ -111,7 +165,21 @@ def qubit_to_qudit_unitary(
     num_qudits: int,
     qubit_unitary: np.ndarray,
 ):
-    """Converts a m-qubit-per-qudit unitary to a qudit-space quantum unitary."""
+    """Converts a m-qubit-per-qudit unitary to a qudit-space quantum unitary.
+
+    Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger
+    state space than the qudit, the unitary needs to be sliced up to the qudit length in each
+    dimension. A unitary is treated similar to a 2*num_qudits system's state vector.
+
+    Args:
+        qudit_dimension: The dimension of a single qudit i.e. the number of states it can
+          represent.
+        num_qudits: The number of qudits in the given/output unitary.
+        qubit_unitary: A 2-D numpy array representing the unitary in m-qubit-per-qudit form.
+
+    Returns:
+        A numpy array representing the input unitary using qudits.
+    """
     mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
     # Treat unitary as a `num_qudits*2` qudit system state vector.
     effective_num_qudits = num_qudits * 2
diff --git a/unitary/alpha/qudit_state_transform_test.py b/unitary/alpha/qudit_state_transform_test.py
index bab0c4e1..81678d75 100644
--- a/unitary/alpha/qudit_state_transform_test.py
+++ b/unitary/alpha/qudit_state_transform_test.py
@@ -24,9 +24,9 @@ def test_qudit_state_and_unitary_transform_equivalence(qudit_dim, num_qudits):
         transformed_state = qudit_state_transform.qudit_to_qubit_state(
             qudit_dim, num_qudits, random_state
         )
-        # Qubit space representation of the qudit unitary.
+        # Qubit space representation of the qudit unitary. Alternate between memoizing or not.
         transformed_unitary = qudit_state_transform.qudit_to_qubit_unitary(
-            qudit_dim, num_qudits, random_unitary, memoize=True
+            qudit_dim, num_qudits, random_unitary, memoize=(i % 2)
         )
         # Apply the transformed unitary on the transformed state vector.
         transformed_product = np.matmul(transformed_unitary, transformed_state)
@@ -51,3 +51,118 @@ def test_qudit_state_and_unitary_transform_equivalence(qudit_dim, num_qudits):
         # Assert that the operations in the qubit space are equivalent to the operations in the
         # qudit space.
         np.testing.assert_allclose(product_in_qudit_space, expected_product)
+
+
+@pytest.mark.parametrize(
+    "qudit_dim, num_qudits, qudit_representation, qubit_representation",
+    [
+        (
+            3,
+            2,
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 1]),
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]),
+        ),
+        (
+            4,
+            2,
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]),
+            np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0]),
+        ),
+        (
+            3,
+            2,
+            np.array([1, 0, 0, 0, 0, 0, 0, 0, 1], dtype=np.complex_),
+            np.array(
+                [
+                    1,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                    1,
+                    0,
+                    0,
+                    0,
+                    0,
+                    0,
+                ]
+            ),
+        ),
+    ],
+)
+def test_specific_transformations_vectors(
+    qudit_dim, num_qudits, qudit_representation, qubit_representation
+):
+    transformed_vector = qudit_state_transform.qudit_to_qubit_state(
+        qudit_dim, num_qudits, qudit_representation
+    )
+    np.testing.assert_allclose(transformed_vector, qubit_representation)
+    untransformed_vector = qudit_state_transform.qubit_to_qudit_state(
+        qudit_dim, num_qudits, transformed_vector
+    )
+    np.testing.assert_allclose(untransformed_vector, qudit_representation)
+
+
+a = complex(0.5, 0.5)
+b = complex(0.5, -0.5)
+
+
+@pytest.mark.parametrize(
+    "qudit_dim, num_qudits, qudit_representation, qubit_representation",
+    [
+        # Qutrit square root of iSwap.
+        (
+            3,
+            2,
+            np.array(
+                [
+                    [1, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, a, 0, b, 0, 0, 0, 0, 0],
+                    [0, 0, a, 0, 0, 0, b, 0, 0],
+                    [0, b, 0, a, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 1, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 1, 0, 0, 0],
+                    [0, 0, b, 0, 0, 0, a, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 1, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 1],
+                ]
+            ),
+            np.array(
+                [
+                    [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, a, 0, 0, b, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, a, 0, 0, 0, 0, 0, b, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, b, 0, 0, a, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, b, 0, 0, 0, 0, 0, a, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
+                ]
+            ),
+        ),
+    ],
+)
+def test_specific_transformations_unitaries(
+    qudit_dim, num_qudits, qudit_representation, qubit_representation
+):
+    transformed_unitary = qudit_state_transform.qudit_to_qubit_unitary(
+        qudit_dim, num_qudits, qudit_representation
+    )
+    np.testing.assert_allclose(transformed_unitary, qubit_representation)
+    untransformed_unitary = qudit_state_transform.qubit_to_qudit_unitary(
+        qudit_dim, num_qudits, transformed_unitary
+    )
+    np.testing.assert_allclose(untransformed_unitary, qudit_representation)

From abd0c7ce112e1e337844d0a0030b9846d0d95342 Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Tue, 25 Oct 2022 00:25:32 -0700
Subject: [PATCH 09/12] Attribution

---
 unitary/alpha/qudit_state_transform.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
index 1690547a..2028f03e 100644
--- a/unitary/alpha/qudit_state_transform.py
+++ b/unitary/alpha/qudit_state_transform.py
@@ -37,6 +37,8 @@ def qudit_to_qubit_state(
 
     Each qudit is replaced by a set of qubits. Since the set of qubits can represent a larger state
     space than the qudit, the state vector needs to be padded with 0s for those extra elements.
+    Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and
+    https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
@@ -73,6 +75,8 @@ def qubit_to_qudit_state(
     Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger
     state space than the qudit, the state vector needs to be sliced up to the qudit length in each
     dimension.
+    Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and
+    https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
@@ -106,6 +110,8 @@ def qudit_to_qubit_unitary(
     unitary is treated similar to a 2*num_qudits system's state vector and padded using the state
     vector protocol. The resulting unitary is updated to have the extra dimensions map to
     themselves (identity) to preserve unitarity.
+    Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and
+    https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
@@ -170,6 +176,8 @@ def qubit_to_qudit_unitary(
     Each qudit was replaced by a set of qubits. Since the set of qubits could represent a larger
     state space than the qudit, the unitary needs to be sliced up to the qudit length in each
     dimension. A unitary is treated similar to a 2*num_qudits system's state vector.
+    Based on https://drive.google.com/file/d/1n1Ym7JdnM44NnvNQDUXSk5rBLTQZzJ9t and
+    https://colab.research.google.com/drive/1MC6D4FyOXG0RjvzyV9IFcdRsLbUwwAcx
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can

From b072dbfc5d7529b6eae1f92ac9b324f8af065c4a Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Tue, 25 Oct 2022 00:31:13 -0700
Subject: [PATCH 10/12] helper method cleanup

---
 unitary/alpha/qudit_state_transform.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
index 2028f03e..f8daa160 100644
--- a/unitary/alpha/qudit_state_transform.py
+++ b/unitary/alpha/qudit_state_transform.py
@@ -21,10 +21,10 @@ def _nearest_power_of_two_ceiling(qudit_dim: int) -> int:
     """Returns the smallest power of two greater than or equal to qudit_dim."""
     if qudit_dim == 0:
         return 0
-    bits = 0
-    while 2**bits < qudit_dim:
-        bits += 1
-    return 2**bits
+    result = 1
+    while result < qudit_dim:
+        result = result << 1
+    return result
 
 
 def qudit_to_qubit_state(

From 3235ffca720ede5d180b8edf82e14342169244de Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Mon, 31 Oct 2022 17:54:25 -0700
Subject: [PATCH 11/12] remove numpy req and add expected shapes for args

---
 requirements.txt                       |  1 -
 unitary/alpha/qudit_state_transform.py | 40 ++++++++++++++++----------
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 1210459c..a76f17a4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,6 @@ cirq-core>=0.15.0
 cirq-google>=0.15.0
 # When changing Cirq requirements be sure to update dev_tools/write-ci-requirements.py
 
-numpy
 seaborn
 sphinx
 ipython
diff --git a/unitary/alpha/qudit_state_transform.py b/unitary/alpha/qudit_state_transform.py
index f8daa160..f9a8d522 100644
--- a/unitary/alpha/qudit_state_transform.py
+++ b/unitary/alpha/qudit_state_transform.py
@@ -31,7 +31,7 @@ def qudit_to_qubit_state(
     qudit_dimension: int,
     num_qudits: int,
     qudit_state_vector: np.ndarray,
-    pad_value: np.complex_ = 0,
+    _pad_value: np.complex_ = 0,
 ) -> np.ndarray:
     """Converts a qudit-space quantum state vector to m-qubit-per-qudit column vector.
 
@@ -42,24 +42,26 @@ def qudit_to_qubit_state(
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
-          represent.
+            represent.
         num_qudits: The number of qudits in the given state vector.
-        qudit_state_vector: A flat, 1-D numpy array representing the state vector in qudit form.
-        pad_value: The value to set for the elements in the extra state space created in the
-          conversion. This field is mostly private, mostly for use by other methods. Do not set
-          this unless you really need to.
+        qudit_state_vector: A numpy array representing the state vector in qudit form.
+            Expected shape: `(qudit_dimension ^ num_qudits,)`.
+        _pad_value: The value to set for the elements in the extra state space created in the
+            conversion. This field is mostly private, mostly for use by other methods. Do not set
+            this unless you really need to.
 
     Returns:
-        A flat numpy array representing the input state vector using qubits.
+        A flat numpy array representing the input state vector using m-qubits-per-qudit.
+            Expected shape: `((2 ^ m) ^ num_qudits,)`.
     """
     # Reshape the state vector to a `num_qudits` rank tensor.
     state_tensor = qudit_state_vector.reshape((qudit_dimension,) * num_qudits)
     # Number of extra elements needed in each dimension if represented using qubits.
     padding_amount = _nearest_power_of_two_ceiling(qudit_dimension) - qudit_dimension
     # Expand the number of elements in each dimension by the padding_amount. Fill
-    # the new elements with the pad_value.
+    # the new elements with the _pad_value.
     padded_state_tensor = np.pad(
-        state_tensor, pad_width=(0, padding_amount), constant_values=pad_value
+        state_tensor, pad_width=(0, padding_amount), constant_values=_pad_value
     )
     # Return a flattened state vector view of the final tensor.
     return np.ravel(padded_state_tensor)
@@ -80,13 +82,14 @@ def qubit_to_qudit_state(
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
-          represent.
+            represent.
         num_qudits: The number of qudits in the given/output state vector.
-        qubit_state_vector: A flat, 1-D numpy array representing the state vector in an
-          m-qubit-per-qudit form.
+        qubit_state_vector: A numpy array representing the state vector in an
+            m-qubit-per-qudit form. Expected shape: `((2 ^ m) ^ num_qudits,)`.
 
     Returns:
         A flat numpy array representing the input state vector using qudits.
+            Expected shape: `(qudit_dimension ^ num_qudits,)`.
     """
     mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
     # Reshape the state vector to a `num_qudits` rank tensor.
@@ -118,13 +121,15 @@ def qudit_to_qubit_unitary(
             represent.
         num_qudits: The number of qudits in the given unitary.
         qudit_unitary: A 2-D numpy array representing the unitary in qudit form.
+            Expected shape: `(qudit_dimension ^ num_qudits, qudit_dimension ^ num_qudits)`.
         memoize: Currently, this method has two independent implementations. If memoize is True, an
             alternate implementation than above is used. A special state vector is passed to the
             state vector protocol to get a mapping from qudit state indices to qubit state indices.
             This mapping is then iteratively applied to the input unitary's elements.
 
     Returns:
-        A numpy array representing the input unitary using qubits.
+        A numpy array representing the input unitary using m-qubits-per-qudit.
+            Expected shape: `((2 ^ m) ^ num_qudits, (2 ^ m) ^ num_qudits)`.
     """
     dim_qubit_space = _nearest_power_of_two_ceiling(qudit_dimension) ** num_qudits
 
@@ -156,7 +161,10 @@ def qudit_to_qubit_unitary(
     # A qubit-based state vector with the extra padding bits having 1s and rest having 0s. This
     # vector marks only the bits that are padded.
     pad_qubits_vector = qudit_to_qubit_state(
-        qudit_dimension, num_qudits, np.zeros(qudit_dimension**num_qudits), 1
+        qudit_dimension,
+        num_qudits,
+        np.zeros(qudit_dimension**num_qudits),
+        _pad_value=1,
     )
     # Reshape the padded unitary to the final shape and add a diagonal matrix corresponding to the
     # pad_qubits_vector. This addition ensures that the invalid states with the "padding" bits map
@@ -181,12 +189,14 @@ def qubit_to_qudit_unitary(
 
     Args:
         qudit_dimension: The dimension of a single qudit i.e. the number of states it can
-          represent.
+            represent.
         num_qudits: The number of qudits in the given/output unitary.
         qubit_unitary: A 2-D numpy array representing the unitary in m-qubit-per-qudit form.
+            Expected shape: `((2 ^ m) ^ num_qudits, (2 ^ m) ^ num_qudits)`.
 
     Returns:
         A numpy array representing the input unitary using qudits.
+            Expected shape: `(qudit_dimension ^ num_qudits, qudit_dimension ^ num_qudits)`.
     """
     mbit_dimension = _nearest_power_of_two_ceiling(qudit_dimension)
     # Treat unitary as a `num_qudits*2` qudit system state vector.

From 0487b5da877b87f962ea1fcbe6cd740fcf08625f Mon Sep 17 00:00:00 2001
From: Smit <smitsanghavi22@gmail.com>
Date: Wed, 2 Nov 2022 18:36:01 -0700
Subject: [PATCH 12/12] Hack TicTacToe to work using qubits via the qudit
 transform helpers

---
 unitary/alpha/qudit_gates.py                | 31 ++++++++++++++-------
 unitary/examples/tictactoe/enums.py         |  1 +
 unitary/examples/tictactoe/tic_tac_split.py |  7 +++--
 unitary/examples/tictactoe/tic_tac_toe.py   |  2 ++
 4 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/unitary/alpha/qudit_gates.py b/unitary/alpha/qudit_gates.py
index fa0a7035..11a45e39 100644
--- a/unitary/alpha/qudit_gates.py
+++ b/unitary/alpha/qudit_gates.py
@@ -14,6 +14,17 @@
 #
 import numpy as np
 import cirq
+from unitary.alpha.qudit_state_transform import qudit_to_qubit_unitary
+
+
+def _nearest_power_of_two_ceiling(qudit_dim: int) -> int:
+    """Returns the smallest power of two greater than or equal to qudit_dim."""
+    if qudit_dim == 0:
+        return 0
+    result = 1
+    while result < qudit_dim:
+        result = result << 1
+    return result
 
 
 class QuditXGate(cirq.Gate):
@@ -36,7 +47,7 @@ def __init__(
         self.destination_state = destination_state
 
     def _qid_shape_(self):
-        return (self.dimension,)
+        return (_nearest_power_of_two_ceiling(self.dimension),)
 
     def _unitary_(self):
         arr = np.zeros((self.dimension, self.dimension))
@@ -45,7 +56,7 @@ def _unitary_(self):
         for i in range(self.dimension):
             if i != self.source_state and i != self.destination_state:
                 arr[i, i] = 1
-        return arr
+        return qudit_to_qubit_unitary(self.dimension, 1, arr)
 
     def _circuit_diagram_info_(self, args):
         return f"X({self.source_state}_{self.destination_state})"
@@ -63,13 +74,13 @@ def __init__(self, dimension: int, addend: int = 1):
         self.addend = addend
 
     def _qid_shape_(self):
-        return (self.dimension,)
+        return (_nearest_power_of_two_ceiling(self.dimension),)
 
     def _unitary_(self):
         arr = np.zeros((self.dimension, self.dimension))
         for i in range(self.dimension):
             arr[(i + self.addend) % self.dimension, i] = 1
-        return arr
+        return qudit_to_qubit_unitary(self.dimension, 1, arr)
 
     def _circuit_diagram_info_(self, args):
         return f"[+{self.addend}]"
@@ -99,7 +110,7 @@ def __init__(self, dimension: int, control_state: int = 1, state: int = 1):
         self.control_state = control_state
 
     def _qid_shape_(self):
-        return (self.dimension, self.dimension)
+        return (_nearest_power_of_two_ceiling(self.dimension), _nearest_power_of_two_ceiling(self.dimension))
 
     def _unitary_(self):
         size = self.dimension * self.dimension
@@ -111,7 +122,7 @@ def _unitary_(self):
             for y in range(self.dimension):
                 if x != self.control_state or (y != self.state and y != 0):
                     arr[x * self.dimension + y, x * self.dimension + y] = 1
-        return arr
+        return qudit_to_qubit_unitary(self.dimension, 2, arr)
 
 
 class QuditSwapPowGate(cirq.Gate):
@@ -134,7 +145,7 @@ def __init__(self, dimension: int, exponent: float = 1):
         self.exponent = exponent
 
     def _qid_shape_(self):
-        return (self.dimension, self.dimension)
+        return (_nearest_power_of_two_ceiling(self.dimension), _nearest_power_of_two_ceiling(self.dimension))
 
     def _unitary_(self):
         size = self.dimension * self.dimension
@@ -149,7 +160,7 @@ def _unitary_(self):
                 diag = g * np.cos(np.pi * self.exponent / 2)
                 arr[x * self.dimension + y, y * self.dimension + x] = coeff
                 arr[x * self.dimension + y, x * self.dimension + y] = diag
-        return arr
+        return qudit_to_qubit_unitary(self.dimension, 2, arr)
 
     def _circuit_diagram_info_(self, args):
         if not args.use_unicode_characters:
@@ -181,7 +192,7 @@ def __init__(self, dimension: int, exponent: float = 1):
         self.exponent = exponent
 
     def _qid_shape_(self):
-        return (self.dimension, self.dimension)
+        return (_nearest_power_of_two_ceiling(self.dimension), _nearest_power_of_two_ceiling(self.dimension))
 
     def _unitary_(self):
         size = self.dimension * self.dimension
@@ -196,7 +207,7 @@ def _unitary_(self):
                 arr[x * self.dimension + y, y * self.dimension + x] = coeff
                 arr[x * self.dimension + y, x * self.dimension + y] = diag
 
-        return arr
+        return qudit_to_qubit_unitary(self.dimension, 2, arr)
 
     def _circuit_diagram_info_(self, args):
         return cirq.CircuitDiagramInfo(
diff --git a/unitary/examples/tictactoe/enums.py b/unitary/examples/tictactoe/enums.py
index 8bbd9ac8..c0ac3f30 100644
--- a/unitary/examples/tictactoe/enums.py
+++ b/unitary/examples/tictactoe/enums.py
@@ -19,6 +19,7 @@ class TicTacSquare(enum.Enum):
     EMPTY = 0
     X = 1
     O = 2
+    PADDING = 3
 
 
 class TicTacResult(enum.Enum):
diff --git a/unitary/examples/tictactoe/tic_tac_split.py b/unitary/examples/tictactoe/tic_tac_split.py
index 897437e9..dc894660 100644
--- a/unitary/examples/tictactoe/tic_tac_split.py
+++ b/unitary/examples/tictactoe/tic_tac_split.py
@@ -20,6 +20,7 @@
 from unitary.alpha import QuantumEffect, QuantumObject
 from unitary.alpha.qudit_gates import QuditXGate, QuditISwapPowGate
 from unitary.examples.tictactoe.enums import TicTacSquare, TicTacRules
+from unitary.alpha.qudit_state_transform import qudit_to_qubit_unitary
 
 
 class QuditSplitGate(cirq.Gate):
@@ -40,7 +41,7 @@ def __init__(self, square: TicTacSquare):
             raise ValueError("Not a valid square: {self.square}")
 
     def _qid_shape_(self):
-        return (3, 3)
+        return (4, 4)
 
     def _unitary_(self):
         arr = np.zeros((9, 9), dtype=np.complex64)
@@ -59,7 +60,7 @@ def _unitary_(self):
             arr[3, 1] = coeff
             arr[3, 3] = diag
             arr[1, 1] = diag
-        return arr
+        return qudit_to_qubit_unitary(3, 2, arr)
 
     def _circuit_diagram_info_(self, args):
         if not args.use_unicode_characters:
@@ -81,7 +82,7 @@ def __init__(self, tic_tac_type: TicTacSquare, rules: TicTacRules):
         self.rules = rules
 
     def num_dimension(self) -> Optional[int]:
-        return 3
+        return 4
 
     def num_objects(self) -> Optional[int]:
         return 2
diff --git a/unitary/examples/tictactoe/tic_tac_toe.py b/unitary/examples/tictactoe/tic_tac_toe.py
index db22fdfc..a8b2a61b 100644
--- a/unitary/examples/tictactoe/tic_tac_toe.py
+++ b/unitary/examples/tictactoe/tic_tac_toe.py
@@ -228,6 +228,8 @@ def print(self) -> str:
         output = "\n"
         for row in range(3):
             for mark in TicTacSquare:
+                if mark == TicTacSquare.PADDING:
+                    continue
                 output += " "
                 for col in range(3):
                     idx = row * 3 + col