wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

also require jaxlib

Browse Source
This commit is contained in:
Andy Rock
2024-03-08 18:17:50 +01:00
committed by Patrick Kidger
parent 102e499d61
commit 696cc5b90f
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/_decorator.py
+1
View File
@@ -192,6 +192,7 @@ def jaxtyped(fn=_sentinel, *, typechecker=_sentinel):
if (
_tb_flag
and importlib.util.find_spec("jax") is not None
and importlib.util.find_spec("jaxlib") is not None
and importlib.util.find_spec("jax._src.traceback_util") is not None
):
import jax._src.traceback_util as traceback_util