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

jax.process_index#

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

Returns the integer process index of this process.

On most platforms, this will always be 0. This will vary on multi-process platforms though.

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:

Integer process index.

Return type:

int