jax.process_index

jax.process_index(backend=None)

返回此进程的整数进程索引。

在大多数平台上，这将始终为 0。不过，这在多进程平台上会有所不同。

参数：

backend (str | xla_client.Client | None | None) – 这是一个实验性功能，API 可能会发生变化。可选，表示 xla 后端的字符串：'cpu'、'gpu' 或 'tpu'。

返回值：

整数进程索引。

返回类型：

int