Sergei Lebedev f936613b06 Upgrade remaining sources to Python 3.9 ...

This PR is a follow up to #18881.

The changes were generated by adding

    from __future__ import annotations

to the files which did not already have them and running

    pyupgrade --py39-plus --keep-percent-format {jax,tests,jaxlib,examples,benchmarks}/**/*.py

2023-12-13 10:29:45 +00: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

[Pallas/Mosaic] Fixes for interpret mode on TPU

2023-12-06 11:03:05 -08:00

primitives.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

types.py

Upgrade remaining sources to Python 3.9

2023-12-13 10:29:45 +00:00

utils.py

Add initial implementation of a run_state primitive

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