Pop the script directory from sys.path.

Quoting http://python-notes.curiousefficiency.org/en/latest/python_concepts/import_traps.html#the-double-import-trap:

"""
There?s a reason the general ?no package directories on sys.path? guideline
exists, and the fact that the interpreter itself doesn?t follow it when
determining sys.path[0] is the root cause of all sorts of grief.
"""

Also relevant: bazelbuild/bazel#7091

This is addressed toward issue tensorflow#539

PiperOrigin-RevId: 268058051

tensorflow_probability/python/experimental/substrates/numpy/rewrite.py

@@ -18,11 +18,17 @@
 from __future__ import division
 from __future__ import print_function
 
+import sys
+if not sys.path[0].endswith('.runfiles'):
+  sys.path.pop(0)
+
+# pylint: disable=g-import-not-at-top,g-bad-import-order
 import collections
 
 # Dependency imports
 from absl import app
 from absl import flags
+# pylint: enable=g-import-not-at-top,g-bad-import-order
 
 flags.DEFINE_boolean('numpy_to_jax', False,
                      'Whether or not to rewrite numpy imports to jax.numpy')

tensorflow_probability/python/internal/backend/numpy/gen_linear_operators.py

@@ -18,14 +18,18 @@
 from __future__ import division
 from __future__ import print_function
 
+import sys
+if not sys.path[0].endswith('.runfiles'):
+  sys.path.pop(0)
+
+# pylint: disable=g-import-not-at-top,g-bad-import-order
 import importlib
 import inspect
 import re
 
-# Dependency imports
-
 from absl import app
 from absl import flags
+# pylint: enable=g-import-not-at-top,g-bad-import-order
 
 FLAGS = flags.FLAGS