jax.lib.xla_bridge.process_index

jax.lib.xla_bridge.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 (Union[str, Client, None]) – This is an experimental feature and the API is likely to change. Optional, a string representing the xla backend: 'cpu', 'gpu', or 'tpu'.

Return type

int

Returns

Integer process index.