dev(hansbug): fix bug of jax integration

test/tree/integration/test_jax.py

@@ -7,6 +7,7 @@ from treevalue import FastTreeValue, register_for_jax
 
 try:
     import jax
+    from jax.tree_util import register_pytree_node
 except (ModuleNotFoundError, ImportError):
     jax = None
 

test/tree/integration/test_torch.py

@@ -6,6 +6,7 @@ from treevalue import FastTreeValue, register_for_torch
 
 try:
     import torch
+    from torch.utils._pytree import _register_pytree_node
 except (ImportError, ModuleNotFoundError):
     torch = None
 

treevalue/tree/integration/jax.py

@@ -3,13 +3,16 @@ from functools import wraps
 
 try:
     import jax
+    from jax.tree_util import register_pytree_node
 except (ModuleNotFoundError, ImportError):
     from .cjax import register_for_jax as _original_register_for_jax
 
 
     @wraps(_original_register_for_jax)
     def register_for_jax(cls):
-        warnings.warn(f'Jax is not installed, registration of {cls!r} will be ignored.')
+        warnings.warn(f'Jax doesn\'t have tree_util module due to either not installed '
+                      f'or the installed version is too low, '
+                      f'so the registration of {cls!r} will be ignored.')
 else:
     from .cjax import register_for_jax
     from ..tree import TreeValue

treevalue/tree/integration/torch.py

@@ -3,13 +3,16 @@ from functools import wraps
 
 try:
     import torch
+    from torch.utils._pytree import _register_pytree_node
 except (ModuleNotFoundError, ImportError):
     from .ctorch import register_for_torch as _original_register_for_torch
 
 
     @wraps(_original_register_for_torch)
     def register_for_torch(cls):
-        warnings.warn(f'Torch is not installed, registration of {cls!r} will be ignored.')
+        warnings.warn(f'Pytree module is not included in the Torch installation '
+                      f'or the installed version is too low, '
+                      f'so the registration of {cls!r} will be ignored.')
 else:
     from .ctorch import register_for_torch
     from ..tree import TreeValue