new

Author: Huang-Xu-Yang

tensorcircuit/timeevol.py

@@ -438,22 +438,22 @@ def _solve_ode(
     rtol = solver_kws.get("rtol", 1e-12)
     atol = solver_kws.get("atol", 1e-12)
     ode_backend = solver_kws.get("ode_backend", "jaxode")
+    max_steps = solver_kws.get("max_steps", 10000)
 
     ts = backend.convert_to_tensor(times)
     ts = backend.cast(ts, dtype=rdtypestr)
 
     if ode_backend == "jaxode":
         from jax.experimental.ode import odeint
 
-        s1 = odeint(f, s, ts, rtol=rtol, atol=atol, *args)
+        s1 = odeint(f, s, ts, rtol=rtol, atol=atol, mxstep=max_steps, *args)
         return s1
 
     import diffrax
 
     # Ignore complex warning
     warnings.simplefilter("ignore", category=UserWarning, append=True)
 
-    max_steps = solver_kws.get("max_steps", 10000)
     solver = solver_kws.get("solver", "Tsit5")
     dt0 = solver_kws.get("dt0", 0.01)
     all_solvers = {

tests/test_timeevol.py

@@ -11,7 +11,7 @@
 import tensorcircuit as tc
 
 
-def test_circuit_ode_evol(jaxb):
+def test_circuit_ode_evol(highp, jaxb):
     def h_square(t, b):
         return (tc.backend.sign(t - 1.0) + 1) / 2 * b * tc.gates.x().tensor
 
@@ -33,7 +33,11 @@ def h_square_sparse(t, b):
     c.cx(0, 1)
     c.h(2)
     c = tc.timeevol.evol_global(
-        c, h_square_sparse, 2.0, tc.backend.convert_to_tensor(0.2)
+        c,
+        h_square_sparse,
+        2.0,
+        tc.backend.convert_to_tensor(0.2),
+        ode_backend="diffrax",
     )
     c.rx(1, theta=np.pi - 0.4)
     np.testing.assert_allclose(c.expectation_ps(z=[1]), 1.0, atol=1e-5)