xla icon indicating copy to clipboard operation
xla copied to clipboard

Published 20 hours ago • verified publisher icon openxla

[JAX] Add PyClient::GetAllDevices() and expose it as an internal JAX backend API

Open copybara-service[bot] opened this issue 1 year ago • 0 comments

[JAX] Add PyClient::GetAllDevices() and expose it as an internal JAX backend API

JAX backend forwards xla::ifrt::Client::GetAllDevices() to xla::PyClient::GetAllDevices(), which is accessible via JAX backend.get_all_devices(). This API is an internal JAX API that is used for building an experimental mesh utils API (finding colocated CPU devices) and should not be used by the user code.

copybara-service[bot] avatar Oct 02 '24 02:10 copybara-service[bot]