Fix mypy error caused by cyclic import dependency.

Add tests that source_info_util is using the same directory as the jax root module.
jax-ml · Nov 18, 2020 · 13db2c3 · 13db2c3
1 parent f939e71
commit 13db2c3
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/jax/_src/source_info_util.py b/jax/_src/source_info_util.py
@@ -17,7 +17,7 @@
 import threading
 from typing import Any, Optional
 
-import jax
+import jax.version
 from jax.lib import xla_client
 
 from jax._src import traceback_util
@@ -27,7 +27,7 @@
 Traceback = Any  # xla_client.Traceback
 Frame = Any  # xla_client.Traceback::Frame
 
-_jax_path = os.path.dirname(jax.__file__)
+_jax_path = os.path.dirname(jax.version.__file__)
 
 def user_frame(source_info: Optional[Traceback]) -> Optional[Frame]:
   """Heuristic that guesses the identity of the user's code in a stack trace."""

diff --git a/tests/errors_test.py b/tests/errors_test.py
@@ -12,17 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import os
 import re
 import traceback
 import unittest
 
 from absl.testing import absltest
 
+import jax
 from jax import grad, jit, vmap
 import jax.numpy as jnp
 from jax import test_util as jtu
 from jax._src import traceback_util
+from jax._src import source_info_util
 
 
 from jax.config import config
@@ -145,5 +147,12 @@ def outer(x):
     self.assertIsInstance(e.__cause__.__cause__,
                           traceback_util.FilteredStackTrace)
 
+
+class SourceInfoTest(jtu.JaxTestCase):
+
+  def testJaxPathMatchesSourceInfoPath(self):
+    self.assertEqual(source_info_util._jax_path, os.path.dirname(jax.__file__))
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())