Sergei Lebedev 65d3058944 Migrate a subset of internal modules to use state objects ...

The motivation here is to gradually replace all dynamic lookups on `jax.config`
with statically-typed state objects, which are more type checker/IDE friendly.

PiperOrigin-RevId: 571932143

2023-10-09 07:29:53 -07:00

..

__init__.py

[Pallas] Add async_copy_to and async_remote_copy_to for doing DMAs.

2023-09-14 14:32:08 -07:00

discharge.py

Migrate a subset of internal modules to use state objects

2023-10-09 07:29:53 -07:00

primitives.py

Use lower-case PEP 585 names for types.

2023-06-23 15:12:14 -07:00

types.py

MAINT Clean up leftover Array = Any aliases in jax/_src/**.py

2023-10-01 12:19:21 +01:00

utils.py

Add initial implementation of a run_state primitive

2023-04-03 21:32:32 -07:00