tensorcircuit-nightly 1.0.2.dev20250108__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/gates.py

@@ -34,6 +34,12 @@ one_state = np.array([0.0, 1.0], dtype=npdtype)
 plus_state = 1.0 / np.sqrt(2) * (zero_state + one_state)
 minus_state = 1.0 / np.sqrt(2) * (zero_state - one_state)
 
+# Common elements as np.ndarray objects
+_i00 = np.array([[1.0, 0.0], [0.0, 0.0]])
+_i01 = np.array([[0.0, 1.0], [0.0, 0.0]])
+_i10 = np.array([[0.0, 0.0], [1.0, 0.0]])
+_i11 = np.array([[0.0, 0.0], [0.0, 1.0]])
+
 # Common single qubit gates as np.ndarray objects
 _h_matrix = 1 / np.sqrt(2) * np.array([[1.0, 1.0], [1.0, -1.0]])
 _i_matrix = np.array([[1.0, 0.0], [0.0, 1.0]])
@@ -229,7 +235,7 @@ def num_to_tensor(*num: Union[float, Tensor], dtype: Optional[str] = None) -> An
     # TODO(@YHPeter): fix __doc__ for same function with different names
 
     l = []
-    if not dtype:
+    if dtype is None:
         dtype = dtypestr
     for n in num:
         if not backend.is_tensor(n):
@@ -245,7 +251,7 @@ array_to_tensor = num_to_tensor
 
 
 def gate_wrapper(m: Tensor, n: Optional[str] = None) -> Gate:
-    if not n:
+    if n is None:
         n = "unknowngate"
     m = m.astype(npdtype)
     return Gate(deepcopy(m), name=n)
@@ -255,7 +261,7 @@ class GateF:
     def __init__(
         self, m: Tensor, n: Optional[str] = None, ctrl: Optional[List[int]] = None
     ):
-        if not n:
+        if n is None:
             n = "unknowngate"
         self.m = m
         self.n = n
@@ -310,7 +316,7 @@ class GateF:
 
             return Gate(cu, name="c" + self.n)
 
-        if not self.ctrl:
+        if self.ctrl is None:
             ctrl = [1]
         else:
             ctrl = [1] + self.ctrl
@@ -330,7 +336,7 @@ class GateF:
             # TODO(@refraction-ray): ctrl convention to be finally determined
             return Gate(ocu, name="o" + self.n)
 
-        if not self.ctrl:
+        if self.ctrl is None:
             ctrl = [0]
         else:
             ctrl = [0] + self.ctrl
@@ -349,7 +355,7 @@ class GateVF(GateF):
         n: Optional[str] = None,
         ctrl: Optional[List[int]] = None,
     ):
-        if not n:
+        if n is None:
             n = "unknowngate"
         self.f = f
         self.n = n
@@ -483,7 +489,7 @@ def phase_gate(theta: float = 0) -> Gate:
     :rtype: Gate
     """
     theta = array_to_tensor(theta)
-    i00, i11 = array_to_tensor(np.array([[1, 0], [0, 0]]), np.array([[0, 0], [0, 1]]))
+    i00, i11 = array_to_tensor(_i00, _i11)
     unitary = i00 + backend.exp(1.0j * theta) * i11
     return Gate(unitary)
 
@@ -512,7 +518,7 @@ def get_u_parameter(m: Tensor) -> Tuple[float, float, float]:
     return theta, phi, lbd
 
 
-def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
+def u_gate(theta: float = 0.0, phi: float = 0.0, lbd: float = 0.0) -> Gate:
     r"""
     IBMQ U gate following the converntion of OpenQASM3.0.
     See `OpenQASM doc <https://openqasm.com/language/gates.html#built-in-gates>`_
@@ -533,12 +539,7 @@ def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
     :rtype: Gate
     """
     theta, phi, lbd = array_to_tensor(theta, phi, lbd)
-    i00, i01, i10, i11 = array_to_tensor(
-        np.array([[1, 0], [0, 0]]),
-        np.array([[0, 1], [0, 0]]),
-        np.array([[0, 0], [1, 0]]),
-        np.array([[0, 0], [0, 1]]),
-    )
+    i00, i01, i10, i11 = array_to_tensor(_i00, _i01, _i10, _i11)
     unitary = (
         backend.cos(theta / 2) * i00
         - backend.exp(1.0j * lbd) * backend.sin(theta / 2) * i01
@@ -548,7 +549,7 @@ def u_gate(theta: float = 0, phi: float = 0, lbd: float = 0) -> Gate:
     return Gate(unitary)
 
 
-def r_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def r_gate(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     General single qubit rotation gate
 
@@ -582,7 +583,7 @@ def r_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
 # r = r_gate
 
 
-def rx_gate(theta: float = 0) -> Gate:
+def rx_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`x` axis.
 
@@ -603,7 +604,7 @@ def rx_gate(theta: float = 0) -> Gate:
 # rx = rx_gate
 
 
-def ry_gate(theta: float = 0) -> Gate:
+def ry_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`y` axis.
 
@@ -624,7 +625,7 @@ def ry_gate(theta: float = 0) -> Gate:
 # ry = ry_gate
 
 
-def rz_gate(theta: float = 0) -> Gate:
+def rz_gate(theta: float = 0.0) -> Gate:
     r"""
     Rotation gate along :math:`z` axis.
 
@@ -645,7 +646,7 @@ def rz_gate(theta: float = 0) -> Gate:
 # rz = rz_gate
 
 
-def rgate_theoretical(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def rgate_theoretical(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     Rotation gate implemented by matrix exponential. The output is the same as `rgate`.
 
@@ -723,7 +724,7 @@ def iswap_gate(theta: float = 1.0) -> Gate:
 # iswap = iswap_gate
 
 
-def cr_gate(theta: float = 0, alpha: float = 0, phi: float = 0) -> Gate:
+def cr_gate(theta: float = 0.0, alpha: float = 0.0, phi: float = 0.0) -> Gate:
     r"""
     Controlled rotation gate. When the control qubit is 1, `rgate` is applied to the target qubit.
 
@@ -775,7 +776,7 @@ def random_two_qubit_gate() -> Gate:
     return Gate(deepcopy(unitary), name="R2Q")
 
 
-def any_gate(unitary: Tensor, name: str = "any") -> Gate:
+def any_gate(unitary: Tensor, name: str = "any", dim: Optional[int] = None) -> Gate:
     """
     Note one should provide the gate with properly reshaped.
 
@@ -783,6 +784,8 @@ def any_gate(unitary: Tensor, name: str = "any") -> Gate:
     :type unitary: Tensor
     :param name: The name of the gate.
     :type name: str
+    :param dim: The dimension of the gate.
+    :type dim: int
     :return: the resulted gate
     :rtype: Gate
     """
@@ -791,7 +794,10 @@ def any_gate(unitary: Tensor, name: str = "any") -> Gate:
         unitary.tensor = backend.cast(unitary.tensor, dtypestr)
         return unitary
     unitary = backend.cast(unitary, dtypestr)
-    unitary = backend.reshape2(unitary)
+    if dim is None or dim == 2:
+        unitary = backend.reshape2(unitary)
+    else:
+        unitary = backend.reshaped(unitary, dim)
     # nleg = int(np.log2(backend.sizen(unitary)))
     # unitary = backend.reshape(unitary, [2 for _ in range(nleg)])
     return Gate(unitary, name=name)
@@ -864,6 +870,43 @@ def exponential_gate_unity(
     return Gate(mat, name="exp1-" + name)
 
 
+def su4_gate(theta: Tensor, name: str = "su(4)") -> Gate:
+    r"""
+    Two-qubit general SU(4) gate.
+
+    :param theta: the angle tensor (15 components) of the gate.
+    :type theta: Tensor
+    :param name: the name of the gate.
+    :type name: str
+    :return: a gate object.
+    :rtype: Gate
+    """
+    theta = num_to_tensor(theta)
+    pauli_ops = array_to_tensor(
+        _ix_matrix,
+        _iy_matrix,
+        _iz_matrix,
+        _xi_matrix,
+        _xx_matrix,
+        _xy_matrix,
+        _xz_matrix,
+        _yi_matrix,
+        _yx_matrix,
+        _yy_matrix,
+        _yz_matrix,
+        _zi_matrix,
+        _zx_matrix,
+        _zy_matrix,
+        _zz_matrix,
+    )
+    generator = backend.sum(
+        backend.stack([theta[i] * pauli_ops[i] for i in range(15)]), axis=0
+    )
+    mat = backend.expm(-1j * generator)
+    mat = backend.reshape2(mat)
+    return Gate(mat, name=name)
+
+
 exp1_gate = exponential_gate_unity
 # exp1 = exponential_gate_unity
 rzz_gate = partial(exp1_gate, unitary=_zz_matrix, half=True)
@@ -968,6 +1011,7 @@ def meta_vgate() -> None:
         "rzz",
         "rxx",
         "ryy",
+        "su4",
     ]:
         for funcname in [f, f + "gate"]:
             setattr(thismodule, funcname, GateVF(getattr(thismodule, f + "_gate"), f))

tensorcircuit/interfaces/__init__.py

@@ -14,6 +14,4 @@ from .numpy import numpy_interface, np_interface
 from .scipy import scipy_interface, scipy_optimize_interface
 from .torch import torch_interface, pytorch_interface, torch_interface_kws
 from .tensorflow import tensorflow_interface, tf_interface
-
-
-# TODO(@refraction-ray): jax interface using puer_callback and custom_vjp
+from .jax import jax_interface

tensorcircuit/interfaces/jax.py

@@ -0,0 +1,189 @@
+"""
+Interface wraps quantum function as a jax function
+"""
+
+from typing import Any, Callable, Tuple, Optional, Union, Sequence
+from functools import wraps, partial
+
+from ..cons import backend
+from .tensortrans import general_args_to_backend
+
+Tensor = Any
+
+
+def jax_wrapper(
+    fun: Callable[..., Any],
+    enable_dlpack: bool = False,
+    output_shape: Optional[
+        Union[Tuple[int, ...], Tuple[int, ...], Sequence[Tuple[int, ...]]]
+    ] = None,
+    output_dtype: Optional[Union[Any, Sequence[Any]]] = None,
+) -> Callable[..., Any]:
+    import jax
+
+    @wraps(fun)
+    def fun_jax(*x: Any) -> Any:
+        def wrapped_fun(*args: Any) -> Any:
+            args = general_args_to_backend(args, enable_dlpack=enable_dlpack)
+            y = fun(*args)
+            y = general_args_to_backend(
+                y, target_backend="jax", enable_dlpack=enable_dlpack
+            )
+            return y
+
+        # Use provided shape and dtype if available, otherwise run test
+        if output_shape is not None and output_dtype is not None:
+            if isinstance(output_shape, Sequence) and not isinstance(
+                output_shape[0], int
+            ):
+                # Multiple outputs case
+                out_shape = tuple(
+                    jax.ShapeDtypeStruct(s, d)
+                    for s, d in zip(output_shape, output_dtype)
+                )
+            else:
+                # Single output case
+                out_shape = jax.ShapeDtypeStruct(output_shape, output_dtype)  # type: ignore
+        else:
+            # Get expected output shape by running function once
+            test_out = wrapped_fun(*x)
+            if isinstance(test_out, tuple):
+                # Multiple outputs case
+                out_shape = tuple(
+                    jax.ShapeDtypeStruct(
+                        t.shape if hasattr(t, "shape") else (),
+                        t.dtype if hasattr(t, "dtype") else x[0].dtype,
+                    )
+                    for t in test_out
+                )
+            else:
+                # Single output case
+                out_shape = jax.ShapeDtypeStruct(  # type: ignore
+                    test_out.shape if hasattr(test_out, "shape") else (),
+                    test_out.dtype if hasattr(test_out, "dtype") else x[0].dtype,
+                )
+
+        # Use pure_callback with correct output shape
+        result = jax.pure_callback(wrapped_fun, out_shape, *x)
+        return result
+
+    return fun_jax
+
+
+def jax_interface(
+    fun: Callable[..., Any],
+    jit: bool = False,
+    enable_dlpack: bool = False,
+    output_shape: Optional[Union[Tuple[int, ...], Tuple[()]]] = None,
+    output_dtype: Optional[Any] = None,
+) -> Callable[..., Any]:
+    """
+    Wrap a function on different ML backend with a jax interface.
+
+    :Example:
+
+    .. code-block:: python
+
+        tc.set_backend("tensorflow")
+
+        def f(params):
+            c = tc.Circuit(1)
+            c.rx(0, theta=params[0])
+            c.ry(0, theta=params[1])
+            return tc.backend.real(c.expectation([tc.gates.z(), [0]]))
+
+        f = tc.interfaces.jax_interface(f, jit=True)
+
+        params = jnp.ones(2)
+        value, grad = jax.value_and_grad(f)(params)
+
+    :param fun: The quantum function with tensor in and tensor out
+    :type fun: Callable[..., Any]
+    :param jit: whether to jit ``fun``, defaults to False
+    :type jit: bool, optional
+    :param enable_dlpack: whether transform tensor backend via dlpack, defaults to False
+    :type enable_dlpack: bool, optional
+    :param output_shape: Optional shape of the function output, defaults to None
+    :type output_shape: Optional[Union[Tuple[int, ...], Tuple[()]]], optional
+    :param output_dtype: Optional dtype of the function output, defaults to None
+    :type output_dtype: Optional[Any], optional
+    :return: The same quantum function but now with jax array in and jax array out
+        while AD is also supported
+    :rtype: Callable[..., Any]
+    """
+    jax_fun = create_jax_function(
+        fun,
+        enable_dlpack=enable_dlpack,
+        jit=jit,
+        output_shape=output_shape,
+        output_dtype=output_dtype,
+    )
+    return jax_fun
+
+
+def create_jax_function(
+    fun: Callable[..., Any],
+    enable_dlpack: bool = False,
+    jit: bool = False,
+    output_shape: Optional[Union[Tuple[int, ...], Tuple[()]]] = None,
+    output_dtype: Optional[Any] = None,
+) -> Callable[..., Any]:
+    import jax
+    from jax import custom_vjp
+
+    if jit:
+        fun = backend.jit(fun)
+
+    wrapped = jax_wrapper(
+        fun,
+        enable_dlpack=enable_dlpack,
+        output_shape=output_shape,
+        output_dtype=output_dtype,
+    )
+
+    @custom_vjp
+    def f(*x: Any) -> Any:
+        return wrapped(*x)
+
+    def f_fwd(*x: Any) -> Tuple[Any, Tuple[Any, ...]]:
+        y = wrapped(*x)
+        return y, x
+
+    def f_bwd(res: Tuple[Any, ...], g: Any) -> Tuple[Any, ...]:
+        x = res
+
+        if len(x) == 1:
+            x = x[0]
+
+        vjp_fun = partial(backend.vjp, fun)
+        if jit:
+            vjp_fun = backend.jit(vjp_fun)  # type: ignore
+
+        def vjp_wrapped(args: Any) -> Any:
+            args = general_args_to_backend(args, enable_dlpack=enable_dlpack)
+            gb = general_args_to_backend(g, enable_dlpack=enable_dlpack)
+            r = vjp_fun(args, gb)[1]
+            r = general_args_to_backend(
+                r, target_backend="jax", enable_dlpack=enable_dlpack
+            )
+            return r
+
+        # Handle gradient shape for both single input and tuple inputs
+        if isinstance(x, tuple):
+            # Create a tuple of ShapeDtypeStruct for each input
+            grad_shape = tuple(jax.ShapeDtypeStruct(xi.shape, xi.dtype) for xi in x)
+        else:
+            grad_shape = jax.ShapeDtypeStruct(x.shape, x.dtype)
+
+        dx = jax.pure_callback(
+            vjp_wrapped,
+            grad_shape,
+            x,
+        )
+
+        if not isinstance(dx, tuple):
+            dx = (dx,)
+        return dx  # type: ignore
+
+    f.defvjp(f_fwd, f_bwd)
+    return f

tensorcircuit/keras.py

@@ -24,7 +24,7 @@ class QuantumLayer(Layer):  # type: ignore
         initializer: Union[Text, Sequence[Text]] = "glorot_uniform",
         constraint: Optional[Union[Text, Sequence[Text]]] = None,
         regularizer: Optional[Union[Text, Sequence[Text]]] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> None:
         """
         `QuantumLayer` wraps the quantum function `f` as a `keras.Layer`
@@ -103,7 +103,7 @@ class QuantumLayer(Layer):  # type: ignore
         inputs: tf.Tensor,
         training: Optional[bool] = None,
         mask: Optional[tf.Tensor] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> tf.Tensor:
         # input_shape = list(inputs.shape)
         # inputs = tf.reshape(inputs, (-1, input_shape[-1]))
@@ -154,7 +154,7 @@ class HardwareLayer(QuantumLayer):
         inputs: tf.Tensor,
         training: Optional[bool] = None,
         mask: Optional[tf.Tensor] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> tf.Tensor:
         if inputs is None:  # not possible
             result = self.f(*self.pqc_weights, **kwargs)