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) – 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.