tensorcircuit-nightly 1.3.0.dev20250730__py3-none-any.whl → 1.3.0.dev20250801__py3-none-any.whl

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0.dev20250730"
+__version__ = "1.3.0.dev20250801"
 __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.
@@ -1400,10 +1435,37 @@ class ExtendedBackend:
                 carry = f(carry, x)
 
         return carry
-        # carry = init
-        # 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:
         """

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
@@ -245,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
 
@@ -330,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')

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")
 

tensorcircuit/templates/lattice.py

@@ -15,6 +15,7 @@ from typing import (
     Union,
     TYPE_CHECKING,
     cast,
+    Set,
 )
 
 logger = logging.getLogger(__name__)
@@ -1446,3 +1447,54 @@ class CustomizeLattice(AbstractLattice):
         logger.info(
             f"{len(ids_to_remove)} sites removed. Lattice now has {self.num_sites} sites."
         )
+
+
+def get_compatible_layers(bonds: List[Tuple[int, int]]) -> List[List[Tuple[int, int]]]:
+    """
+    Partitions a list of pairs (bonds) into compatible layers for parallel
+    gate application using a greedy edge-coloring algorithm.
+
+    This function takes a list of pairs, representing connections like
+    nearest-neighbor (NN) or next-nearest-neighbor (NNN) bonds, and
+    partitions them into the minimum number of sets ("layers") where no two
+    pairs in a set share an index. This is a general utility for scheduling
+    non-overlapping operations.
+
+    :Example:
+
+    >>> from tensorcircuit.templates.lattice import SquareLattice
+    >>> sq_lattice = SquareLattice(size=(2, 2), pbc=False)
+    >>> nn_bonds = sq_lattice.get_neighbor_pairs(k=1, unique=True)
+
+    >>> gate_layers = get_compatible_layers(nn_bonds)
+    >>> print(gate_layers)
+    [[[0, 1], [2, 3]], [[0, 2], [1, 3]]]
+
+    :param bonds: A list of tuples, where each tuple represents a bond (i, j)
+        of site indices to be scheduled.
+    :type bonds: List[Tuple[int, int]]
+    :return: A list of layers. Each layer is a list of tuples, where each
+        tuple represents a bond. All bonds within a layer are non-overlapping.
+    :rtype: List[List[Tuple[int, int]]]
+    """
+    uncolored_edges: Set[Tuple[int, int]] = {(min(bond), max(bond)) for bond in bonds}
+
+    layers: List[List[Tuple[int, int]]] = []
+
+    while uncolored_edges:
+        current_layer: List[Tuple[int, int]] = []
+        qubits_in_this_layer: Set[int] = set()
+
+        edges_to_process = sorted(list(uncolored_edges))
+
+        for edge in edges_to_process:
+            i, j = edge
+            if i not in qubits_in_this_layer and j not in qubits_in_this_layer:
+                current_layer.append(edge)
+                qubits_in_this_layer.add(i)
+                qubits_in_this_layer.add(j)
+
+        uncolored_edges -= set(current_layer)
+        layers.append(sorted(current_layer))
+
+    return layers

tensorcircuit/timeevol.py

@@ -32,6 +32,8 @@ def lanczos_iteration_scan(
     :rtype: Tuple[Tensor, Tensor]
     """
     state_size = backend.shape_tuple(initial_vector)[0]
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
 
     # Main scan body for the outer loop (iterating j)
     def lanczos_step(carry: Tuple[Any, ...], j: int) -> Tuple[Any, ...]:
@@ -196,6 +198,9 @@ def lanczos_iteration(
     # Add first basis vector
     basis_vectors.append(vector)
 
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
+
     # Lanczos iteration (fixed number of iterations for JIT compatibility)
     for j in range(subspace_dimension):
         # Calculate H|v_j>
@@ -592,3 +597,215 @@ def ode_evol_global(
     if not callback:
         return s1
     return backend.stack([callback(s1[i]) for i in range(len(s1))])
+
+
+def chebyshev_evol(
+    hamiltonian: Any,
+    initial_state: Tensor,
+    t: float,
+    spectral_bounds: Tuple[float, float],
+    k: int,
+    M: int,
+) -> Any:
+    """
+    Chebyshev evolution method by expanding the time evolution exponential operator
+    in Chebyshev series.
+    Note the state returned is not normalized. But the norm should be very close to 1 for
+    sufficiently large k and M, which can serve as a accuracy check of the final result.
+
+    :param hamiltonian: Hamiltonian matrix (sparse or dense)
+    :type hamiltonian: Any
+    :param initial_state: Initial state vector
+    :type initial_state: Tensor
+    :param time: Time to evolve
+    :type time: float
+    :param spectral_bounds: Spectral bounds for the Hamiltonian (Emax, Emin)
+    :type spectral_bounds: Tuple[float, float]
+    :param k: Number of Chebyshev coefficients, a good estimate is k > t*(Emax-Emin)/2
+    :type k: int
+    :param M: Number of iterations to estimate Bessel function, a good estimate is given
+        by `estimate_M` helper method.
+    :type M: int
+    :return: Evolved state
+    :rtype: Tensor
+    """
+    # TODO(@refraction-ray): no support for tf backend as bessel function has no implementation
+    E_max, E_min = spectral_bounds
+    if E_max <= E_min:
+        raise ValueError("E_max must be > E_min.")
+
+    a = (E_max - E_min) / 2.0
+    b = (E_max + E_min) / 2.0
+    tau = a * t  # Rescaled time parameter
+
+    if backend.is_sparse(hamiltonian):
+        hamiltonian = backend.sparse_csr_from_coo(hamiltonian)
+
+    def apply_h_norm(psi: Any) -> Any:
+        """Applies the normalized Hamiltonian to a state."""
+        return ((hamiltonian @ psi) - b * psi) / a
+
+    # Handle edge case where no evolution is needed.
+    if k == 0:
+        # The phase factor still applies even for zero evolution of the series part.
+        phase = backend.exp(-1j * b * t)
+        return phase * backend.zeros_like(initial_state)
+
+    # --- 2. Calculate Chebyshev Expansion Coefficients ---
+    k_indices = backend.arange(k)
+    bessel_vals = backend.special_jv(k, tau, M)
+
+    # Prefactor is 1 for k=0 and 2 for k>0.
+    prefactor = backend.ones([k])
+    if k > 1:
+        # Using concat for backend compatibility (vs. jax's .at[1:].set(2.0))
+        prefactor = backend.concat(
+            [backend.ones([1]), backend.ones([k - 1]) * 2.0], axis=0
+        )
+
+    ik_powers = backend.power(0 - 1j, k_indices)
+    coeffs = prefactor * ik_powers * bessel_vals
+
+    # --- 3. Iteratively build the result using a scan ---
+
+    # Handle the simple case of k=1 separately.
+    if k == 1:
+        psi_unphased = coeffs[0] * initial_state
+    else:  # k >= 2, use the scan operation.
+        # Initialize the first two Chebyshev vectors and the initial sum.
+        T0 = initial_state
+        T1 = apply_h_norm(T0)
+        initial_sum = coeffs[0] * T0 + coeffs[1] * T1
+
+        # The carry for the scan holds the state needed for the next iteration:
+        # (current vector T_k, previous vector T_{k-1}, and the running sum).
+        initial_carry = (T1, T0, initial_sum)
+
+        def scan_body(carry, i):  # type: ignore
+            """The body of the scan operation."""
+            Tk, Tkm1, current_sum = carry
+
+            # Calculate the next Chebyshev vector using the recurrence relation.
+            Tkp1 = 2 * apply_h_norm(Tk) - Tkm1
+
+            # Add its contribution to the running sum.
+            new_sum = current_sum + coeffs[i] * Tkp1
+
+            # Return the updated carry for the next step. No intermediate output is needed.
+            return (Tkp1, Tk, new_sum)
+
+        # Run the scan over the remaining coefficients (from index 2 to k-1).
+        final_carry = backend.scan(scan_body, backend.arange(2, k), initial_carry)
+
+        # The final result is the sum accumulated in the last carry state.
+        psi_unphased = final_carry[2]
+
+    # --- 4. Final Step: Apply Phase Correction ---
+    # This undoes the energy shift from the Hamiltonian normalization.
+    phase = backend.exp(-1j * b * t)
+    psi_final = phase * psi_unphased
+
+    return psi_final
+
+
+def estimate_k(t: float, spectral_bounds: Tuple[float, float]) -> int:
+    """
+    estimate k for chebyshev expansion
+
+    :param t: time
+    :type t: float
+    :param spectral_bounds: spectral bounds (Emax, Emin)
+    :type spectral_bounds: Tuple[float, float]
+    :return: k
+    :rtype: int
+    """
+    E_max, E_min = spectral_bounds
+    a = (E_max - E_min) / 2.0
+    tau = a * t  # tau is now a scalar
+    return max(int(1.1 * tau), int(tau + 20))
+
+
+def estimate_M(t: float, spectral_bounds: Tuple[float, float], k: int) -> int:
+    """
+    estimate M for Bessel function iterations
+
+    :param t: time
+    :type t: float
+    :param spectral_bounds: spectral bounds (Emax, Emin)
+    :type spectral_bounds: Tuple[float, float]
+    :param k: k
+    :type k: int
+    :return: M
+    :rtype: int
+    """
+    E_max, E_min = spectral_bounds
+    a = (E_max - E_min) / 2.0
+    tau = a * t  # tau is now a scalar
+    safety_factor = 15
+    M = max(k, int(abs(tau))) + int(safety_factor * np.sqrt(abs(tau)))
+    M = max(M, k + 30)
+    return M
+
+
+def estimate_spectral_bounds(
+    h: Any, n_iter: int = 30, psi0: Optional[Any] = None
+) -> Tuple[float, float]:
+    """
+    Lanczos algorithm to estimate the spectral bounds of a Hamiltonian.
+    Just for quick run before `chebyshev_evol`, non jit-able.
+
+    :param h: Hamiltonian matrix.
+    :type h: Any
+    :param n_iter: iteration number.
+    :type n_iter: int
+    :param psi0: Optional initial state.
+    :type psi0: Optional[Any]
+    :return: (E_max, E_min)。
+    """
+    shape = h.shape
+    D = shape[-1]
+    if psi0 is None:
+        psi0 = np.random.normal(size=[D])
+
+    psi0 = backend.convert_to_tensor(psi0) / backend.norm(psi0)
+    psi0 = backend.cast(psi0, dtypestr)
+
+    # Lanczos
+    alphas = []
+    betas = []
+    q_prev = backend.zeros(psi0.shape, dtype=psi0.dtype)
+    q = psi0
+    beta = 0
+
+    for _ in range(n_iter):
+        r = h @ q
+        r = backend.convert_to_tensor(r)  # in case np.matrix
+        r = backend.reshape(r, [-1])
+        if beta != 0:
+            r -= backend.cast(beta, dtypestr) * q_prev
+
+        alpha = backend.real(backend.sum(backend.conj(q) * r))
+
+        alphas.append(alpha)
+
+        r -= backend.cast(alpha, dtypestr) * q
+
+        q_prev = q
+        beta = backend.norm(r)
+        q = r / beta
+        beta = backend.abs(beta)
+        betas.append(beta)
+        if beta < 1e-8:
+            break
+
+    alphas = backend.stack(alphas)
+    betas = backend.stack(betas)
+    T = (
+        backend.diagflat(alphas)
+        + backend.diagflat(betas[:-1], k=1)
+        + backend.diagflat(betas[:-1], k=-1)
+    )
+
+    ritz_values, _ = backend.eigh(T)
+
+    return backend.max(ritz_values), backend.min(ritz_values)

{tensorcircuit_nightly-1.3.0.dev20250730.dist-info → tensorcircuit_nightly-1.3.0.dev20250801.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorcircuit-nightly
-Version: 1.3.0.dev20250730
+Version: 1.3.0.dev20250801
 Summary: nightly release for tensorcircuit
 Home-page: https://github.com/refraction-ray/tensorcircuit-dev
 Author: TensorCircuit Authors
@@ -70,7 +70,7 @@ TensorCircuit-NG is the actively maintained official version and a [fully compat
 
 Please begin with [Quick Start](/docs/source/quickstart.rst) in the [full documentation](https://tensorcircuit-ng.readthedocs.io/).
 
-For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 80+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to AI-native docs for tensorcircuit-ng: [Devin Deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) and [Context7 MCP](https://context7.com/tensorcircuit/tensorcircuit-ng).
+For more information on software usage, sota algorithm implementation and engineer paradigm demonstration, please refer to 90+ [example scripts](/examples) and 30+ [tutorial notebooks](https://tensorcircuit-ng.readthedocs.io/en/latest/#tutorials). API docstrings and test cases in [tests](/tests) are also informative. One can also refer to AI-native docs for tensorcircuit-ng: [Devin Deepwiki](https://deepwiki.com/tensorcircuit/tensorcircuit-ng) and [Context7 MCP](https://context7.com/tensorcircuit/tensorcircuit-ng).
 
 For beginners, please refer to [quantum computing lectures with TC-NG](https://github.com/sxzgroup/qc_lecture) to learn both quantum computing basics and representative usage of TensorCircuit-NG.
 

{tensorcircuit_nightly-1.3.0.dev20250730.dist-info → tensorcircuit_nightly-1.3.0.dev20250801.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tensorcircuit/__init__.py,sha256=DrH3d7oEWpI2MmrSmabmwoovMMYBNoLZeCk6SFC360M,2055
+tensorcircuit/__init__.py,sha256=IL3iNxbN-D630LW977RNmxFHWyPUPXy8CxLIq0cZ8f0,2055
 tensorcircuit/about.py,sha256=DazTswU2nAwOmASTaDII3L04PVtaQ7oiWPty5YMI3Wk,5267
 tensorcircuit/abstractcircuit.py,sha256=0osacPqq7B1EJki-cI1aLYoVRmjFaG9q3XevWMs7SsA,44125
 tensorcircuit/asciiart.py,sha256=neY1OWFwtoW5cHPNwkQHgRPktDniQvdlP9QKHkk52fM,8236
@@ -8,17 +8,17 @@ tensorcircuit/circuit.py,sha256=mE4b_9xRu3ydoB8iDffdx35V9GZLhAQD_tkjZDLnLjg,3910
 tensorcircuit/cons.py,sha256=uYKBeYKkDoJEqJTNrOZPRM31tBtkqe5aAg8GtVidJ1Y,33014
 tensorcircuit/densitymatrix.py,sha256=VqMBnWCxO5-OsOp6LOdc5RS2AzmB3U4-w40Vn_lqygo,14865
 tensorcircuit/experimental.py,sha256=TGK4FaS6TS_ZhtjcIZgYVuAkGdRW50LN0DdXp-h4bos,29906
-tensorcircuit/fgs.py,sha256=pzaZuzPIFPpfr5Z-UsBQ_Yp0x7mbSM2sUc4dO2SUmVs,49543
+tensorcircuit/fgs.py,sha256=J1TjAiiqZk9KO1xYX_V0xsgKlYZaUQ7Enm4s5zkRM50,49514
 tensorcircuit/gates.py,sha256=x-wA7adVpP7o0AQLt_xYUScFKj8tU_wUOV2mR1GyrPc,29322
 tensorcircuit/keras.py,sha256=5OF4dfhEeS8sRYglpqYtQsWPeqp7uK0i7-P-6RRJ7zQ,10126
 tensorcircuit/mps_base.py,sha256=UZ-v8vsr_rAsKrfun8prVgbXJ-qsdqKy2DZIHpq3sxo,15400
 tensorcircuit/mpscircuit.py,sha256=COO9xzvA2Whe7Ncp6OqrgtXKmahHgTHxXTELAVHzFSY,36777
 tensorcircuit/noisemodel.py,sha256=vzxpoYEZbHVC4a6g7_Jk4dxsHi4wvhpRFwud8b616Qo,11878
-tensorcircuit/quantum.py,sha256=LNkIv5cJ2KG6puC18zTuXi-5cojW1Tnz-N-WjZ0Qu5Q,90217
+tensorcircuit/quantum.py,sha256=1fZJJJ_o3NF7LaPPPU5Abd82NmMSnNQ8uUp_ClZdp3c,90244
 tensorcircuit/shadows.py,sha256=6XmWNubbuaxFNvZVWu-RXd0lN9Jkk-xwong_K8o8_KE,17014
 tensorcircuit/simplify.py,sha256=O11G3UYiVAc30GOfwXXmhLXwGZrQ8OVwLTMQMZp_XBc,9414
-tensorcircuit/stabilizercircuit.py,sha256=yNqcEKtYzRYrgqGil8QEyKN4OEMp9g6uOG2zuRaU8uc,15465
-tensorcircuit/timeevol.py,sha256=8p4C3nhUQ9eC2wYfZ9w5BGFIt25NPEJTXmZU_iZy4tM,21607
+tensorcircuit/stabilizercircuit.py,sha256=KbrBVSo2pXnf5JHIrxwRPSPTm7bJVMIcyE4d7-dIfCM,15545
+tensorcircuit/timeevol.py,sha256=GDD6IgzkRhZIyvLkSKJ7ZihnkHQBqoWOBsmBwRAC1UI,28574
 tensorcircuit/torchnn.py,sha256=z_QpM0QC3mydGyWpyp877j-tSFCPyzynCwqrTWaw-IA,4637
 tensorcircuit/translation.py,sha256=VnU7DnYmbk1cWjqa7N68WNLNDn3DwENrMzmbG4_CQco,28611
 tensorcircuit/utils.py,sha256=nEDR1wTh1WF_yV6UyZYlifqOPWdKk_Krr4HjhrWHnGQ,7228
@@ -40,15 +40,15 @@ tensorcircuit/applications/physics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQe
 tensorcircuit/applications/physics/baseline.py,sha256=RWrzMGnC0PtmpYSFkvCE7r1llR88gncXuCakAAhFE-w,1775
 tensorcircuit/applications/physics/fss.py,sha256=ny3U9ZDmT459PXjA1oUGfarBOlSKSy6fs04vD9s1XH4,3633
 tensorcircuit/backends/__init__.py,sha256=WiUmbUFzM29w3hKfhuKxVUk3PpqDFiXf4za9g0ctpZA,80
-tensorcircuit/backends/abstract_backend.py,sha256=pmuf8NFmtRXnlBTsRrv2s6weCNiq3cd9gA7wiKvIobM,59589
+tensorcircuit/backends/abstract_backend.py,sha256=fAzyKvZ-1Mw2wvbyuR5IScVGhwZjY7AxBVtRo6viVY0,61743
 tensorcircuit/backends/backend_factory.py,sha256=Z0aQ-RnxOnQzp-SRw8sefAH8XyBSlj2NXZwOlHinbfY,1713
 tensorcircuit/backends/cupy_backend.py,sha256=4vgO3lnQnsvWL5hukhskjJp37EAHqio6z6TVXTQcdjs,15077
-tensorcircuit/backends/jax_backend.py,sha256=dkDQ380CJHIdlt1fZvlN_g8DIowWPEcTTV_XBcs0YB0,26088
-tensorcircuit/backends/jax_ops.py,sha256=o7tLlQMRnaKWcr5rVnOMqwG6KZVpR8M8ryNQ-ceXVxs,4789
-tensorcircuit/backends/numpy_backend.py,sha256=PhbpXeATQ6X4kZ3xA-RvQVO_dKrpFh5vyXcCKKDMC7U,14197
+tensorcircuit/backends/jax_backend.py,sha256=GIWsJwhlz0PD_KcypRkNWcjkegdNdoFMsmo0u0RQqrk,26704
+tensorcircuit/backends/jax_ops.py,sha256=WyUGavch2R9uEFsI1Ap7eP1UcU4s2TItBgGsrVS3Hzs,9320
+tensorcircuit/backends/numpy_backend.py,sha256=hhjrm0GK0d05TBYHNCZrBBBIJQ7V6qN99m1KLl0WKro,14408
 tensorcircuit/backends/pytorch_backend.py,sha256=yhfZSrm99yNW-dmijk8t6zAkbVgLRd4b_aIWKrpT7bY,24230
 tensorcircuit/backends/pytorch_ops.py,sha256=lLxpK6OqfpVwifyFlgsqhpnt-oIn4R5paPMVg51WaW0,3826
-tensorcircuit/backends/tensorflow_backend.py,sha256=T2BmFxOyl2QU4dSIwUMCLyPspLrNFLx3hVfD7TD1No0,36598
+tensorcircuit/backends/tensorflow_backend.py,sha256=CQhdWjUoqCNHxmJgfYOhUVpwDqjSo0RYrmsiaVWmxCU,37842
 tensorcircuit/backends/tf_ops.py,sha256=FJwDU7LhZrt0VUIx12DJU0gZnWhMv7B7r9sAKG710As,3378
 tensorcircuit/cloud/__init__.py,sha256=n0Lx07GYF6YbdIa6AJCLJk4zlAm5CqaeHszvkxxuoI4,139
 tensorcircuit/cloud/abstraction.py,sha256=6aSxbz0MP21jBVdFbSMrvJPLQH117vGz9sSHbMFoodE,14582
@@ -84,12 +84,12 @@ tensorcircuit/templates/conversions.py,sha256=D3chiKDr7G1ekCJngiol91k9iqrMag1DZQ
 tensorcircuit/templates/dataset.py,sha256=ldPvCUlwjHU_S98E2ISQp34KqJzJPpPHmDIKJ4K-qYo,1933
 tensorcircuit/templates/graphs.py,sha256=cPYrxjoem0xZ-Is9dZKAvEzWZL_FejfIRiCEOTA4qd4,3935
 tensorcircuit/templates/hamiltonians.py,sha256=Ag8djD6lckTeU7I99gCbXiQAb2VYqzm_p7-hpXo-5u4,5554
-tensorcircuit/templates/lattice.py,sha256=F35ebANk0DSmSHLR0-Q_hUbcznyCmZjb4fKmvCMywmA,58575
+tensorcircuit/templates/lattice.py,sha256=P64OGUedE3o8vWekhM8XAs5nUe5CdG-gojLlTGA20TI,60534
 tensorcircuit/templates/measurements.py,sha256=pzc5Aa9S416Ilg4aOY77Z6ZhUlYcXnAkQNQFTuHjFFs,10943
-tensorcircuit_nightly-1.3.0.dev20250730.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+tensorcircuit_nightly-1.3.0.dev20250801.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/conftest.py,sha256=J9nHlLE3Zspz1rMyzadEuBWhaS5I4Q9sq0lnWybcdIA,1457
-tests/test_backends.py,sha256=rClxb2gyAoGeXd_ZYVSAJ0zEvJ7z_2btAeFM_Iy_wwY,33925
+tests/test_backends.py,sha256=MLxpRWavWF0qWcjVI61lIa3CYGbztvTO8ITFOYX47ao,38312
 tests/test_calibrating.py,sha256=D1Tlv8mucUhg3ULvB5QlYyaDfw7aEERwq69-aGSb1A4,3805
 tests/test_channels.py,sha256=BL4CirU8ku9-_NrI6PZAS5xZ0wrL1UEC1S3wPI9dYQM,12628
 tests/test_circuit.py,sha256=IsSIFEs7hUCSYexMb-ESt1ZUpztHtLA0qz0CZolGdc4,52240
@@ -102,7 +102,7 @@ tests/test_gates.py,sha256=rAIV2QFpFsA5bT1QivTSkhdarvwu5t0N3IOz4SEDrzg,4593
 tests/test_hamiltonians.py,sha256=E0E5ABhUeG7XLMLRkb3AIAPi7aJgnIeMWTgqzF1Q6yc,5724
 tests/test_interfaces.py,sha256=iJPmes8S8HkA9_PGjsu4Ike-vCXYyS1EMgnNKKXDNaU,16938
 tests/test_keras.py,sha256=U453jukavmx0RMeTSDEgPzrNdHNEfK1CW0CqO3XCNKo,4841
-tests/test_lattice.py,sha256=_ptDVK3EhS-X5fCQWiP8sHk3azdyGFuwqg6KMkBTkDE,65789
+tests/test_lattice.py,sha256=DJoQ3Dr6uAHrdaKofEApc2LD8FgjYAH_a3Ux0cIkgO8,68917
 tests/test_miscs.py,sha256=4fXKsW0kYu2JYO0iGlwWLAYlkFD1rfeVc4xG4Zjn5FQ,8935
 tests/test_mpscircuit.py,sha256=mDXX8oQeFeHr_PdZvwqyDs_tVcVAqLmCERqlTAU7590,10552
 tests/test_noisemodel.py,sha256=UYoMtCjwDaB-CCn5kLosofz-qTMiY4KGAFBjVtqqLPE,5637
@@ -115,10 +115,10 @@ tests/test_shadows.py,sha256=1T3kJesVJ5XfZrSncL80xdq-taGCSnTDF3eL15UlavY,5160
 tests/test_simplify.py,sha256=35tbOu1QANsPvY1buLwNhqPnMkBOsnBtHn82qaukmgI,1175
 tests/test_stabilizer.py,sha256=MivuZ5pY7GOcEPTanhtrflXostyLBToHyjfPqCU0tG0,5450
 tests/test_templates.py,sha256=Xm9otFFaaBWG9TZpgJ-nNh9MBfRipTzFWL8fBOnie2k,7192
-tests/test_timeevol.py,sha256=N2x5BjmOwDStQ1sd0mkdENt7Y_MNES6_3JyMyOF43Iw,14780
+tests/test_timeevol.py,sha256=zz17x21C-5f8ZvcgkXm30JzLgZMhsKaOCzyHCyS43h0,20333
 tests/test_torchnn.py,sha256=CHLTfWkF7Ses5_XnGFN_uv_JddfgenFEFzaDtSH8XYU,2848
 tests/test_van.py,sha256=kAWz860ivlb5zAJuYpzuBe27qccT-Yf0jatf5uXtTo4,3163
-tensorcircuit_nightly-1.3.0.dev20250730.dist-info/METADATA,sha256=o5DEwWjsxTEMdi4Rxhl8KHljen9REwS9D_MtlG8SFdU,34922
-tensorcircuit_nightly-1.3.0.dev20250730.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tensorcircuit_nightly-1.3.0.dev20250730.dist-info/top_level.txt,sha256=O_Iqeh2x02lasEYMI9iyPNNNtMzcpg5qvwMOkZQ7n4A,20
-tensorcircuit_nightly-1.3.0.dev20250730.dist-info/RECORD,,
+tensorcircuit_nightly-1.3.0.dev20250801.dist-info/METADATA,sha256=nHOy6uHlzXTlcH8Iq1t8lJKer9bJ0bTrkamRBgO2l1Q,34922
+tensorcircuit_nightly-1.3.0.dev20250801.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tensorcircuit_nightly-1.3.0.dev20250801.dist-info/top_level.txt,sha256=O_Iqeh2x02lasEYMI9iyPNNNtMzcpg5qvwMOkZQ7n4A,20
+tensorcircuit_nightly-1.3.0.dev20250801.dist-info/RECORD,,

tests/test_backends.py

@@ -9,6 +9,7 @@ os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"
 import numpy as np
 import pytest
 from pytest_lazyfixture import lazy_fixture as lf
+import scipy
 import tensorflow as tf
 
 thisfile = os.path.abspath(__file__)
@@ -60,6 +61,126 @@ def test_grad_torch(torchb):
     np.testing.assert_allclose(f(a), np.ones([2]), atol=1e-5)
 
 
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_sparse_csr_from_coo(backend):
+    # Create a sparse matrix in COO format
+    values = tc.backend.convert_to_tensor(np.array([1.0, 2.0, 3.0]))
+    values = tc.backend.cast(values, "complex64")
+    indices = tc.backend.convert_to_tensor(np.array([[0, 0], [1, 1], [2, 3]]))
+    indices = tc.backend.cast(indices, "int64")
+    coo_matrix = tc.backend.coo_sparse_matrix(indices, values, shape=[4, 4])
+
+    # Convert COO to CSR
+    csr_matrix = tc.backend.sparse_csr_from_coo(coo_matrix)
+
+    # Check that the result is still recognized as sparse
+    assert tc.backend.is_sparse(csr_matrix) is True
+
+    # Check that the conversion preserves values by comparing dense representations
+    coo_dense = tc.backend.to_dense(coo_matrix)
+    csr_dense = tc.backend.to_dense(csr_matrix)
+    np.testing.assert_allclose(coo_dense, csr_dense, atol=1e-5)
+
+
+def test_sparse_tensor_matmul_monkey_patch(tfb):
+    """
+    Test the monkey-patched __matmul__ method for tf.SparseTensor.
+    This test specifically targets the line:
+    tf.SparseTensor.__matmul__ = sparse_tensor_matmul
+    """
+    # Create a sparse matrix in COO format
+    indices = tf.constant([[0, 0], [1, 1], [2, 3]], dtype=tf.int64)
+    values = tf.constant([1.0, 2.0, 3.0], dtype=tf.complex64)
+    shape = [4, 4]
+    sparse_matrix = tf.SparseTensor(indices=indices, values=values, dense_shape=shape)
+
+    # Test 1: Matrix-vector multiplication with 1D vector
+    vector_1d = tf.constant([1.0, 2.0, 3.0, 4.0], dtype=tf.complex64)
+    result_1d = sparse_matrix @ vector_1d  # Using the monkey-patched @ operator
+
+    expected_1d = tf.constant([1.0, 4.0, 12.0, 0.0], dtype=tf.complex64)
+
+    np.testing.assert_allclose(result_1d, expected_1d, atol=1e-6)
+    vector_1d = tc.backend.reshape(vector_1d, [4, 1])
+    result_1dn = sparse_matrix @ vector_1d  # Using the monkey-patched @ operator
+    expected_1d = tc.backend.reshape(expected_1d, [4, 1])
+
+    np.testing.assert_allclose(result_1dn, expected_1d, atol=1e-6)
+
+    # Test 2: Matrix-matrix multiplication with 2D matrix
+    matrix_2d = tf.constant(
+        [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], dtype=tf.complex64
+    )
+    result_2d = sparse_matrix @ matrix_2d  # Using the monkey-patched @ operator
+
+    expected_2d = tf.sparse.sparse_dense_matmul(sparse_matrix, matrix_2d)
+
+    np.testing.assert_allclose(result_2d.numpy(), expected_2d.numpy(), atol=1e-6)
+
+    # Test 3: Verify that the operation is consistent with sparse_dense_matmul
+
+    reference_result = tc.backend.sparse_dense_matmul(sparse_matrix, vector_1d)
+    reference_result_squeezed = tc.backend.reshape(reference_result, [-1])
+
+    np.testing.assert_allclose(result_1d, reference_result_squeezed, atol=1e-6)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("jaxb")])
+def test_backend_jv(backend, highp):
+    def calculate_M(k, x_val):
+        safety_factor = 15
+        M = max(k, int(abs(x_val))) + int(safety_factor * np.sqrt(abs(x_val)))
+        M = max(M, k + 30)
+        return M
+
+    k_values = [5, 20, 50, 200, 500, 3000]
+    x_values = [0.0, 0.1, 1.0, 10.0, 100, 1000, 6000]
+    for k in k_values:
+        for x_val in x_values:
+            M = calculate_M(k, x_val)
+            f_vals = tc.backend.special_jv(k, x_val, M)
+            np.testing.assert_allclose(
+                f_vals, scipy.special.jv(np.arange(k), x_val), atol=1e-6
+            )
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("jaxb")])
+def test_backend_jaxy_scan(backend):
+    def body_fun(carry, x):
+        counter, decrementor = carry
+
+        # 更新状态
+        new_counter = counter + 1
+        new_decrementor = decrementor - 1
+        new_carry = (new_counter, new_decrementor)
+
+        y = counter + decrementor
+
+        return new_carry, y
+
+    init_val = (0, 100)
+
+    final_carry, stacked_ys = tc.backend.jaxy_scan(
+        f=body_fun,
+        init=init_val,
+        xs=tc.backend.arange(5),
+    )
+
+    expected_final_carry = (5, 95)
+    expected_stacked_ys = np.array([100, 100, 100, 100, 100])
+
+    assert final_carry == expected_final_carry
+
+    np.testing.assert_array_equal(np.asarray(stacked_ys), expected_stacked_ys)
+
+
+def test_backend_jv_grad(jaxb, highp):
+    def f(x):
+        return tc.backend.sum(tc.backend.special_jv(5, x, 100))
+
+    print(tc.backend.jit(tc.backend.value_and_grad(f))(0.2))
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
 def test_backend_scatter(backend):
     np.testing.assert_allclose(

tests/test_lattice.py

@@ -23,6 +23,8 @@ from tensorcircuit.templates.lattice import (
     RectangularLattice,
     SquareLattice,
     TriangularLattice,
+    AbstractLattice,
+    get_compatible_layers,
 )
 
 
@@ -1664,3 +1666,85 @@ class TestDistanceMatrix:
 #             "The specialized PBC implementation is significantly slower "
 #             "than the general-purpose implementation."
 #         )
+
+
+def _validate_layers(bonds, layers) -> None:
+    """
+    A helper function to scientifically validate the output of get_compatible_layers.
+    """
+    # MODIFICATION: This function now takes the original bonds list for comparison.
+    expected_edges = set(tuple(sorted(b)) for b in bonds)
+    actual_edges = set(tuple(sorted(edge)) for layer in layers for edge in layer)
+
+    assert (
+        expected_edges == actual_edges
+    ), "Completeness check failed: The set of all edges in the layers must "
+    "exactly match the input bonds."
+
+    for i, layer in enumerate(layers):
+        qubits_in_layer: set[int] = set()
+        for edge in layer:
+            q1, q2 = edge
+            assert (
+                q1 not in qubits_in_layer
+            ), f"Compatibility check failed: Qubit {q1} is reused in layer {i}."
+            qubits_in_layer.add(q1)
+            assert (
+                q2 not in qubits_in_layer
+            ), f"Compatibility check failed: Qubit {q2} is reused in layer {i}."
+            qubits_in_layer.add(q2)
+
+
+@pytest.mark.parametrize(
+    "lattice_instance",
+    [
+        SquareLattice(size=(3, 2), pbc=False),
+        SquareLattice(size=(3, 3), pbc=True),
+        HoneycombLattice(size=(2, 2), pbc=False),
+    ],
+    ids=[
+        "SquareLattice_3x2_OBC",
+        "SquareLattice_3x3_PBC",
+        "HoneycombLattice_2x2_OBC",
+    ],
+)
+def test_layering_on_various_lattices(lattice_instance: AbstractLattice):
+    """Tests gate layering for various standard lattice types."""
+    bonds = lattice_instance.get_neighbor_pairs(k=1, unique=True)
+    layers = get_compatible_layers(bonds)
+
+    assert len(layers) > 0, "Layers should not be empty for non-trivial lattices."
+    _validate_layers(bonds, layers)
+
+
+def test_layering_on_1d_chain_pbc():
+    """Test layering on a 1D chain with periodic boundaries (a cycle graph)."""
+    lattice_even = ChainLattice(size=(6,), pbc=True)
+    bonds_even = lattice_even.get_neighbor_pairs(k=1, unique=True)
+    layers_even = get_compatible_layers(bonds_even)
+    _validate_layers(bonds_even, layers_even)
+
+    lattice_odd = ChainLattice(size=(5,), pbc=True)
+    bonds_odd = lattice_odd.get_neighbor_pairs(k=1, unique=True)
+    layers_odd = get_compatible_layers(bonds_odd)
+    assert len(layers_odd) == 3, "A 5-site cycle graph should be 3-colorable."
+    _validate_layers(bonds_odd, layers_odd)
+
+
+def test_layering_on_custom_star_graph():
+    """Test layering on a custom lattice forming a star graph."""
+    star_edges = [(0, 1), (0, 2), (0, 3)]
+    layers = get_compatible_layers(star_edges)
+    assert len(layers) == 3, "A star graph S_4 requires 3 layers."
+    _validate_layers(star_edges, layers)
+
+
+def test_layering_on_edge_cases():
+    """Test various edge cases: empty, single-site, and no-edge lattices."""
+    layers_empty = get_compatible_layers([])
+    assert layers_empty == [], "Layers should be empty for an empty set of bonds."
+
+    single_edge = [(0, 1)]
+    layers_single = get_compatible_layers(single_edge)
+    assert layers_single == [[(0, 1)]]
+    _validate_layers(single_edge, layers_single)

tests/test_timeevol.py

@@ -279,7 +279,6 @@ def test_krylov_evol_heisenberg_6_sites(backend):
 
     # Generate Heisenberg Hamiltonian
     h = tc.quantum.heisenberg_hamiltonian(g, hzz=1.0, hxx=1.0, hyy=1.0, sparse=False)
-    print(h.dtype)
     # Initial state - all spins up except last one down
     psi0 = np.zeros((2**n,))
     psi0[62] = 1.0
@@ -452,3 +451,191 @@ def test_krylov_evol_gradient(backend):
     grad_fn = tc.backend.jit(tc.backend.grad(loss_function))
     gradient = grad_fn(t)
     print(gradient)
+
+
+@pytest.mark.parametrize(
+    "backend, sparse",
+    [[lf("npb"), True], [lf("npb"), False], [lf("jaxb"), True], [lf("jaxb"), False]],
+)
+def test_chebyshev_evol_basic(backend, highp, sparse):
+    n = 6
+    # Create a 1D chain graph
+    g = tc.templates.graphs.Line1D(n, pbc=False)
+
+    # Generate Heisenberg Hamiltonian (dense for better compatibility)
+    h = tc.quantum.heisenberg_hamiltonian(
+        g, hzz=1.0, hxx=1.0, hyy=1.0, hx=0.2, sparse=sparse
+    )
+
+    # Initial Neel state: |↑↓↑↓⟩
+    c = tc.Circuit(n)
+    c.x([1, 3, 5])  # Apply X gates to qubits 1 and 3
+    psi0 = c.state()
+
+    # Evolution time
+    t = 2.0
+
+    # Estimate spectral bounds
+    e_max, e_min = tc.timeevol.estimate_spectral_bounds(h, n_iter=30)
+
+    # Estimate parameters
+    k = tc.timeevol.estimate_k(t, (e_max, e_min))
+    m = tc.timeevol.estimate_M(t, (e_max, e_min), k)
+
+    # Evolve using Chebyshev method
+    psi_chebyshev = tc.timeevol.chebyshev_evol(
+        h, psi0, t, (float(e_max) + 0.1, float(e_min) - 0.1), k, m
+    )
+
+    # Check that state is normalized (or close to it)
+    norm = tc.backend.norm(psi_chebyshev)
+    np.testing.assert_allclose(norm, 1.0, atol=1e-3)
+
+    # Compare with exact evolution for small system
+    if sparse is True:
+        h = tc.backend.to_dense(h)
+    psi_exact = tc.timeevol.ed_evol(h, psi0, 1.0j * tc.backend.convert_to_tensor([t]))[
+        0
+    ]
+
+    # States should be close (up to global phase)
+    fidelity = np.abs(np.vdot(np.asarray(psi_exact), np.asarray(psi_chebyshev))) ** 2
+    assert fidelity > 0.95  # Should be close, but not exact due to approximations
+
+
+def test_chebyshev_evol_vmap_on_t(jaxb, highp):
+    n = 4
+    # Create a 1D chain graph
+    g = tc.templates.graphs.Line1D(n, pbc=False)
+
+    # Generate Heisenberg Hamiltonian
+    h = tc.quantum.heisenberg_hamiltonian(g, hzz=1.0, hxx=1.0, hyy=1.0, sparse=False)
+
+    # Initial Neel state
+    c = tc.Circuit(n)
+    c.x([1, 3])
+    psi0 = c.state()
+
+    # Estimate spectral bounds
+    e_max, e_min = tc.timeevol.estimate_spectral_bounds(h, n_iter=20)
+
+    # Fixed parameters
+    k = 50
+    m = 150
+
+    # Define vectorized evolution function
+    def evolve_single_time(t):
+        return tc.timeevol.chebyshev_evol(
+            h, psi0, t, (float(e_max) + 0.1, float(e_min) - 0.1), k, m
+        )
+
+    # Vectorize over times
+    times = tc.backend.convert_to_tensor([0.5, 1.0, 1.5])
+    vmap_evolve = tc.backend.jit(tc.backend.vmap(evolve_single_time))
+    states_vmap = vmap_evolve(times)
+
+    # Check output shape
+    assert states_vmap.shape == (3, 2**n)
+
+    # Compare with sequential execution
+    states_sequential = []
+    for t in times:
+        state = tc.timeevol.chebyshev_evol(
+            h, psi0, float(t), (e_max + 0.1, e_min - 0.1), k, m
+        )
+        states_sequential.append(state)
+
+    states_sequential = tc.backend.stack(states_sequential)
+
+    # Results should be the same
+    np.testing.assert_allclose(states_vmap, states_sequential, atol=1e-5)
+
+
+def test_chebyshev_evol_jit_on_psi(jaxb, highp):
+    """Test JIT compilation capability of chebyshev_evol on psi parameter"""
+    n = 4
+    # Create a 1D chain graph
+    g = tc.templates.graphs.Line1D(n, pbc=False)
+
+    # Generate Heisenberg Hamiltonian
+    h = tc.quantum.heisenberg_hamiltonian(g, hzz=1.0, hxx=0.6, hyy=1.0, sparse=True)
+
+    # Estimate spectral bounds
+    e_max, e_min = tc.timeevol.estimate_spectral_bounds(h, n_iter=20)
+
+    # Fixed parameters
+    t = 1.0
+    k = 50
+    m = 150
+
+    # Define JIT-compiled evolution function with psi as argument
+    def evolve_state(psi):
+        return tc.timeevol.chebyshev_evol(
+            h, psi, t, (float(e_max) + 0.1, float(e_min) - 0.1), k, m
+        )
+
+    jit_evolve = tc.backend.jit(evolve_state)
+
+    # Test with different initial states
+    c1 = tc.Circuit(n)
+    c1.x([0, 2])
+    psi1 = c1.state()
+
+    c2 = tc.Circuit(n)
+    c2.h(0)
+    for i in range(n - 1):
+        c2.cnot(i, i + 1)
+    psi2 = c2.state()
+
+    # Run JIT-compiled evolution
+    result1_jit = jit_evolve(psi1)
+    result2_jit = jit_evolve(psi2)
+
+    # Run regular evolution for comparison
+    result1_regular = tc.timeevol.chebyshev_evol(
+        h, psi1, t, (e_max + 0.1, e_min - 0.1), k, m
+    )
+    result2_regular = tc.timeevol.chebyshev_evol(
+        h, psi2, t, (e_max + 0.1, e_min - 0.1), k, m
+    )
+    print(result1_jit)
+    # Results should be the same
+    np.testing.assert_allclose(result1_jit, result1_regular, atol=1e-5)
+    np.testing.assert_allclose(result2_jit, result2_regular, atol=1e-5)
+
+
+def test_chebyshev_evol_ad_on_t(jaxb, highp):
+    n = 5
+    # Create a 1D chain graph
+    g = tc.templates.graphs.Line1D(n, pbc=True)
+
+    # Generate Heisenberg Hamiltonian
+    h = tc.quantum.heisenberg_hamiltonian(g, hzz=1.0, hxx=1.0, hyy=1.0, sparse=True)
+
+    # Initial state
+    c = tc.Circuit(n)
+    c.x([1, 3])
+    psi0 = c.state()
+
+    # Estimate spectral bounds
+    e_max, e_min = tc.timeevol.estimate_spectral_bounds(h, n_iter=20)
+
+    # Fixed parameters
+    k = 50
+    m = 100
+
+    # Define loss function for gradient computation
+    def loss_function(t):
+        psi_t = tc.timeevol.chebyshev_evol(
+            h, psi0, t, (float(e_max) + 0.1, float(e_min) - 0.1), k, m
+        )
+        c = tc.Circuit(5, inputs=psi_t)
+        return tc.backend.real(c.expectation_ps(z=[2]))
+
+    # Compute gradient
+    grad_fn = tc.backend.jit(tc.backend.grad(loss_function))
+    t_test = tc.backend.convert_to_tensor(1.0)
+    gradient = grad_fn(t_test)
+    print(gradient)
+    # Gradient should be a scalar
+    assert gradient.shape == ()