tensorcircuit-nightly 1.3.0.dev20250729__py3-none-any.whl → 1.3.0.dev20250731__py3-none-any.whl

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0.dev20250729"
+__version__ = "1.3.0.dev20250731"
 __author__ = "TensorCircuit Authors"
 __creator__ = "refraction-ray"
 

tensorcircuit/backends/abstract_backend.py

@@ -808,6 +808,21 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `solve`.".format(self.name)
         )
 
+    def special_jv(self: Any, v: int, z: Tensor, M: int) -> Tensor:
+        """
+        Special function: Bessel function of the first kind.
+
+        :param v: The order of the Bessel function.
+        :type v: int
+        :param z: The argument of the Bessel function.
+        :type z: Tensor
+        :return: The value of the Bessel function [J_0, ...J_{v-1}(z)].
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `special_jv`.".format(self.name)
+        )
+
     def searchsorted(self: Any, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         """
         Find indices where elements should be inserted to maintain order.
@@ -1251,6 +1266,26 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `sparse_dense_matmul`.".format(self.name)
         )
 
+    def sparse_csr_from_coo(self: Any, coo: Tensor, strict: bool = False) -> Tensor:
+        """
+        transform a coo matrix to a csr matrix
+
+        :param coo: a coo matrix
+        :type coo: Tensor
+        :param strict: whether to enforce the transform, defaults to False,
+            corresponding to return the coo matrix if there is no implementation for specific backend.
+        :type strict: bool, optional
+        :return: a csr matrix
+        :rtype: Tensor
+        """
+        if strict:
+            raise NotImplementedError(
+                "Backend '{}' has not implemented `sparse_csr_from_coo`.".format(
+                    self.name
+                )
+            )
+        return coo
+
     def to_dense(self: Any, sp_a: Tensor) -> Tensor:
         """
         Convert a sparse matrix to dense tensor.
@@ -1389,10 +1424,49 @@ class ExtendedBackend:
         :rtype: Tensor
         """
         carry = init
-        for x in xs:
-            carry = f(carry, x)
+        # Check if `xs` is a PyTree (tuple or list) of arrays.
+        if isinstance(xs, (tuple, list)):
+            for x_slice_tuple in zip(*xs):
+                # x_slice_tuple will be (k_elems[i], j_elems[i]) at each step.
+                carry = f(carry, x_slice_tuple)
+        else:
+            # If xs is a single array, iterate normally.
+            for x in xs:
+                carry = f(carry, x)
+
         return carry
 
+    def jaxy_scan(
+        self: Any, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        """
+        This API follows jax scan style. TF use plain for loop
+
+        :param f: _description_
+        :type f: Callable[[Tensor, Tensor], Tensor]
+        :param init: _description_
+        :type init: Tensor
+        :param xs: _description_
+        :type xs: Tensor
+        :raises ValueError: _description_
+        :return: _description_
+        :rtype: Tensor
+        """
+        if xs is None:
+            raise ValueError("Either xs or length must be provided.")
+        if xs is not None:
+            length = len(xs)
+        carry, outputs_to_stack = init, []
+        for i in range(length):
+            if isinstance(xs, (tuple, list)):
+                x = [ele[i] for ele in xs]
+            else:
+                x = xs[i]
+            new_carry, y = f(carry, x)
+            carry = new_carry
+            outputs_to_stack.append(y)
+        return carry, self.stack(outputs_to_stack)
+
     def stop_gradient(self: Any, a: Tensor) -> Tensor:
         """
         Stop backpropagation from ``a``.

tensorcircuit/backends/jax_backend.py

@@ -418,6 +418,11 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type: ignore
         return jsp.linalg.solve(A, b, assume_a=assume_a)
 
+    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
+        from .jax_ops import bessel_jv_jax_rescaled
+
+        return bessel_jv_jax_rescaled(v, z, M)
+
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         if not self.is_tensor(a):
             a = self.convert_to_tensor(a)
@@ -615,6 +620,11 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
         carry, _ = libjax.lax.scan(f_jax, init, xs)
         return carry
 
+    def jaxy_scan(
+        self, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        return libjax.lax.scan(f, init, xs)
+
     def scatter(self, operand: Tensor, indices: Tensor, updates: Tensor) -> Tensor:
         # updates = jnp.reshape(updates, indices.shape)
         # return operand.at[indices].set(updates)
@@ -639,11 +649,20 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        try:
+            return sparse.BCSR.from_bcoo(coo)  # type: ignore
+        except AttributeError as e:
+            if not strict:
+                return coo
+            else:
+                raise e
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 
     def is_sparse(self, a: Tensor) -> bool:
-        return isinstance(a, sparse.BCOO)  # type: ignore
+        return isinstance(a, sparse.JAXSparse)  # type: ignore
 
     def device(self, a: Tensor) -> str:
         (dev,) = a.devices()

tensorcircuit/backends/jax_ops.py

@@ -3,8 +3,11 @@ Customized ops for ML framework
 """
 
 # pylint: disable=invalid-name
+# pylint: disable=unused-variable
+
 
 from typing import Any, Tuple, Sequence
+from functools import partial
 
 import jax
 import jax.numpy as jnp
@@ -174,3 +177,108 @@ def jaxeigh_bwd(r: Array, tangents: Array) -> Array:
 
 adaware_eigh.defvjp(jaxeigh_fwd, jaxeigh_bwd)
 adaware_eigh_jit = jax.jit(adaware_eigh)
+
+
+@partial(jax.jit, static_argnums=[0, 2])
+def bessel_jv_jax_rescaled(k: int, x: jnp.ndarray, M: int) -> jnp.ndarray:
+    """
+    Computes Bessel function Jv using Miller's algorithm with dynamic rescaling,
+    implemented in JAX.
+    """
+    if M <= k:
+        raise ValueError(
+            f"Recurrence length M ({M}) must be greater than the required order k ({k})."
+        )
+
+    # Use vmap to handle array inputs for x efficiently.
+    # We map _bessel_jv_scalar_rescaled over the last dimension of x.
+    return _bessel_jv_scalar_rescaled(k, M, x)
+
+
+def _bessel_jv_scalar_rescaled(k: int, M: int, x_val: jnp.ndarray) -> jnp.ndarray:
+    """
+    JAX implementation for a scalar input x_val.
+    This function will be vmapped for array inputs.
+    """
+    rescale_threshold = 1e250
+
+    # Define the body of the recurrence loop
+    def recurrence_body(i, state):  # type: ignore
+        # M - i is the current 'm' value in the original loop.
+        # Loop from M down to 1. jax.lax.fori_loop goes from lower to upper-1.
+        # So for m from M down to 1, we map i from 0 to M-1.
+        # Current m_val = M - i
+        # The loop range for m in numpy was `range(M, 0, -1)`, which means m goes from M, M-1, ..., 1.
+        # For lax.fori_loop (start, stop, body_fn, init_val), start is inclusive, stop is exclusive.
+        # So to iterate M times for m from M down to 1, we do i from 0 to M-1.
+        # m_val = M - i means that for i=0, m_val=M; for i=M-1, m_val=1.
+        m_val = M - i
+        f_m, f_m_p1, f_vals = state
+
+        # If x_val is near zero, this division could be an issue,
+        # but the outer lax.cond handles the x_val near zero case before this loop runs.
+        f_m_m1 = (2.0 * m_val / x_val) * f_m - f_m_p1
+
+        # --- Rescaling Step ---
+        # jax.lax.cond requires all branches to return the exact same type and shape.
+        def rescale_branch(vals):  # type: ignore
+            f_m_val, f_m_p1_val, f_vals_arr = vals
+            scale_factor = f_m_m1
+            # Return new f_m, f_m_p1, updated f_vals_arr, and the new f_m_m1 value (which is 1.0)
+            return (
+                f_m_val / scale_factor,
+                f_m_p1_val / scale_factor,
+                f_vals_arr / scale_factor,
+                1.0,
+            )
+
+        def no_rescale_branch(vals):  # type: ignore
+            f_m_val, f_m_p1_val, f_vals_arr = (
+                vals  # Unpack to keep signatures consistent
+            )
+            # Return original f_m, f_m_p1, original f_vals_arr, and the computed f_m_m1
+            return (f_m_val, f_m_p1_val, f_vals_arr, f_m_m1)
+
+        f_m_rescaled, f_m_p1_rescaled, f_vals_rescaled, f_m_m1_effective = jax.lax.cond(
+            jnp.abs(f_m_m1) > rescale_threshold,
+            rescale_branch,
+            no_rescale_branch,
+            (f_m, f_m_p1, f_vals),  # Arguments passed to branches
+        )
+
+        # Update f_vals at index m_val - 1. JAX uses .at[idx].set(val) for non-in-place updates.
+        f_vals_updated = f_vals_rescaled.at[m_val - 1].set(f_m_m1_effective)
+
+        # Return new state for the next iteration: (new f_m, new f_m_p1, updated f_vals)
+        return (f_m_m1_effective, f_m_rescaled, f_vals_updated)
+
+    # Initial state for the recurrence loop
+    f_m_p1_init = 0.0
+    f_m_init = 1e-30  # Start with a very small number
+    f_vals_init = jnp.zeros(M + 1).at[M].set(f_m_init)
+
+    # Use jax.lax.fori_loop for the backward recurrence
+    # Loop from i = 0 to M-1 (total M iterations)
+    # The 'body' function gets current 'i' and 'state', returns 'new_state'.
+    # We don't need the final f_m_p1, only f_m and f_vals.
+    final_f_m, _, f_vals = jax.lax.fori_loop(
+        0, M, recurrence_body, (f_m_init, f_m_p1_init, f_vals_init)
+    )
+
+    # Normalization using Neumann's sum rule
+    even_sum = jnp.sum(f_vals[2::2])
+    norm_const = f_vals[0] + 2.0 * even_sum
+
+    # Handle division by near-zero normalization constant
+    norm_const_safe = jnp.where(jnp.abs(norm_const) < 1e-12, 1e-12, norm_const)
+
+    # Conditional logic for x_val close to zero
+    def x_is_zero_case() -> jnp.ndarray:
+        # For x=0, J_0(0)=1, J_k(0)=0 for k>0
+        return jnp.zeros(k).at[0].set(1.0)
+
+    def x_is_not_zero_case() -> jnp.ndarray:
+        return f_vals[:k] / norm_const_safe  # type: ignore
+
+    # Use lax.cond to select between the two cases based on x_val
+    return jax.lax.cond(jnp.abs(x_val) < 1e-12, x_is_zero_case, x_is_not_zero_case)  # type: ignore

tensorcircuit/backends/numpy_backend.py

@@ -17,7 +17,7 @@ except ImportError:  # np2.0 compatibility
 
 import tensornetwork
 from scipy.linalg import expm, solve, schur
-from scipy.special import softmax, expit
+from scipy.special import softmax, expit, jv
 from scipy.sparse import coo_matrix, issparse
 from tensornetwork.backends.numpy import numpy_backend
 from .abstract_backend import ExtendedBackend
@@ -200,6 +200,7 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
         return softmax(a, axis=axis)
 
     def onehot(self, a: Tensor, num: int) -> Tensor:
+        a = np.asarray(a)
         res = np.eye(num)[a.reshape([-1])]
         return res.reshape(list(a.shape) + [num])
         # https://stackoverflow.com/questions/38592324/one-hot-encoding-using-numpy
@@ -244,6 +245,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
         # https://stackoverflow.com/questions/44672029/difference-between-numpy-linalg-solve-and-numpy-linalg-lu-solve/44710451
         return solve(A, b, assume_a=assume_a)
 
+    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
+        return jv(np.arange(v), z)
+
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         return np.searchsorted(a, v, side=side)  # type: ignore
 
@@ -329,6 +333,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        return coo.tocsr()
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 

tensorcircuit/backends/tensorflow_backend.py

@@ -360,6 +360,38 @@ tensornetwork.backends.tensorflow.tensorflow_backend.TensorFlowBackend.rq = _rq_
 tensornetwork.backends.tensorflow.tensorflow_backend.TensorFlowBackend.svd = _svd_tf
 
 
+def sparse_tensor_matmul(self: Tensor, other: Tensor) -> Tensor:
+    """
+    An implementation of matrix multiplication (@) for tf.SparseTensor.
+
+    This function is designed to be monkey-patched onto the tf.SparseTensor class.
+    It handles multiplication with a dense vector (rank-1 Tensor) by temporarily
+    promoting it to a matrix (rank-2 Tensor) for the underlying TensorFlow call.
+    """
+    # Ensure the 'other' tensor is of a compatible dtype
+    if not other.dtype.is_compatible_with(self.dtype):
+        other = tf.cast(other, self.dtype)
+
+    # tf.sparse.sparse_dense_matmul requires the dense tensor to be a 2D matrix.
+    # If we get a 1D vector, we need to reshape it.
+    is_vector = len(other.shape) == 1
+
+    if is_vector:
+        # Promote the vector to a column matrix [N] -> [N, 1]
+        other_matrix = tf.expand_dims(other, axis=1)
+    else:
+        other_matrix = other
+
+    # Perform the actual multiplication
+    result_matrix = tf.sparse.sparse_dense_matmul(self, other_matrix)
+
+    if is_vector:
+        # Demote the result matrix back to a vector [M, 1] -> [M]
+        return tf.squeeze(result_matrix, axis=1)
+    else:
+        return result_matrix
+
+
 class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):  # type: ignore
     """
     See the original backend API at `tensorflow backend
@@ -378,6 +410,8 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
             )
         tf = tensorflow
         tf.sparse.SparseTensor.__add__ = tf.sparse.add
+        tf.SparseTensor.__matmul__ = sparse_tensor_matmul
+
         self.minor = int(tf.__version__.split(".")[1])
         self.name = "tensorflow"
         logger = tf.get_logger()  # .setLevel('ERROR')
@@ -719,7 +753,10 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def scan(
         self, f: Callable[[Tensor, Tensor], Tensor], xs: Tensor, init: Tensor
     ) -> Tensor:
-        return tf.scan(f, xs, init)[-1]
+        stacked_results = tf.scan(f, xs, init)
+        final_state = tf.nest.map_structure(lambda x: x[-1], stacked_results)
+        return final_state
+        # return tf.scan(f, xs, init)[-1]
 
     def device(self, a: Tensor) -> str:
         dev = a.device

tensorcircuit/fgs.py

@@ -28,7 +28,6 @@ def onehot_matrix(i: int, j: int, N: int) -> Tensor:
 
 # TODO(@refraction-ray): efficiency benchmark with jit
 # TODO(@refraction-ray): FGS mixed state support?
-# TODO(@refraction-ray): overlap?
 # TODO(@refraction-ray): fermionic logarithmic negativity
 
 
@@ -227,7 +226,7 @@ class FGSSimulator:
         return self.alpha
 
     def get_cmatrix(self, now_i: bool = True, now_j: bool = True) -> Tensor:
-        """
+        r"""
         Calculates the correlation matrix.
 
         The correlation matrix is defined as :math:`C_{ij} = \langle c_i^\dagger c_j \rangle`.
@@ -509,7 +508,7 @@ class FGSSimulator:
 
     @staticmethod
     def hopping(chi: Tensor, i: int, j: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the hopping Hamiltonian between two sites.
 
         The hopping Hamiltonian is given by :math:`\chi c_i^\dagger c_j + h.c.`.
@@ -550,7 +549,7 @@ class FGSSimulator:
 
     @staticmethod
     def chemical_potential(chi: Tensor, i: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the chemical potential Hamiltonian for a single site.
 
         The chemical potential Hamiltonian is given by :math:`\chi c_i^\dagger c_i`.
@@ -572,7 +571,7 @@ class FGSSimulator:
 
     @staticmethod
     def sc_pairing(chi: Tensor, i: int, j: int, L: int) -> Tensor:
-        """
+        r"""
         Constructs the superconducting pairing Hamiltonian between two sites.
 
         The superconducting pairing Hamiltonian is given by :math:`\chi c_i^\dagger c_j^\dagger + h.c.`.
@@ -637,7 +636,7 @@ class FGSSimulator:
         self.evol_ihamiltonian(self.chemical_potential(chi, i, self.L))
 
     def get_bogoliubov_uv(self) -> Tuple[Tensor, Tensor]:
-        """
+        r"""
         Returns the u and v matrices of the Bogoliubov transformation.
 
         The Bogoliubov transformation is defined as:

tensorcircuit/quantum.py

@@ -1433,7 +1433,7 @@ def PauliStringSum2Dense(
         return sparsem.todense()
     sparsem = backend.coo_sparse_matrix_from_numpy(sparsem)
     densem = backend.to_dense(sparsem)
-    return densem
+    return backend.convert_to_tensor(densem)
 
 
 # already implemented as backend method

tensorcircuit/stabilizercircuit.py

@@ -96,10 +96,12 @@ class StabilizerCircuit(AbstractCircuit):
 
         if name.lower() in self.gate_map:
             # self._stim_circuit.append(gate_map[name.lower()], list(index))
-            instruction = f"{self.gate_map[name.lower()]} {' '.join(map(str, index))}"
+            gn = self.gate_map[name.lower()]
+            instruction = f"{gn} {' '.join(map(str, index))}"
             self._stim_circuit.append_from_stim_program_text(instruction)
             # append is much slower
-            self.current_sim.do(stim.Circuit(instruction))
+            # self.current_sim.do(stim.Circuit(instruction))
+            getattr(self.current_sim, gn.lower())(*index)
         else:
             raise ValueError(f"Gate {name} is not supported in stabilizer simulation")