Sergei Lebedev 3cea57d9d1 Slice.from_slice now works for slices with a negative start index ...

The implementation still requires the step to be 1, so any slice
with a negative start index has size 0.

2024-01-29 13:13:06 +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

integrate attrs in jax.jit

2024-01-27 17:44:43 -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

[Pallas/Mosaic] Add support for nested ref.at

2024-01-02 21:54:15 -08:00

types.py

Add nested pipeline/pallas_call support for TPU meta-programming of collectives + compute.

2024-01-18 21:44:39 -08:00

utils.py

integrate attrs in jax.jit

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