Sergei Lebedev 5ab05e42c9 MAINT Clean up leftover Array = Any aliases in jax/_src/**.py ...

I had to revert to using `Any` for `RaggedAxis.ragged_axes` because pytype
found more latent type errors, which require the understanding of ragedness
and dynamic shapes internals to fix properly.

2023-10-01 12:19:21 +01: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

Fix typos across the package

2023-09-22 14:54:31 -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