Repository: google/jax

Issue:
Add missing exports to jax.config and jax.debug
### Description

Add missing exports to [jax.config](https://github.com/google/jax/blob/main/jax/config.py) and [jax.debug](https://github.com/google/jax/blob/main/jax/debug.py) (and any other such files). See #7570 for context.

Example:
```python3
from jax.config import config
from jax.debug import print, callback, breakpoint
```

mypy yields
```
example.py:1: error: Module "jax.config" does not explicitly export attribute "config"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "print"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "callback"  [attr-defined]
example.py:2: error: Module "jax.debug" does not explicitly export attribute "breakpoint"  [attr-defined]
```

pyright yields
```
  /Users/carlos/Desktop/example.py:1:24 - error: "config" is not exported from module "jax.config"
    Import from "jax._src.config" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:23 - error: "print" is not exported from module "jax.debug"
    Import from "jax._src.debugging" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:30 - error: "callback" is not exported from module "jax.debug"
    Import from "jax._src.debugging" instead (reportPrivateImportUsage)
  /Users/carlos/Desktop/example.py:2:40 - error: "breakpoint" is not exported from module "jax.debug"
    Import from "jax._src.debugger.core" instead (reportPrivateImportUsage)
```

### What jax/jaxlib version are you using?

jax 0.4.8, jaxlib 0.4.7

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.11.2, macOS 11.7.4

### NVIDIA GPU info

_No response_


Discussion:
> ```python
> from jax.config import config
> ```
I think this is not meant to be a supported import. This is the supported way to get the same thing:
```python
from jax import config
```
(yes, it's strange that we rewrite the `config` module with the `config` object: we need to clean that up at some point)
> ```python
> from jax.debug import print, callback, breakpoint
> ```
We should fix these. `breakpoint` is easy (we just need to add `as breakpoint`), but I'm not sure how to tell `mypy` that `callback` and `print` are meant to be exported, because they are renamed on import: https://github.com/google/jax/blob/87a1fea1c71b9b059b286c86bda4942cea766e3a/jax/debug.py#L14-L20

Do you know how to tell mypy that the renamed  imports are meant as exported names?

Resolving patch:
diff --git a/jax/debug.py b/jax/debug.py
--- a/jax/debug.py
+++ b/jax/debug.py
@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+__all__ = ["callback", "print", "DebugEffect", "visualize_array_sharding",
+           "inspect_array_sharding", "visualize_sharding", "breakpoint"]
 from jax._src.debugging import debug_callback as callback
 from jax._src.debugging import debug_print as print
 from jax._src.debugging import DebugEffect


Test patch: