rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-23 22:26:04 +00:00

History

Ayaka fc1af8d050 Support strided load / store in interpret mode

This is a part of the efforts to fix the indexing implementation in JAX state. This PR adds support for strides in array indexing. In other words, the aim of the PR is to support this test: bb160cf54e/tests/pallas/ops_test.py (L772-L786)

This PR adds a set of test cases that makes it easier to track the completeness of the indexing implementation in JAX state. Test cases that are not yet supported are temporarily commented out.

PiperOrigin-RevId: 668402290

2024-08-28 05:10: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

Support strided load / store in interpret mode

2024-08-28 05:10:54 -07:00

indexing.py

Run pyupgrade --py310-plus.

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

primitives.py

Removed various ununsed functions