rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-24 04:56:07 +00:00

History

Matthew Johnson 3a403f2a0e [mutable-arrays] move MutableArray, add eager, improve tests, fix bug

1. move MutableArray to core.py, and some handlers to their respective files
2. fix a bug in aliasing setup (it was just broken before, now better test coverage)
3. add eager support by enabling get_p, swap_p, and addupdate_p impls
4. improve tests slightly

2024-03-01 15:03:23 -08: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

[mutable-arrays] allow state effects in jit by building in run_state

2024-02-29 21:50:19 -08:00

indexing.py

Slice.from_slice now works for slices with a negative start index

2024-01-29 13:13:06 +00:00

primitives.py

[mutable-arrays] move MutableArray, add eager, improve tests, fix bug