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

tensorcircuit/__init__.py

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

tensorcircuit/backends/abstract_backend.py

@@ -1389,9 +1389,21 @@ 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
+        # carry = init
+        # for x in xs:
+        #     carry = f(carry, x)
+        # return carry
 
     def stop_gradient(self: Any, a: Tensor) -> Tensor:
         """

tensorcircuit/backends/numpy_backend.py

@@ -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

tensorcircuit/backends/tensorflow_backend.py

@@ -719,7 +719,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/timeevol.py

@@ -3,14 +3,165 @@ Analog time evolution engines
 """
 
 from typing import Any, Tuple, Optional, Callable, List, Sequence
+from functools import partial
+
+import numpy as np
 
 from .cons import backend, dtypestr, rdtypestr, contractor
 from .gates import Gate
+from .utils import arg_alias
 
 Tensor = Any
 Circuit = Any
 
 
+def lanczos_iteration_scan(
+    hamiltonian: Any, initial_vector: Any, subspace_dimension: int
+) -> Tuple[Any, Any]:
+    """
+    Use Lanczos algorithm to construct orthogonal basis and projected Hamiltonian
+    of Krylov subspace, using `tc.backend.scan` for JIT compatibility.
+
+    :param hamiltonian: Sparse or dense Hamiltonian matrix
+    :type hamiltonian: Tensor
+    :param initial_vector: Initial quantum state vector
+    :type initial_vector: Tensor
+    :param subspace_dimension: Dimension of Krylov subspace
+    :type subspace_dimension: int
+    :return: Tuple containing (basis matrix, projected Hamiltonian)
+    :rtype: Tuple[Tensor, Tensor]
+    """
+    state_size = backend.shape_tuple(initial_vector)[0]
+
+    # Main scan body for the outer loop (iterating j)
+    def lanczos_step(carry: Tuple[Any, ...], j: int) -> Tuple[Any, ...]:
+        v, basis, alphas, betas = carry
+
+        if backend.is_sparse(hamiltonian):
+            w = backend.sparse_dense_matmul(hamiltonian, v)
+        else:
+            w = backend.matvec(hamiltonian, v)
+
+        alpha = backend.real(backend.sum(backend.conj(v) * w))
+        w = w - backend.cast(alpha, dtypestr) * v
+
+        # Inner scan for re-orthogonalization (iterating k)
+        # def ortho_step(inner_carry: Tuple[Any, Any], k: int) -> Tuple[Any, Any]:
+        #     w_carry, j_val = inner_carry
+
+        #     def do_projection() -> Any:
+        #         # `basis` is available here through closure
+        #         v_k = basis[:, k]
+        #         projection = backend.sum(backend.conj(v_k) * w_carry)
+        #         return w_carry - projection * v_k
+
+        #     def do_nothing() -> Any:
+        #         return w_carry
+
+        #     # Orthogonalize against v_0, ..., v_j
+        #     w_new = backend.cond(k <= j_val, do_projection, do_nothing)
+        #     return (w_new, j_val)  # Return the new carry for the inner loop
+
+        # # Pass `j` into the inner scan's carry
+        # inner_init_carry = (w, j)
+        # final_inner_carry = backend.scan(
+        #     ortho_step, backend.arange(subspace_dimension), inner_init_carry
+        # )
+        # w_ortho = final_inner_carry[0]
+
+        def ortho_step(w_carry: Any, elems_tuple: Tuple[Any, Any]) -> Any:
+            k, j_from_elems = elems_tuple
+
+            def do_projection() -> Any:
+                v_k = basis[:, k]
+                projection = backend.sum(backend.conj(v_k) * w_carry)
+                return w_carry - projection * v_k
+
+            def do_nothing() -> Any:
+                return backend.cast(w_carry, dtype=dtypestr)
+
+            w_new = backend.cond(k <= j_from_elems, do_projection, do_nothing)
+            return w_new
+
+        k_elems = backend.arange(subspace_dimension)
+        j_elems = backend.tile(backend.reshape(j, [1]), [subspace_dimension])
+        inner_elems = (k_elems, j_elems)
+        w_ortho = backend.scan(ortho_step, inner_elems, w)
+
+        beta = backend.norm(w_ortho)
+        beta = backend.real(beta)
+
+        # Update alphas and betas arrays
+        new_alphas = backend.scatter(
+            alphas, backend.reshape(j, [1, 1]), backend.reshape(alpha, [1])
+        )
+        new_betas = backend.scatter(
+            betas, backend.reshape(j, [1, 1]), backend.reshape(beta, [1])
+        )
+
+        def update_state_fn() -> Tuple[Any, Any]:
+            epsilon = 1e-15
+            next_v = w_ortho / backend.cast(beta + epsilon, dtypestr)
+
+            one_hot_update = backend.onehot(j + 1, subspace_dimension)
+            one_hot_update = backend.cast(one_hot_update, dtype=dtypestr)
+
+            # Create a mask to update only the (j+1)-th column
+            mask = 1.0 - backend.reshape(one_hot_update, [1, subspace_dimension])
+            new_basis = basis * mask + backend.reshape(
+                next_v, [-1, 1]
+            ) * backend.reshape(one_hot_update, [1, subspace_dimension])
+
+            return next_v, new_basis
+
+        def keep_state_fn() -> Tuple[Any, Any]:
+            return v, basis
+
+        next_v_carry, new_basis = backend.cond(
+            j < subspace_dimension - 1, update_state_fn, keep_state_fn
+        )
+
+        return (next_v_carry, new_basis, new_alphas, new_betas)
+
+    # Prepare initial state for the main scan
+    v0 = initial_vector / backend.norm(initial_vector)
+
+    init_basis = backend.zeros((state_size, subspace_dimension), dtype=dtypestr)
+    init_alphas = backend.zeros((subspace_dimension,), dtype=rdtypestr)
+    init_betas = backend.zeros((subspace_dimension,), dtype=rdtypestr)
+
+    one_hot_0 = backend.onehot(0, subspace_dimension)
+    one_hot_0 = backend.cast(one_hot_0, dtype=dtypestr)
+    init_basis = init_basis + backend.reshape(v0, [-1, 1]) * backend.reshape(
+        one_hot_0, [1, subspace_dimension]
+    )
+
+    init_carry = (v0, init_basis, init_alphas, init_betas)
+
+    # Run the main scan
+    final_carry = backend.scan(
+        lanczos_step, backend.arange(subspace_dimension), init_carry
+    )
+    basis_matrix, alphas_tensor, betas_tensor = (
+        final_carry[1],
+        final_carry[2],
+        final_carry[3],
+    )
+
+    betas_off_diag = betas_tensor[:-1]
+
+    diag_part = backend.diagflat(alphas_tensor)
+    if backend.shape_tuple(betas_off_diag)[0] > 0:
+        off_diag_part = backend.diagflat(betas_off_diag, k=1)
+        projected_hamiltonian = (
+            diag_part + off_diag_part + backend.conj(backend.transpose(off_diag_part))
+        )
+    else:
+        projected_hamiltonian = diag_part
+
+    return basis_matrix, projected_hamiltonian
+
+
 def lanczos_iteration(
     hamiltonian: Tensor, initial_vector: Tensor, subspace_dimension: int
 ) -> Tuple[Tensor, Tensor]:
@@ -113,9 +264,10 @@ def lanczos_iteration(
 def krylov_evol(
     hamiltonian: Tensor,
     initial_state: Tensor,
-    time_points: Tensor,
+    times: Tensor,
     subspace_dimension: int,
     callback: Optional[Callable[[Any], Any]] = None,
+    scan_impl: bool = False,
 ) -> Any:
     """
     Perform quantum state time evolution using Krylov subspace method.
@@ -124,21 +276,30 @@ def krylov_evol(
     :type hamiltonian: Tensor
     :param initial_state: Initial quantum state
     :type initial_state: Tensor
-    :param time_points: List of time points
-    :type time_points: Tensor
+    :param times: List of time points
+    :type times: Tensor
     :param subspace_dimension: Krylov subspace dimension
     :type subspace_dimension: int
     :param callback: Optional callback function applied to quantum state at
                   each evolution time point, return some observables
     :type callback: Optional[Callable[[Any], Any]], optional
+    :param scan_impl: whether use scan implementation, suitable for jit but may be slow on numpy
+        defaults False, True not work for tensorflow backend + jit, due to stupid issue of tensorflow
+        context separation and the notorious inaccesibletensor error
+    :type scan_impl: bool, optional
     :return: List of evolved quantum states, or list of callback function results
         (if callback provided)
     :rtype: Any
     """
     # TODO(@refraction-ray): stable and efficient AD is to be investigated
-    basis_matrix, projected_hamiltonian = lanczos_iteration(
-        hamiltonian, initial_state, subspace_dimension
-    )
+    if not scan_impl:
+        basis_matrix, projected_hamiltonian = lanczos_iteration(
+            hamiltonian, initial_state, subspace_dimension
+        )
+    else:
+        basis_matrix, projected_hamiltonian = lanczos_iteration_scan(
+            hamiltonian, initial_state, subspace_dimension
+        )
     initial_state = backend.cast(initial_state, dtypestr)
     # Project initial state to Krylov subspace: |psi_proj> = V_m^† |psi(0)>
     projected_state = backend.matvec(
@@ -148,8 +309,9 @@ def krylov_evol(
     # Perform spectral decomposition of projected Hamiltonian: T_m = U D U^†
     eigenvalues, eigenvectors = backend.eigh(projected_hamiltonian)
     eigenvalues = backend.cast(eigenvalues, dtypestr)
-    time_points = backend.convert_to_tensor(time_points)
-    time_points = backend.cast(time_points, dtypestr)
+    eigenvectors = backend.cast(eigenvectors, dtypestr)
+    times = backend.convert_to_tensor(times)
+    times = backend.cast(times, dtypestr)
 
     # Transform projected state to eigenbasis: |psi_coeff> = U^† |psi_proj>
     eigenvectors_projected_state = backend.matvec(
@@ -158,7 +320,7 @@ def krylov_evol(
 
     # Calculate exp(-i*projected_H*t) * projected_state
     results = []
-    for t in time_points:
+    for t in times:
         # Calculate exp(-i*eigenvalues*t)
         exp_diagonal = backend.exp(-1j * eigenvalues * t)
 
@@ -182,22 +344,26 @@ def krylov_evol(
     return backend.stack(results)
 
 
+@partial(
+    arg_alias,
+    alias_dict={"h": ["hamiltonian"], "psi0": ["initial_state"], "tlist": ["times"]},
+)
 def hamiltonian_evol(
-    tlist: Tensor,
     h: Tensor,
     psi0: Tensor,
+    tlist: Tensor,
     callback: Optional[Callable[..., Any]] = None,
 ) -> Tensor:
     """
     Fast implementation of time independent Hamiltonian evolution using eigendecomposition.
     By default, performs imaginary time evolution.
 
-    :param tlist: Time points for evolution
-    :type tlist: Tensor
     :param h: Time-independent Hamiltonian matrix
     :type h: Tensor
     :param psi0: Initial state vector
     :type psi0: Tensor
+    :param tlist: Time points for evolution
+    :type tlist: Tensor
     :param callback: Optional function to process state at each time point
     :type callback: Optional[Callable[..., Any]], optional
     :return: Evolution results at each time point. If callback is None, returns state vectors;
@@ -233,9 +399,10 @@ def hamiltonian_evol(
     psi0 = backend.cast(psi0, dtypestr)
     es, u = backend.eigh(h)
     u = backend.cast(u, dtypestr)
-    utpsi0 = backend.reshape(
-        backend.transpose(u) @ backend.reshape(psi0, [-1, 1]), [-1]
-    )
+    utpsi0 = backend.convert_to_tensor(
+        backend.transpose(u) @ backend.reshape(psi0, [-1, 1])
+    )  # in case np.matrix...
+    utpsi0 = backend.reshape(utpsi0, [-1])
     es = backend.cast(es, dtypestr)
     tlist = backend.cast(backend.convert_to_tensor(tlist), dtypestr)
 
@@ -255,6 +422,7 @@ def hamiltonian_evol(
 ed_evol = hamiltonian_evol
 
 
+@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
 def evol_local(
     c: Circuit,
     index: Sequence[int],
@@ -279,16 +447,58 @@ def evol_local(
     :return: _description_
     :rtype: Circuit
     """
+    s = c.state()
+    n = int(np.log2(s.shape[-1]) + 1e-7)
+    if isinstance(t, float):
+        t = backend.stack([0.0, t])
+    s1 = ode_evol_local(h_fun, s, t, index, None, *args, **solver_kws)
+    return type(c)(n, inputs=s1[-1])
+
+
+def ode_evol_local(
+    hamiltonian: Callable[..., Tensor],
+    initial_state: Tensor,
+    times: Tensor,
+    index: Sequence[int],
+    callback: Optional[Callable[..., Tensor]] = None,
+    *args: Any,
+    **solver_kws: Any,
+) -> Tensor:
+    """
+    ODE-based time evolution for a time-dependent Hamiltonian acting on a subsystem of qubits.
+
+    This function solves the time-dependent Schrodinger equation using numerical ODE integration.
+    The Hamiltonian is applied only to a specific subset of qubits (indices) in the system.
+
+    Note: This function currently only supports the JAX backend.
+
+    :param hamiltonian: A function that returns a dense Hamiltonian matrix for the specified
+        subsystem size. The function signature should be hamiltonian(time, *args) -> Tensor.
+    :type hamiltonian: Callable[..., Tensor]
+    :param initial_state: The initial quantum state vector of the full system.
+    :type initial_state: Tensor
+    :param times: Time points for which to compute the evolution. Should be a 1D array of times.
+    :type times: Tensor
+    :param index: Indices of qubits where the Hamiltonian is applied.
+    :type index: Sequence[int]
+    :param callback: Optional function to apply to the state at each time step.
+    :type callback: Optional[Callable[..., Tensor]]
+    :param args: Additional arguments to pass to the Hamiltonian function.
+    :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+    :return: Evolved quantum states at the specified time points. If callback is provided,
+        returns the callback results; otherwise returns the state vectors.
+    :rtype: Tensor
+    """
     from jax.experimental.ode import odeint
 
-    s = c.state()
-    n = c._nqubits
+    s = initial_state
+    n = int(np.log2(backend.shape_tuple(initial_state)[-1]) + 1e-7)
     l = len(index)
 
     def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
         y = backend.reshape2(y)
         y = Gate(y)
-        h = -1.0j * h_fun(t, *args)
+        h = -1.0j * hamiltonian(t, *args)
         h = backend.reshape2(h)
         h = Gate(h)
         edges = []
@@ -302,15 +512,15 @@ def evol_local(
         y = contractor([y, h], output_edge_order=edges)
         return backend.reshape(y.tensor, [-1])
 
-    ts = backend.stack([0.0, t])
+    ts = backend.convert_to_tensor(times)
     ts = backend.cast(ts, dtype=rdtypestr)
     s1 = odeint(f, s, ts, *args, **solver_kws)
-    return type(c)(n, inputs=s1[-1])
-
-
-ode_evol_local = evol_local
+    if not callback:
+        return s1
+    return backend.stack([callback(s1[i]) for i in range(len(s1))])
 
 
+@partial(arg_alias, alias_dict={"h_fun": ["hamiltonian"], "t": ["times"]})
 def evol_global(
     c: Circuit, h_fun: Callable[..., Tensor], t: float, *args: Any, **solver_kws: Any
 ) -> Circuit:
@@ -328,19 +538,57 @@ def evol_global(
     :return: _description_
     :rtype: Circuit
     """
-    from jax.experimental.ode import odeint
-
     s = c.state()
     n = c._nqubits
+    if isinstance(t, float):
+        t = backend.stack([0.0, t])
+    s1 = ode_evol_global(h_fun, s, t, None, *args, **solver_kws)
+    return type(c)(n, inputs=s1[-1])
 
-    def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
-        h = -1.0j * h_fun(t, *args)
-        return backend.sparse_dense_matmul(h, y)
 
-    ts = backend.stack([0.0, t])
+def ode_evol_global(
+    hamiltonian: Callable[..., Tensor],
+    initial_state: Tensor,
+    times: Tensor,
+    callback: Optional[Callable[..., Tensor]] = None,
+    *args: Any,
+    **solver_kws: Any,
+) -> Tensor:
+    """
+    ODE-based time evolution for a time-dependent Hamiltonian acting on the entire system.
+
+    This function solves the time-dependent Schrodinger equation using numerical ODE integration.
+    The Hamiltonian is applied to the full system and should be provided in sparse matrix format
+    for efficiency.
+
+    Note: This function currently only supports the JAX backend.
+
+    :param hamiltonian: A function that returns a sparse Hamiltonian matrix for the full system.
+        The function signature should be hamiltonian(time, *args) -> Tensor.
+    :type hamiltonian: Callable[..., Tensor]
+    :param initial_state: The initial quantum state vector.
+    :type initial_state: Tensor
+    :param times: Time points for which to compute the evolution. Should be a 1D array of times.
+    :type times: Tensor
+    :param callback: Optional function to apply to the state at each time step.
+    :type callback: Optional[Callable[..., Tensor]]
+    :param args: Additional arguments to pass to the Hamiltonian function.
+    :param solver_kws: Additional keyword arguments to pass to the ODE solver.
+    :return: Evolved quantum states at the specified time points. If callback is provided,
+        returns the callback results; otherwise returns the state vectors.
+    :rtype: Tensor
+    """
+    from jax.experimental.ode import odeint
+
+    s = initial_state
+    ts = backend.convert_to_tensor(times)
     ts = backend.cast(ts, dtype=rdtypestr)
-    s1 = odeint(f, s, ts, *args, **solver_kws)
-    return type(c)(n, inputs=s1[-1])
 
+    def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
+        h = -1.0j * hamiltonian(t, *args)
+        return backend.sparse_dense_matmul(h, y)
 
-ode_evol_global = evol_global
+    s1 = odeint(f, s, ts, *args, **solver_kws)
+    if not callback:
+        return s1
+    return backend.stack([callback(s1[i]) for i in range(len(s1))])

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorcircuit-nightly
-Version: 1.3.0.dev20250729
+Version: 1.3.0.dev20250730
 Summary: nightly release for tensorcircuit
 Home-page: https://github.com/refraction-ray/tensorcircuit-dev
 Author: TensorCircuit Authors

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

@@ -1,4 +1,4 @@
-tensorcircuit/__init__.py,sha256=GNg45VO8J-m5ApDmmQbFvTPIzXvBpREXDg5MiVgCqqI,2055
+tensorcircuit/__init__.py,sha256=DrH3d7oEWpI2MmrSmabmwoovMMYBNoLZeCk6SFC360M,2055
 tensorcircuit/about.py,sha256=DazTswU2nAwOmASTaDII3L04PVtaQ7oiWPty5YMI3Wk,5267
 tensorcircuit/abstractcircuit.py,sha256=0osacPqq7B1EJki-cI1aLYoVRmjFaG9q3XevWMs7SsA,44125
 tensorcircuit/asciiart.py,sha256=neY1OWFwtoW5cHPNwkQHgRPktDniQvdlP9QKHkk52fM,8236
@@ -18,7 +18,7 @@ tensorcircuit/quantum.py,sha256=LNkIv5cJ2KG6puC18zTuXi-5cojW1Tnz-N-WjZ0Qu5Q,9021
 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=LWKOw4Y4CUtUO4a_72BVzYPH26PbubGYOaJlXkmA350,11439
+tensorcircuit/timeevol.py,sha256=8p4C3nhUQ9eC2wYfZ9w5BGFIt25NPEJTXmZU_iZy4tM,21607
 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=tzhOFZSOkZ_iraMvY7g84VbWBpYcELX1O0bhUipZ7YQ,59131
+tensorcircuit/backends/abstract_backend.py,sha256=pmuf8NFmtRXnlBTsRrv2s6weCNiq3cd9gA7wiKvIobM,59589
 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=sd1migp_E2FWjchvOeYRuyM47yexegT2_SW_ukSYSF8,14171
+tensorcircuit/backends/numpy_backend.py,sha256=PhbpXeATQ6X4kZ3xA-RvQVO_dKrpFh5vyXcCKKDMC7U,14197
 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=N8iL2P76ZRp3USJSyFdZ2LWG8muXYe1kOuengv1ijXo,36444
+tensorcircuit/backends/tensorflow_backend.py,sha256=T2BmFxOyl2QU4dSIwUMCLyPspLrNFLx3hVfD7TD1No0,36598
 tensorcircuit/backends/tf_ops.py,sha256=FJwDU7LhZrt0VUIx12DJU0gZnWhMv7B7r9sAKG710As,3378
 tensorcircuit/cloud/__init__.py,sha256=n0Lx07GYF6YbdIa6AJCLJk4zlAm5CqaeHszvkxxuoI4,139
 tensorcircuit/cloud/abstraction.py,sha256=6aSxbz0MP21jBVdFbSMrvJPLQH117vGz9sSHbMFoodE,14582
@@ -86,7 +86,7 @@ tensorcircuit/templates/graphs.py,sha256=cPYrxjoem0xZ-Is9dZKAvEzWZL_FejfIRiCEOTA
 tensorcircuit/templates/hamiltonians.py,sha256=Ag8djD6lckTeU7I99gCbXiQAb2VYqzm_p7-hpXo-5u4,5554
 tensorcircuit/templates/lattice.py,sha256=F35ebANk0DSmSHLR0-Q_hUbcznyCmZjb4fKmvCMywmA,58575
 tensorcircuit/templates/measurements.py,sha256=pzc5Aa9S416Ilg4aOY77Z6ZhUlYcXnAkQNQFTuHjFFs,10943
-tensorcircuit_nightly-1.3.0.dev20250729.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+tensorcircuit_nightly-1.3.0.dev20250730.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
@@ -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=b5fsUzW3kfLZV909jnJpH9vCBz5fefONOUSrm2N6Nxg,7886
+tests/test_timeevol.py,sha256=N2x5BjmOwDStQ1sd0mkdENt7Y_MNES6_3JyMyOF43Iw,14780
 tests/test_torchnn.py,sha256=CHLTfWkF7Ses5_XnGFN_uv_JddfgenFEFzaDtSH8XYU,2848
 tests/test_van.py,sha256=kAWz860ivlb5zAJuYpzuBe27qccT-Yf0jatf5uXtTo4,3163
-tensorcircuit_nightly-1.3.0.dev20250729.dist-info/METADATA,sha256=MijyspDQzOzMCKoGxFN50a89pOhICnhrpUeXdPASpQg,34922
-tensorcircuit_nightly-1.3.0.dev20250729.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tensorcircuit_nightly-1.3.0.dev20250729.dist-info/top_level.txt,sha256=O_Iqeh2x02lasEYMI9iyPNNNtMzcpg5qvwMOkZQ7n4A,20
-tensorcircuit_nightly-1.3.0.dev20250729.dist-info/RECORD,,
+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,,

tests/test_timeevol.py

@@ -11,7 +11,7 @@ sys.path.insert(0, modulepath)
 import tensorcircuit as tc
 
 
-def test_ode_evol(jaxb):
+def test_circuit_ode_evol(jaxb):
     def h_square(t, b):
         return (tc.backend.sign(t - 1.0) + 1) / 2 * b * tc.gates.x().tensor
 
@@ -19,9 +19,7 @@ def test_ode_evol(jaxb):
     c.x(0)
     c.cx(0, 1)
     c.h(2)
-    c = tc.timeevol.ode_evol_local(
-        c, [1], h_square, 2.0, tc.backend.convert_to_tensor(0.2)
-    )
+    c = tc.timeevol.evol_local(c, [1], h_square, 2.0, tc.backend.convert_to_tensor(0.2))
     c.rx(1, theta=np.pi - 0.4)
     np.testing.assert_allclose(c.expectation_ps(z=[1]), 1.0, atol=1e-5)
 
@@ -34,13 +32,156 @@ def test_ode_evol(jaxb):
     c.x(0)
     c.cx(0, 1)
     c.h(2)
-    c = tc.timeevol.ode_evol_global(
+    c = tc.timeevol.evol_global(
         c, h_square_sparse, 2.0, tc.backend.convert_to_tensor(0.2)
     )
     c.rx(1, theta=np.pi - 0.4)
     np.testing.assert_allclose(c.expectation_ps(z=[1]), 1.0, atol=1e-5)
 
 
+def test_ode_evol_local(jaxb):
+    def local_hamiltonian(t, Omega, phi):
+        angle = phi * t
+        coeff = Omega * tc.backend.cos(2.0 * t)  # Amplitude modulation
+
+        # Single-qubit Rabi Hamiltonian (2x2 matrix)
+        hx = coeff * tc.backend.cos(angle) * tc.gates.x().tensor
+        hy = coeff * tc.backend.sin(angle) * tc.gates.y().tensor
+        return hx + hy
+
+    # Initial state: GHZ state |0000⟩ + |1111⟩
+    c = tc.Circuit(4)
+    c.h(0)
+    for i in range(3):
+        c.cnot(i, i + 1)
+    psi0 = c.state()
+
+    # Time points
+    times = tc.backend.arange(0.0, 3.0, 0.1)
+
+    # Evolve with local Hamiltonian acting on qubit 1
+    states = tc.timeevol.ode_evol_local(
+        local_hamiltonian,
+        psi0,
+        times,
+        [1],  # Apply to qubit 1
+        None,
+        1.0,
+        2.0,  # Omega=1.0, phi=2.0
+    )
+    assert tc.backend.shape_tuple(states) == (30, 16)
+
+
+def test_ode_evol_global(jaxb):
+    # Create a time-dependent transverse field Hamiltonian
+    # H(t) = -∑ᵢ Jᵢ(t) ZᵢZᵢ₊₁ - ∑ᵢ hᵢ(t) Xᵢ
+
+    # Time-dependent coefficients
+    def time_dep_J(t):
+        return 1.0 + 0.5 * tc.backend.sin(2.0 * t)
+
+    def time_dep_h(t):
+        return 0.5 * tc.backend.cos(1.5 * t)
+
+    zz_ham = tc.quantum.PauliStringSum2COO(
+        [[3, 3, 0, 0], [0, 3, 3, 0], [0, 0, 3, 3]], [1, 1, 1]
+    )
+    x_ham = tc.quantum.PauliStringSum2COO(
+        [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], [1, 1, 1, 1]
+    )
+
+    # Hamiltonian construction function
+    def hamiltonian_func(t):
+        # Create time-dependent ZZ terms
+        zz_coeff = time_dep_J(t)
+
+        # Create time-dependent X terms
+        x_coeff = time_dep_h(t)
+
+        return zz_coeff * zz_ham + x_coeff * x_ham
+
+    # Initial state: |↑↓↑↓⟩
+    c = tc.Circuit(4)
+    c.x([1, 3])
+    psi0 = c.state()
+
+    # Time points for evolution
+    times = tc.backend.arange(0, 5, 0.5)
+
+    def zobs(state):
+        n = int(np.log2(state.shape[-1]))
+        c = tc.Circuit(n, inputs=state)
+        return tc.backend.real(c.expectation_ps(z=[0]))
+
+    # Perform global ODE evolution
+    states = tc.timeevol.ode_evol_global(hamiltonian_func, psi0, times, zobs)
+    assert tc.backend.shape_tuple(states) == (10,)
+
+    zz_ham = tc.quantum.PauliStringSum2COO([[3, 3, 0, 0], [0, 3, 3, 0]], [1, 1])
+    x_ham = tc.quantum.PauliStringSum2COO([[1, 0, 0, 0], [0, 1, 0, 0]], [1, 1])
+
+    # Example with parameterized Hamiltonian and optimization
+    def parametrized_hamiltonian(t, params):
+        # params = [J0, J1, h0, h1] - parameters to optimize
+        J_t = params[0] + params[1] * tc.backend.sin(2.0 * t)
+        h_t = params[2] + params[3] * tc.backend.cos(1.5 * t)
+
+        return J_t * zz_ham + h_t * x_ham
+
+    # Observable function: measure ZZ correlation
+    def zz_correlation(state):
+        n = int(np.log2(state.shape[0]))
+        circuit = tc.Circuit(n, inputs=state)
+        return circuit.expectation_ps(z=[0, 1])
+
+    @tc.backend.jit
+    @tc.backend.value_and_grad
+    def objective_function(params):
+        states = tc.timeevol.ode_evol_global(
+            parametrized_hamiltonian,
+            psi0,
+            tc.backend.convert_to_tensor([0, 1.0]),
+            None,
+            params,
+        )
+        # Measure ZZ correlation at final time
+        final_state = states[-1]
+        return tc.backend.real(zz_correlation(final_state))
+
+    print(objective_function(tc.backend.ones([4])))
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_ed_evol(backend):
+    n = 4
+    g = tc.templates.graphs.Line1D(n, pbc=False)
+    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])  # Apply X gates to qubits 1 and 3
+    psi0 = c.state()
+
+    # Imaginary time evolution times
+    times = tc.backend.convert_to_tensor([0.0, 0.5, 1.0, 2.0])
+
+    # Evolve and get states
+    states = tc.timeevol.ed_evol(h, psi0, times)
+    print(states)
+
+    def evolve_and_measure(params):
+        # Parametrized Hamiltonian
+        h_param = tc.quantum.heisenberg_hamiltonian(
+            g, hzz=params[0], hxx=params[1], hyy=params[2], sparse=False
+        )
+        states = tc.timeevol.ed_evol(h_param, psi0, times)
+        # Measure observable on final state
+        circuit = tc.Circuit(n, inputs=states[-1])
+        return tc.backend.real(circuit.expectation_ps(z=[0]))
+
+    evolve_and_measure(tc.backend.ones([3]))
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
 def test_hamiltonian_evol_basic(backend):
     """Test basic functionality of hamiltonian_evol with a simple 2-qubit Hamiltonian"""
@@ -63,7 +204,7 @@ def test_hamiltonian_evol_basic(backend):
     )
 
     # Evolve and get states
-    states = tc.timeevol.hamiltonian_evol(times, h, psi0)
+    states = tc.timeevol.hamiltonian_evol(h, psi0, times)
 
     # Check output shape
     assert states.shape == (3, 4)
@@ -99,7 +240,7 @@ def test_hamiltonian_evol_with_callback(backend):
         return tc.backend.real(c.expectation_ps(z=[0]))
 
     # Evolve with callback
-    results = tc.timeevol.hamiltonian_evol(times, h, psi0, callback)
+    results = tc.timeevol.hamiltonian_evol(h, psi0, times, callback)
 
     # Check output shape - should be scalar for each time point
     assert results.shape == (3,)
@@ -121,7 +262,7 @@ def test_hamiltonian_evol_imaginary_time(backend):
     times = tc.backend.convert_to_tensor([0.0, 10.0])
 
     # Evolve
-    states = tc.timeevol.hamiltonian_evol(times, h, psi0)
+    states = tc.timeevol.hamiltonian_evol(h, psi0, times)
 
     # Ground state is |1⟩ (eigenvalue 1.0), so after long imaginary time
     # evolution, we should approach this state
@@ -236,3 +377,78 @@ def test_krylov_evol_subspace_accuracy(backend):
     # At least verify they have the correct shape
     assert state_small.shape == (1, 2**n)
     assert state_large.shape == (1, 2**n)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_krylov_evol_scan_impl(backend):
+    """Test krylov_evol with scan_impl=True"""
+    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=True)
+
+    c = tc.Circuit(n)
+    c.x([1, 2])
+    psi0 = c.state()
+
+    # Evolution times
+    times = tc.backend.convert_to_tensor([0.0, 0.5])
+
+    # Perform Krylov evolution with scan implementation
+    states_scan = tc.timeevol.krylov_evol(
+        h, psi0, times, subspace_dimension=8, scan_impl=True
+    )
+
+    states_scan_dense = tc.timeevol.krylov_evol(
+        tc.backend.to_dense(h), psi0, times, subspace_dimension=8, scan_impl=True
+    )
+
+    # Perform Krylov evolution with regular implementation
+    states_regular = tc.timeevol.krylov_evol(
+        h, psi0, times, subspace_dimension=8, scan_impl=False
+    )
+
+    # Check output shapes
+    assert states_scan.shape == (2, 2**n)
+    assert states_regular.shape == (2, 2**n)
+
+    # Results should be the same (up to numerical precision)
+    np.testing.assert_allclose(states_scan, states_regular, atol=1e-5)
+    np.testing.assert_allclose(states_scan_dense, states_regular, atol=1e-5)
+
+    # All states should be normalized
+    for state in states_scan:
+        norm = tc.backend.norm(state)
+        np.testing.assert_allclose(norm, 1.0, atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("jaxb")])
+def test_krylov_evol_gradient(backend):
+    """Test gradient computation with krylov_evol"""
+    n = 5
+    # 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)
+
+    c = tc.Circuit(n)
+    c.x([1, 2])
+    psi0 = c.state()
+
+    # Evolution time
+    t = tc.backend.convert_to_tensor([1.0])
+
+    # Define a simple loss function based on the evolved state
+    def loss_function(t):
+        states = tc.timeevol.krylov_evol(
+            h, psi0, t, subspace_dimension=8, scan_impl=True
+        )
+        # Compute the sum of absolute values of the final state as a simple loss
+        return tc.backend.sum(tc.backend.abs(states[0]))
+
+    grad_fn = tc.backend.jit(tc.backend.grad(loss_function))
+    gradient = grad_fn(t)
+    print(gradient)