Sergei Lebedev b7715e279d Another take at enabling Pallas GPU tests on x64 ...

Note that for_loop_p no longer assumes that the loop index is an int32.

Closes #18847

2024-07-23 09:19:01 +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] 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

Another take at enabling Pallas GPU tests on x64

2024-07-23 09:19:01 +00:00

types.py

Another take at enabling Pallas GPU tests on x64

2024-07-23 09:19:01 +00:00

utils.py

[pallas] Fix the handling of captured consts

2024-07-22 13:34:32 +03:00