version1.2.1

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## Unreleased
 
+## v1.2.1
+
+### Fixed
+
+- Jax can be omited in the env
+
 ## v1.2.0
 
 ### Added

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1"
 __author__ = "TensorCircuit Authors"
 __creator__ = "refraction-ray"
 

tensorcircuit/interfaces/jax.py

@@ -5,9 +5,6 @@
 from typing import Any, Callable, Tuple, Optional, Union, Sequence
 from functools import wraps, partial
 
-import jax
-from jax import custom_vjp
-
 from ..cons import backend
 from .tensortrans import general_args_to_backend
 
@@ -22,6 +19,8 @@ def jax_wrapper(
     ] = None,
     output_dtype: Optional[Union[Any, Sequence[Any]]] = None,
 ) -> Callable[..., Any]:
+    import jax
+
     @wraps(fun)
     def fun_jax(*x: Any) -> Any:
         def wrapped_fun(*args: Any) -> Any:
@@ -129,6 +128,9 @@ def create_jax_function(
     output_shape: Optional[Union[Tuple[int, ...], Tuple[()]]] = None,
     output_dtype: Optional[Any] = None,
 ) -> Callable[..., Any]:
+    import jax
+    from jax import custom_vjp
+
     if jit:
         fun = backend.jit(fun)