rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-25 02:36:06 +00:00

History

Hyeontaek Lim c88ea23035 [JAX] Add caching to colocated_python.colocated_cpu_devices()

For a deployment with many devices, `colocated_python.colocated_cpu_devices()`
can take some time to find colocated devices as it needs to find matching
devices one by one in Python.

This change adds caching as an optimization to reduce the overall cost of API
calls.

PiperOrigin-RevId: 740930124

2025-03-26 15:48:48 -07:00

__init__.py

Update references to JAX's GitHub repo

2025-03-04 18:31:09 -08:00

api.py

[JAX] Add caching to colocated_python.colocated_cpu_devices()

2025-03-26 15:48:48 -07:00

func_backend.py

[JAX] Add the function API of jax.experimental.colocated_python

2024-10-28 12:18:48 -07:00

func.py

Updates the Colocated Python's serialization (and deserialization) implementation to utilize the recently added support for string arrays.