jax.process_indices — JAX documentation (original) (raw)

Contents

• process_indices()

jax.process_indices#

jax.process_indices(backend=None)[source]#

Returns the list of all JAX process indices associated with the backend.

Parameters:

backend (str | xla_client.Client | None | None) – This is an experimental feature and the API is likely to change. Optional, a string representing the xla backend: 'cpu', 'gpu', or'tpu'.

Returns:

List of integer process indices.

Return type:

list[int]