Peter Hawkins 335b2cfb26 [JAX] Prepare not to export jax._src by default. ...

Currently
```
import jax
```
populates `jax._src` in the names exported from JAX. This change prepares for not exporting `jax._src` by default.

In particular, explicitly import modules from jax._src and refer to those imports rather than assuming jax._src contents will be around later. This is a common pattern in tests.

This change does not yet remove any exported names.

Issue https://github.com/google/jax/issues/11951

PiperOrigin-RevId: 469480816

2022-08-23 09:36:47 -07:00

..

mlir

Fix ModuleNotFoundError for phawkins only with version

2022-06-28 22:42:45 +00:00

__init__.py

Add PyShardedToken for sharded execution.

2022-08-18 13:41:50 -07:00

xla_bridge.py

[JAX] Prepare not to export jax._src by default.

2022-08-23 09:36:47 -07:00