Use the jax module to determine the jax source directory base, rather than the source_info module.
Use the jax module to determine the jax source directory base, rather than the source_info module.

Fixes some misattributed code source information.

google/jax,google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability