jax
9ceb3f41 - [Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file.

Commit

70 days ago

[Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file. This is the first step in a list of changes that aim to make the data structures used by the interpreter re-usable. PiperOrigin-RevId: 818595222

References

#31381 - Remove old ROCm build code

#32529 - [Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file.

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency

#585 - update a test for checking zero ROCm GPU event

Author

Google-ML-Automation

Committer

Google-ML-Automation

Parents

f44bc6e3

jax 9ceb3f41 - [Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file.

jax
9ceb3f41 - [Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file.