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

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

@@ -1,4 +1,4 @@
-tensorcircuit/__init__.py,sha256=GNg45VO8J-m5ApDmmQbFvTPIzXvBpREXDg5MiVgCqqI,2055
+tensorcircuit/__init__.py,sha256=3yUehn2Xo7YNGDZZMC3I-Pz8sbvVzWsuXV0XDFWkGoQ,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=LWKOw4Y4CUtUO4a_72BVzYPH26PbubGYOaJlXkmA350,11439
+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=tzhOFZSOkZ_iraMvY7g84VbWBpYcELX1O0bhUipZ7YQ,59131
+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=sd1migp_E2FWjchvOeYRuyM47yexegT2_SW_ukSYSF8,14171
+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=N8iL2P76ZRp3USJSyFdZ2LWG8muXYe1kOuengv1ijXo,36444
+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
@@ -86,10 +86,10 @@ 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.dev20250731.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
@@ -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=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.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.dev20250731.dist-info/METADATA,sha256=QqZGkP34B7fJ3L38AiVOyyAgoYx_pZlrkWyWx8DcWEs,34922
+tensorcircuit_nightly-1.3.0.dev20250731.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tensorcircuit_nightly-1.3.0.dev20250731.dist-info/top_level.txt,sha256=O_Iqeh2x02lasEYMI9iyPNNNtMzcpg5qvwMOkZQ7n4A,20
+tensorcircuit_nightly-1.3.0.dev20250731.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_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
@@ -138,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
@@ -236,3 +376,266 @@ 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)
+
+
+@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 == ()