jax authors 70f2ef211f Merge pull request #20971 from google:mutable-array-scan ...

PiperOrigin-RevId: 630130893

2024-05-02 11:40:54 -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

Merge pull request #20971 from google:mutable-array-scan

2024-05-02 11:40:54 -07:00

indexing.py

Only perform checks on slice sizes if they're static.

2024-04-23 18:02:07 -07:00

primitives.py

[Pallas] Fix a typo in error message of swap rule.

2024-03-29 13:04:44 -07:00

types.py

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

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

utils.py

integrate attrs in jax.jit

2024-01-27 17:44:43 -08:00