PR Draft for Qudit inferface

tensorcircuit/applications/layers.py

@@ -13,7 +13,8 @@
 
 from ..circuit import Circuit
 from ..densitymatrix import DMCircuit
-from ..gates import num_to_tensor, array_to_tensor, _swap_matrix
+from .. import gates
+from ..gates import num_to_tensor, array_to_tensor
 from ..channels import depolarizingchannel
 from ..abstractcircuit import sgates
 
@@ -25,6 +26,7 @@
     logger.warning(e)
     logger.warning("Therefore some functionality in %s may not work" % __name__)
 
+_swap_matrix = gates._swap_matrix  # type: ignore
 
 thismodule = sys.modules[__name__]
 
@@ -298,7 +300,7 @@ def anyswaplayer_bitflip_mc(
     generate_gate_layer(gate)
     generate_any_gate_layer(gate)
 
-for gates in itertools.product(*[["x", "y", "z"] for _ in range(2)]):
+for gates in itertools.product(*[["x", "y", "z"] for _ in range(2)]):  # type: ignore
     gates = gates[0] + gates[1]  # type: ignore
     generate_double_gate(gates)  # type: ignore
     generate_double_gate_layer(gates)  # type: ignore
@@ -308,7 +310,7 @@ def anyswaplayer_bitflip_mc(
     generate_any_double_gate_layer_bitflip_mc(gates)  # type: ignore
 
 
-for gates in itertools.product(
+for gates in itertools.product(  # type: ignore
     *[["rx", "ry", "rz", "xx", "yy", "zz"] for _ in range(2)]
 ):
     generate_double_layer_block(gates)  # type: ignore
@@ -504,7 +506,7 @@ def f(
         if gate != "H":
             generate_cirq_any_gate_layer(gate)
 
-    for gates in itertools.product(*[["x", "y", "z"] for _ in range(2)]):
+    for gates in itertools.product(*[["x", "y", "z"] for _ in range(2)]):  # type: ignore
         gates = gates[0] + gates[1]  # type: ignore
         generate_cirq_double_gate(gates)  # type: ignore
         generate_cirq_double_gate_layer(gates)  # type: ignore

tensorcircuit/applications/utils.py

@@ -217,15 +217,15 @@ def train_qml_vag(
                 c.exp(  # type: ignore
                     i,
                     (i + 1) % 10,
-                    unitary=array_to_tensor(G._swap_matrix),
+                    unitary=array_to_tensor(G._swap_matrix),  # type: ignore
                     theta=cnnp[3 * epoch + 2, i],
                 )
             for i in range(1, nqubits, 2):
                 # c.swap(i, (i + 1) % 10)  # type: ignore
                 c.exp(  # type: ignore
                     i,
                     (i + 1) % 10,
-                    unitary=array_to_tensor(G._swap_matrix),
+                    unitary=array_to_tensor(G._swap_matrix),  # type: ignore
                     theta=cnnp[3 * epoch + 2, i],
                 )
         for i in range(nqubits):
@@ -287,15 +287,15 @@ def validate_qml_vag(
             c.exp(  # type: ignore
                 i,
                 (i + 1) % 10,
-                unitary=array_to_tensor(G._swap_matrix),
+                unitary=array_to_tensor(G._swap_matrix),  # type: ignore
                 theta=cnnp[3 * epoch + 2, i],
             )
         for i in range(1, nqubits, 2):
             # c.swap(i, (i + 1) % 10)  # type: ignore
             c.exp(  # type: ignore
                 i,
                 (i + 1) % 10,
-                unitary=array_to_tensor(G._swap_matrix),
+                unitary=array_to_tensor(G._swap_matrix),  # type: ignore
                 theta=cnnp[3 * epoch + 2, i],
             )
     for i in range(nqubits):

tensorcircuit/backends/abstract_backend.py

@@ -9,6 +9,7 @@
 from operator import mul
 from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 
+import math
 import numpy as np
 from ..utils import return_partial
 
@@ -405,6 +406,31 @@ def reshape2(self: Any, a: Tensor) -> Tensor:
         a = self.reshape(a, [2 for _ in range(nleg)])
         return a
 
+    def reshaped(self: Any, a: Tensor, d: int) -> Tensor:
+        """
+        Reshape a tensor to the [d, d, ...] shape.
+
+        :param a: Input tensor
+        :type a: Tensor
+        :param d: edge length for each dimension
+        :type d: int
+        :return: the reshaped tensor
+        :rtype: Tensor
+        """
+        if not isinstance(d, int) or d <= 0:
+            raise ValueError("d must be a positive integer.")
+
+        size = self.sizen(a)
+        if size == 0:
+            return self.reshape(a, (0,))
+
+        nleg_float = math.log(size, d)
+        nleg = int(round(nleg_float))
+        if d**nleg != size:
+            raise ValueError(f"cannot reshape: size {size} is not a power of d={d}")
+
+        return self.reshape(a, (d,) * nleg)
+
     def reshapem(self: Any, a: Tensor) -> Tensor:
         """
         Reshape a tensor to the [l, l] shape.
@@ -748,6 +774,36 @@ def mod(self: Any, x: Tensor, y: Tensor) -> Tensor:
             "Backend '{}' has not implemented `mod`.".format(self.name)
         )
 
+    def floor(self: Any, x: Tensor) -> Tensor:
+        """
+        Compute
+
+        :param x: input values
+        :type x: Tensor
+        :return: results
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `floor`.".format(self.name)
+        )
+
+    def clip(self: Any, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        """
+        Clip values of x into [a_min, a_max], preserving dtype & device.
+
+        :param a: input values
+        :type a: Tensor
+        :param a_min: minimum value
+        :type a_min: Tensor
+        :param a_max: maximum value
+        :type a_max: Tensor
+        :return: results
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `clip`.".format(self.name)
+        )
+
     def reverse(self: Any, a: Tensor) -> Tensor:
         """
         return ``a[::-1]``, only 1D tensor is guaranteed for consistent behavior

tensorcircuit/backends/cupy_backend.py

@@ -11,7 +11,7 @@
 import numpy as np
 
 try:
-    from numpy import ComplexWarning
+    from numpy import ComplexWarning  # type: ignore
 except ImportError:  # np2.0 compatibility
     from numpy.exceptions import ComplexWarning  # type: ignore
 

tensorcircuit/backends/jax_backend.py

@@ -12,7 +12,7 @@
 import numpy as np
 
 try:
-    from numpy import ComplexWarning
+    from numpy import ComplexWarning  # type: ignore
 except ImportError:  # np2.0 compatibility
     from numpy.exceptions import ComplexWarning  # type: ignore
 
@@ -342,6 +342,12 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return jnp.mod(x, y)
 
+    def floor(self, x: Tensor) -> Tensor:
+        return jnp.floor(x)
+
+    def clip(self, x: Tensor, lower: Tensor, upper: Tensor) -> Tensor:
+        return jnp.clip(x, lower, upper)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return jnp.right_shift(x, y)
 

tensorcircuit/backends/numpy_backend.py

@@ -11,7 +11,7 @@
 import numpy as np
 
 try:
-    from numpy import ComplexWarning
+    from numpy import ComplexWarning  # type: ignore
 except ImportError:  # np2.0 compatibility
     from numpy.exceptions import ComplexWarning  # type: ignore
 
@@ -234,6 +234,12 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return np.mod(x, y)
 
+    def floor(self, x: Tensor) -> Tensor:
+        return np.floor(x)
+
+    def clip(self, x: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return np.clip(x, a_min, a_max)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return np.right_shift(x, y)
 

tensorcircuit/backends/pytorch_backend.py

@@ -409,6 +409,12 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return torchlib.fmod(x, y)
 
+    def floor(self, x: Tensor) -> Tensor:
+        return torchlib.floor(x)
+
+    def clip(self, x: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return torchlib.clamp(x, a_min, a_max)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return torchlib.bitwise_right_shift(x, y)
 

tensorcircuit/backends/tensorflow_backend.py

@@ -115,7 +115,7 @@ def _random_choice_tf(
     else:
         if not (isinstance(p, tf.Tensor) or isinstance(p, tf.Variable)):
             p = tf.constant(p)
-        dtype = p.dtype
+        dtype = p.dtype  # type: ignore
     shape1 = reduce(mul, shape)
     p_cuml = tf.cumsum(p)
     r = p_cuml[-1] * (1 - g.uniform([shape1], dtype=dtype))
@@ -544,6 +544,9 @@ def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
     def stack(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return tf.stack(a, axis=axis)
 
+    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return tf.clip_by_value(a, a_min, a_max)
+
     def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return tf.concat(a, axis=axis)
 
@@ -612,6 +615,9 @@ def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tenso
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return tf.math.mod(x, y)
 
+    def floor(self, x: Tensor) -> Tensor:
+        return tf.math.floor(x)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return tf.bitwise.right_shift(x, y)