Sergei Lebedev c033653e28 Deduped three identical implementations of `hoist_consts_to_refs`

2024-07-12 12:48:28 +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

[Pallas] Simplify DMA discharge rule by calling new helpers from JAX state machinery

2024-07-08 19:09:53 -07:00

indexing.py

Run pyupgrade --py310-plus.

2024-06-26 16:10:18 -04:00

primitives.py

Removed noop # type: ignore comments

2024-05-19 21:01:29 +01:00

types.py

[Pallas TPU] Add support for dynamic sized (tile aligned) DMAs

2024-03-22 16:59:32 -07:00

utils.py

Deduped three identical implementations of `hoist_consts_to_refs`

2024-07-12 12:48:28 +01:00